A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 16:30 UTC
[Attention][2/n] Remove usage of deprecated `seq_lens_cpu` and `num_computed_tokens_cpu` CommonAttentionMetadata properties (#31774)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
Lucas Wilkinson committed
Jan 6, 2026 at 12:32 UTC
4c73be14e0397e99162ca13a8b559670c5abd3b0
3 files changed
+13
-14
vllm/v1/attention/backends/gdn_attn.py
+1
-3
@@ -142,8 +142,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
142
m = common_attn_metadata
143
144
query_start_loc = m.query_start_loc
145
- context_lens = m.num_computed_tokens_cpu
146
- context_lens_tensor = context_lens.to(query_start_loc.device, non_blocking=True)
145
+ context_lens_tensor = m.compute_num_computed_tokens()
146
nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
147
148
if (
@@ -370,6 +369,5 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
369
370
num_accepted_tokens = torch.diff(m.query_start_loc)
371
num_decode_draft_tokens_cpu = (num_accepted_tokens - 1).cpu()
373
- m._num_computed_tokens_cpu = m.seq_lens_cpu - num_accepted_tokens.cpu()
372
373
return self.build(0, m, num_accepted_tokens, num_decode_draft_tokens_cpu)
vllm/v1/attention/backends/mamba2_attn.py
+4
-1
@@ -215,7 +215,10 @@ class Mamba2AttentionMetadataBuilder(
215
num_prefills = common.num_prefills
216
num_decode_tokens = common.num_decode_tokens
217
218
- num_computed_tokens_p_cpu = common_attn_metadata.num_computed_tokens_cpu[
218
+ num_computed_tokens_cpu = (
219
+ common_attn_metadata.compute_num_computed_tokens().cpu()
220
+ )
221
+ num_computed_tokens_p_cpu = num_computed_tokens_cpu[
222
num_reqs - num_prefills : num_reqs
223
]
224
query_start_loc_p_cpu = (
vllm/v1/attention/backends/mamba_attn.py
+8
-10
@@ -138,9 +138,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
138
common_attn_metadata: CommonAttentionMetadata,
139
mamba_block_size: int,
140
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
141
- num_computed_tokens = common_attn_metadata.num_computed_tokens_cpu.to(
142
- self.device
143
- )
141
+ num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
142
# Block index of the last computed token
143
block_idx_last_computed_token = cdiv(num_computed_tokens, mamba_block_size) - 1
144
# which is <= block index for the first scheduled token
@@ -193,13 +191,12 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
191
nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
192
193
if self.vllm_config.cache_config.enable_prefix_caching:
194
+ num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
195
+
196
# Return a tensor of shape (#requests, #max blocks)
197
state_indices_tensor = common_attn_metadata.block_table_tensor
198
# Additional cache-related varaiables:
199
mamba_block_size = self.kv_cache_spec.block_size
200
- num_computed_tokens = common_attn_metadata.num_computed_tokens_cpu.to(
201
- self.device
202
- )
200
(
201
block_idx_last_computed_token,
202
block_idx_first_scheduled_token,
@@ -212,15 +209,16 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
209
state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
210
211
if num_prefills > 0:
212
+ if num_computed_tokens is None:
213
+ num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
214
+ num_computed_tokens_cpu = num_computed_tokens.cpu()
215
+
216
query_start_loc_p = (
217
common_attn_metadata.query_start_loc[-num_prefills - 1 :]
218
- num_decode_tokens
219
)
220
has_initial_states_cpu = (
220
- common_attn_metadata.num_computed_tokens_cpu[
221
- num_reqs - num_prefills : num_reqs
222
- ]
223
- > 0
221
+ num_computed_tokens_cpu[num_reqs - num_prefills : num_reqs] > 0
222
)
223
has_initial_states_p = has_initial_states_cpu.to(
224
common_attn_metadata.query_start_loc.device