Skip to content

[v1] Expose num_prompt_tokens in CommonAttentionMetadata #39744

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

Closed
asadafa123 wants to merge 1 commit into from
Closed

[v1] Expose num_prompt_tokens in CommonAttentionMetadata #39744

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions vllm/v1/attention/backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -388,6 +388,12 @@ class CommonAttentionMetadata:
(num_computed_tokens < num_prompt_tokens). Used by some backends to
distinguish actual decodes from short extends."""

num_prompt_tokens: torch.Tensor | None = None
"""(batch_size,) int tensor: original prompt length for each request.
Needed by dual-cache RoPE implementations (e.g. LongRoPE SplitByLength)
to select the correct cache per-sequence under chunked prefill, where
positions.max() of a chunk does not reflect the full prompt length."""

# WARNING: Deprecated fields. Will be removed in a future release (v0.15.0)
_seq_lens_cpu: torch.Tensor | None = None
_num_computed_tokens_cpu: torch.Tensor | None = None
Expand Down Expand Up @@ -470,6 +476,7 @@ def unpadded(
dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens),
dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu),
is_prefilling=maybe_slice_reqs(self.is_prefilling),
num_prompt_tokens=maybe_slice_reqs(self.num_prompt_tokens),
)


Expand Down
1 change: 1 addition & 0 deletions vllm/v1/worker/gpu_model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2192,6 +2192,7 @@ def _get_block_table(kv_cache_gid: int):
slot_mapping=slot_mapping_gid_0,
causal=True,
is_prefilling=is_prefilling,
num_prompt_tokens=num_prompt_tokens_cpu,

@gemini-code-assist gemini-code-assist Bot Apr 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

high

The num_prompt_tokens field is being assigned a CPU tensor (num_prompt_tokens_cpu). In CommonAttentionMetadata, fields without the _cpu suffix are expected to be device tensors. Since this metadata is intended for use in model layers (like RoPE) during the forward pass, using a CPU tensor will cause host-device synchronizations or graph breaks in torch.compile, leading to significant performance degradation. This should be a GPU tensor to allow efficient device-side access. Since the source tensor in InputBatch is pinned, you can use a non-blocking transfer here, or ideally, use a persistent GPU buffer if one is available in GPUModelRunner.

Suggested change
num_prompt_tokens=num_prompt_tokens_cpu,
num_prompt_tokens=num_prompt_tokens_cpu.to(device=self.device, non_blocking=True),

)

if self.dcp_world_size > 1:
Expand Down
Loading