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 14, 05:30 UTC
[Misc] Add max_seq_len to CommonAttentionMetadata (#23216)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Aug 20, 2025 at 09:05 UTC
d6d13bd49ed7fda56ac6a1b0aa53621490c975ac
12 files changed
+22
-7
tests/v1/attention/utils.py
+2
@@ -58,6 +58,7 @@ def create_common_attn_metadata(
58
dtype=torch.int32,
59
device=device)
60
seq_lens_cpu = seq_lens.cpu()
61
+ max_seq_len = int(seq_lens_cpu.max())
62
63
# Create computed tokens (context length for each sequence)
64
context_lens = [
@@ -101,6 +102,7 @@ def create_common_attn_metadata(
102
num_reqs=batch_spec.batch_size,
103
num_actual_tokens=num_tokens,
104
max_query_len=max_query_len,
105
+ max_seq_len=max_seq_len,
106
block_table_tensor=block_table_tensor,
107
slot_mapping=slot_mapping,
108
causal=True,
tests/v1/spec_decode/test_tree_attention.py
+2
@@ -50,6 +50,7 @@ def forward_attention(
50
dtype=torch.int32,
51
)
52
context_lens = seq_lens - query_lens
53
+ max_seq_len = int(seq_lens.max())
54
max_query_len = q_len
55
num_actual_tokens = query_start_loc[-1]
56
@@ -81,6 +82,7 @@ def forward_attention(
82
num_reqs=batch_size,
83
num_actual_tokens=num_actual_tokens,
84
max_query_len=max_query_len,
85
+ max_seq_len=max_seq_len,
86
block_table_tensor=block_table,
87
slot_mapping=slot_mapping,
88
)
vllm/v1/attention/backends/flash_attn.py
+1
-1
@@ -233,7 +233,7 @@ class FlashAttentionMetadataBuilder(
233
num_reqs = common_attn_metadata.num_reqs
234
num_actual_tokens = common_attn_metadata.num_actual_tokens
235
max_query_len = common_attn_metadata.max_query_len
236
- max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
236
+ max_seq_len = common_attn_metadata.max_seq_len
237
query_start_loc = common_attn_metadata.query_start_loc
238
seq_lens = common_attn_metadata.seq_lens
239
seq_lens_cpu = common_attn_metadata.seq_lens_cpu
vllm/v1/attention/backends/flashinfer.py
+1
-1
@@ -463,7 +463,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
463
464
page_size = self.page_size
465
max_q_len = common_attn_metadata.max_query_len
466
- max_seq_len = common_attn_metadata.seq_lens_cpu.max().item()
466
+ max_seq_len = common_attn_metadata.max_seq_len
467
seq_lens = common_attn_metadata.seq_lens
468
seq_lens_cpu = common_attn_metadata.seq_lens_cpu
469
block_table_tensor = common_attn_metadata.block_table_tensor
vllm/v1/attention/backends/flex_attention.py
+1
-1
@@ -305,7 +305,7 @@ class FlexAttentionMetadataBuilder(
305
num_actual_tokens = common_attn_metadata.num_actual_tokens
306
max_query_len = common_attn_metadata.max_query_len
307
308
- max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
308
+ max_seq_len = common_attn_metadata.max_seq_len
309
query_start_loc = common_attn_metadata.query_start_loc
310
seq_lens = common_attn_metadata.seq_lens
311
block_table_tensor = common_attn_metadata.block_table_tensor
vllm/v1/attention/backends/rocm_aiter_fa.py
+1
-1
@@ -270,7 +270,7 @@ class AiterFlashAttentionMetadataBuilder(
270
271
num_actual_tokens = common_attn_metadata.num_actual_tokens
272
max_query_len = common_attn_metadata.max_query_len
273
- max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
273
+ max_seq_len = common_attn_metadata.max_seq_len
274
query_start_loc = common_attn_metadata.query_start_loc
275
seq_lens = common_attn_metadata.seq_lens
276
block_table_tensor = common_attn_metadata.block_table_tensor
vllm/v1/attention/backends/tree_attn.py
+1
-1
@@ -205,7 +205,7 @@ class TreeAttentionMetadataBuilder(
205
q_start_loc = common_attn_metadata.query_start_loc
206
max_query_len = common_attn_metadata.max_query_len
207
kv_seqlens = common_attn_metadata.seq_lens
208
- max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
208
+ max_seq_len = common_attn_metadata.max_seq_len
209
block_table = common_attn_metadata.block_table_tensor
210
slot_mapping = common_attn_metadata.slot_mapping
211
vllm/v1/attention/backends/triton_attn.py
+1
-1
@@ -90,7 +90,7 @@ class TritonAttentionMetadataBuilder(
90
num_actual_tokens = common_attn_metadata.num_actual_tokens
91
max_query_len = common_attn_metadata.max_query_len
92
93
- max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
93
+ max_seq_len = common_attn_metadata.max_seq_len
94
query_start_loc = common_attn_metadata.query_start_loc
95
seq_lens = common_attn_metadata.seq_lens
96
block_table_tensor = common_attn_metadata.block_table_tensor
vllm/v1/attention/backends/utils.py
+6
@@ -58,6 +58,8 @@ class CommonAttentionMetadata:
58
"""Total number of tokens in batch"""
59
max_query_len: int
60
"""Longest query in batch"""
61
+ max_seq_len: int
62
+ """Longest context length in batch"""
63
64
block_table_tensor: torch.Tensor
65
slot_mapping: torch.Tensor
@@ -107,6 +109,7 @@ def _make_metadata_with_slice(
109
110
seq_lens = attn_metadata.seq_lens[request_slice]
111
seq_lens_cpu = attn_metadata.seq_lens_cpu[request_slice]
112
+ max_seq_len = int(seq_lens_cpu.max())
113
num_computed_tokens_cpu = attn_metadata.num_computed_tokens_cpu[
114
request_slice]
115
@@ -128,6 +131,7 @@ def _make_metadata_with_slice(
131
num_reqs=num_requests,
132
num_actual_tokens=num_actual_tokens,
133
max_query_len=max_query_len,
134
+ max_seq_len=max_seq_len,
135
block_table_tensor=block_table_tensor,
136
slot_mapping=slot_mapping,
137
)
@@ -520,6 +524,7 @@ def make_local_attention_virtual_batches(
524
525
query_start_loc_cpu = torch.from_numpy(cu_seqlens_q_local)
526
seq_lens_cpu = torch.from_numpy(seqlens_k_local)
527
+ max_seq_len = int(seq_lens_cpu.max())
528
529
return CommonAttentionMetadata(
530
query_start_loc_cpu=query_start_loc_cpu,
@@ -531,6 +536,7 @@ def make_local_attention_virtual_batches(
536
num_reqs=len(seq_lens_cpu),
537
num_actual_tokens=common_attn_metadata.num_actual_tokens,
538
max_query_len=seqlens_q_local.max(),
539
+ max_seq_len=max_seq_len,
540
block_table_tensor=block_table_local,
541
slot_mapping=common_attn_metadata.slot_mapping,
542
causal=True,
vllm/v1/attention/backends/xformers.py
+1
-1
@@ -231,7 +231,7 @@ class XFormersAttentionMetadataBuilder(
231
q_seqlens = torch.diff(q_start_loc)
232
max_query_len = common_attn_metadata.max_query_len
233
kv_seqlens = common_attn_metadata.seq_lens
234
- max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
234
+ max_seq_len = common_attn_metadata.max_seq_len
235
block_table = common_attn_metadata.block_table_tensor
236
slot_mapping = common_attn_metadata.slot_mapping
237
vllm/v1/spec_decode/eagle.py
+1
@@ -582,6 +582,7 @@ class EagleProposer:
582
num_reqs=common_attn_metadata.num_reqs,
583
num_actual_tokens=total_num_tokens,
584
max_query_len=new_query_len_per_req.max().item(),
585
+ max_seq_len=new_seq_lens_cpu.max().item(),
586
block_table_tensor=common_attn_metadata.block_table_tensor,
587
slot_mapping=common_attn_metadata.slot_mapping[token_indices],
588
causal=True,
vllm/v1/worker/gpu_model_runner.py
+4
@@ -774,6 +774,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
774
self.seq_lens_np[num_reqs:].fill(0)
775
self.seq_lens.copy_(self.seq_lens_cpu, non_blocking=True)
776
seq_lens = self.seq_lens[:num_reqs]
777
+ max_seq_len = self.seq_lens_np[:num_reqs].max().item()
778
779
# Copy the tensors to the GPU.
780
self.input_ids[:total_num_scheduled_tokens].copy_(
@@ -886,6 +887,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
887
num_reqs=num_reqs,
888
num_actual_tokens=total_num_scheduled_tokens,
889
max_query_len=max_num_scheduled_tokens,
890
+ max_seq_len=max_seq_len,
891
block_table_tensor=blk_table_tensor,
892
slot_mapping=slot_mapping,
893
causal=True,
@@ -2338,6 +2340,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2340
num_reqs=num_reqs,
2341
num_actual_tokens=num_tokens,
2342
max_query_len=max_query_len,
2343
+ max_seq_len=self.max_model_len,
2344
block_table_tensor=self.input_batch.block_table[
2345
kv_cache_group_id].get_device_tensor()[:num_reqs],
2346
slot_mapping=self.input_batch.
@@ -3343,6 +3346,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3346
num_reqs=num_reqs,
3347
num_actual_tokens=total_num_scheduled_tokens,
3348
max_query_len=max_num_scheduled_tokens,
3349
+ max_seq_len=self.seq_lens_cpu[:num_reqs].max().item(),
3350
block_table_tensor=dummy_block_table,
3351
slot_mapping=dummy_slot_mapping,
3352
causal=False,