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][1/n] Remove usage of deprecated `seq_lens_cpu` and `num_computed_tokens_cpu` CommonAttentionMetadata properties (#31773)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Jan 6, 2026 at 07:05 UTC
e0327c9db20652b3cef6ffc29adb3a2e1cc3fba1
9 files changed
+23
-12
tests/v1/attention/test_attention_backends.py
+2
-2
@@ -126,12 +126,12 @@ def create_and_prepopulate_kv_cache(
126
Tuple of (kv_cache, updated_block_table)
127
"""
128
batch_size = len(k_contexts)
129
- seq_lens = common_attn_metadata.seq_lens_cpu
129
+ seq_lens = common_attn_metadata.seq_lens.cpu()
130
query_lens = (
131
common_attn_metadata.query_start_loc_cpu[1:]
132
- common_attn_metadata.query_start_loc_cpu[:-1]
133
)
134
- context_lens = common_attn_metadata.num_computed_tokens_cpu
134
+ context_lens = seq_lens - query_lens
135
block_table = common_attn_metadata.block_table_tensor
136
slot_mapping = common_attn_metadata.slot_mapping
137
tests/v1/attention/test_mla_backends.py
+2
-2
@@ -154,12 +154,12 @@ def create_and_prepopulate_kv_cache(
154
MLA KV cache tensor
155
"""
156
batch_size = len(kv_c_contexts)
157
- seq_lens = common_attn_metadata.seq_lens_cpu
157
+ seq_lens = common_attn_metadata.seq_lens.cpu()
158
query_lens = (
159
common_attn_metadata.query_start_loc_cpu[1:]
160
- common_attn_metadata.query_start_loc_cpu[:-1]
161
)
162
- context_lens = common_attn_metadata.num_computed_tokens_cpu
162
+ context_lens = seq_lens - query_lens
163
block_table = common_attn_metadata.block_table_tensor
164
slot_mapping = common_attn_metadata.slot_mapping
165
tests/v1/attention/test_sparse_mla_backends.py
+1
-1
@@ -297,7 +297,7 @@ def test_sparse_backend_decode_correctness(
297
positions = np.arange(starts[-1], dtype=np.int32) - np.repeat(
298
starts[:-1], seg_lengths
299
)
300
- seq_lengths = np.asarray(common_attn_metadata.seq_lens_cpu, dtype=np.int32)
300
+ seq_lengths = np.asarray(common_attn_metadata.seq_lens.cpu(), dtype=np.int32)
301
prefix_lengths = seq_lengths - seg_lengths
302
positions += np.repeat(prefix_lengths, seg_lengths)
303
vllm/v1/attention/backends/flashinfer.py
+3
-1
@@ -870,7 +870,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
870
# Guard access to seq_lens_cpu, which may not always be needed
871
# and can be expensive to retrieve in async mode.
872
needs_seq_lens_cpu = self.use_dcp or use_cascade or not is_only_trtllm_decode
873
- seq_lens_cpu = common_attn_metadata.seq_lens_cpu if needs_seq_lens_cpu else None
873
+ seq_lens_cpu = (
874
+ common_attn_metadata.seq_lens.cpu() if needs_seq_lens_cpu else None
875
+ )
876
seq_lens_np = seq_lens_cpu.numpy() if seq_lens_cpu is not None else None
877
num_blocks_np = (
878
(seq_lens_np + (page_size - 1)) // page_size
vllm/v1/attention/backends/flex_attention.py
+1
-3
@@ -727,9 +727,7 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
727
block_table_tensor, seq_lens, block_size, num_gpu_blocks
728
)
729
730
- offset_tensor = common_attn_metadata.num_computed_tokens_cpu.to(
731
- self.device, non_blocking=True
732
- )
730
+ offset_tensor = common_attn_metadata.compute_num_computed_tokens()
731
732
out = FlexAttentionMetadata(
733
causal=common_attn_metadata.causal,
vllm/v1/attention/backends/mla/common.py
+3
-1
@@ -791,7 +791,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
791
792
prefill_metadata = None
793
if num_prefills > 0:
794
- num_computed_tokens_cpu = common_attn_metadata.num_computed_tokens_cpu
794
+ num_computed_tokens_cpu = (
795
+ common_attn_metadata.compute_num_computed_tokens().cpu()
796
+ )
797
798
reqs_start = num_decodes # prefill_start
799
vllm/v1/attention/backends/mla/flashmla_sparse.py
+1
-1
@@ -511,7 +511,7 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
511
# For pure decode batches, prefill_request_id will be None
512
# For mixed batches, it will have -1 for decode and request_id for prefill
513
if num_prefills > 0:
514
- seq_lens_cpu = common_attn_metadata.seq_lens_cpu
514
+ seq_lens_cpu = common_attn_metadata.seq_lens.cpu()
515
seq_lens = common_attn_metadata.seq_lens
516
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
517
vllm/v1/attention/backends/triton_attn.py
+1
-1
@@ -221,7 +221,7 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
221
prefix_kv_lens = torch.tensor(
222
[common_prefix_len], dtype=torch.int32, device=self.device
223
)
224
- suffix_kv_lens = common_attn_metadata.seq_lens_cpu - common_prefix_len
224
+ suffix_kv_lens = common_attn_metadata.seq_lens.cpu() - common_prefix_len
225
suffix_kv_lens = suffix_kv_lens.to(self.device)
226
else:
227
cu_prefix_query_lens = None
vllm/v1/attention/backends/utils.py
+9
@@ -100,6 +100,8 @@ class CommonAttentionMetadata:
100
_seq_lens_cpu: torch.Tensor | None = None
101
_num_computed_tokens_cpu: torch.Tensor | None = None
102
103
+ _num_computed_tokens_cache: torch.Tensor | None = None
104
+
105
@property
106
@deprecated(
107
"""
@@ -130,6 +132,13 @@ class CommonAttentionMetadata:
132
self._num_computed_tokens_cpu = self.seq_lens_cpu - query_seq_lens
133
return self._num_computed_tokens_cpu
134
135
+ def compute_num_computed_tokens(self) -> torch.Tensor:
136
+ """Compute num_computed_tokens on device (seq_lens - query_lens)."""
137
+ if self._num_computed_tokens_cache is None:
138
+ query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
139
+ self._num_computed_tokens_cache = self.seq_lens - query_lens
140
+ return self._num_computed_tokens_cache
141
+
142
# TODO(lucas): remove once we have FULL-CG spec-decode support
143
def unpadded(
144
self, num_actual_tokens: int, num_actual_reqs: int