Skip to content

[perf] Use CPU tensor to reduce GPU->CPU sync #25884

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 2 commits into from
Sep 30, 2025
Merged

[perf] Use CPU tensor to reduce GPU->CPU sync #25884

merged 2 commits into from
Sep 30, 2025
+1 −1

Conversation

lhtin
Copy link
Contributor

@lhtin lhtin commented Sep 29, 2025 •
edited by github-actions bot
Loading

Use seq_lens_cpu instead of seq_lens to reduce GPU->CPU sync.

Purpose

Reduce unnecessary GPU->CPU sync, since it will affect the perf of Async Scheduling+MTP.
Clipboard_Screenshot_1759149439

Test Plan

Test Result

Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
  • (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

@lhtin
[perf] Use CPU tensor to reduce GPU->CPU sync
9a6d1fc 
Signed-off-by: Lehua Ding <lehuading@tencent.com>
@mergify mergify bot added the v1 label Sep 29, 2025
gemini-code-assist[bot]
gemini-code-assist bot reviewed Sep 29, 2025
View reviewed changes
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a performance optimization to reduce GPU-to-CPU synchronization during speculative decoding. The change replaces a call to .max() on a GPU tensor (seq_lens) with its CPU counterpart (seq_lens_cpu) within a conditional check. This avoids a blocking operation, which is particularly beneficial for asynchronous scheduling. The change is correct and aligns with the stated goal of improving performance. I have no further comments.

@lhtin
Copy link
Contributor Author

lhtin commented Sep 29, 2025

Which introduce by this pr(#24662), @AlonKejzman can you help review this too? thanks.

benchislett
benchislett approved these changes Sep 29, 2025
View reviewed changes
Copy link
Collaborator

@benchislett benchislett left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks!

benchislett
benchislett reviewed Sep 29, 2025
View reviewed changes
vllm/v1/worker/gpu_model_runner.py Outdated
self.speculative_config.draft_model_config.max_model_len)
input_fits_in_drafter = spec_decode_common_attn_metadata and (
spec_decode_common_attn_metadata.seq_lens.max() +
spec_decode_common_attn_metadata.seq_lens_cpu.max() +
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to use .max_seq_len here?

Copy link
Contributor Author

@lhtin lhtin Sep 29, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Indeed, and it is simpler! spec_decode_common_attn_metadata.max_seq_len is cmoputed from self.seq_lens.np[:num_reqs].max().item(), they are equivalent.

@benchislett benchislett added the ready ONLY add when PR is ready to merge/full CI is needed label Sep 29, 2025
@lhtin
[perf] Use .max_seq_len replace .seq_lens_cpu.max()
0617ab6 
Signed-off-by: Lehua Ding <lehuading@tencent.com>
@DarkLight1337 DarkLight1337 merged commit e184c9c into vllm-project:main Sep 30, 2025
42 checks passed
pdasigi pushed a commit to pdasigi/vllm that referenced this pull request Oct 2, 2025
@lhtin @pdasigi
[perf] Use CPU tensor to reduce GPU->CPU sync (vllm-project#25884)
1116b82 
Signed-off-by: Lehua Ding <lehuading@tencent.com>
yewentao256 pushed a commit that referenced this pull request Oct 3, 2025
@lhtin @yewentao256
[perf] Use CPU tensor to reduce GPU->CPU sync (#25884)
eea2536 
Signed-off-by: Lehua Ding <lehuading@tencent.com>
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@benchislett benchislett benchislett approved these changes

@WoosukKwon WoosukKwon Awaiting requested review from WoosukKwon WoosukKwon is a code owner

@robertgshaw2-redhat robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat robertgshaw2-redhat is a code owner

@njhill njhill Awaiting requested review from njhill njhill is a code owner

@ywang96 ywang96 Awaiting requested review from ywang96 ywang96 is a code owner

@comaniac comaniac Awaiting requested review from comaniac comaniac is a code owner

@alexm-redhat alexm-redhat Awaiting requested review from alexm-redhat alexm-redhat is a code owner

+1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements
Assignees
No one assigned
Labels
ready ONLY add when PR is ready to merge/full CI is needed v1
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@lhtin @benchislett @DarkLight1337