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 2, 06:00 UTC
[Model Runner V2] Use actual batch max_seq_len for attn metadata (#43991)
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>
zhrrr committed
Jun 2, 2026 at 14:07 UTC
1edfd09ffd1ff3c15cea35d50441b2b9d78629f1
3 files changed
+14
-3
vllm/v1/worker/gpu/model_states/default.py
+1
-1
@@ -178,7 +178,7 @@ class DefaultModelState(ModelState):
178
# Capture with worst-case max_seq_len so the graph is valid at any replay.
179
max_seq_len = self.max_model_len
180
else:
181
- max_seq_len = int(seq_lens_cpu_upper_bound[:num_reqs].max().item())
181
+ max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item()
182
attn_metadata = build_attn_metadata(
183
attn_groups=attn_groups,
184
num_reqs=num_reqs,
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+7
-1
@@ -86,6 +86,12 @@ class MambaHybridModelState(DefaultModelState):
86
num_tokens = input_batch.num_tokens
87
query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
88
max_query_len = input_batch.num_scheduled_tokens.max().item()
89
+ seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound
90
+ if for_capture:
91
+ # Capture with worst-case max_seq_len so the graph is valid at any replay.
92
+ max_seq_len = self.max_model_len
93
+ else:
94
+ max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item()
95
96
is_prefilling = torch.zeros(num_reqs, dtype=torch.bool, device="cpu")
97
is_prefilling[: input_batch.num_reqs] = torch.from_numpy(
@@ -129,7 +135,7 @@ class MambaHybridModelState(DefaultModelState):
135
query_start_loc_cpu=query_start_loc_cpu,
136
max_query_len=max_query_len,
137
seq_lens=input_batch.seq_lens,
132
- max_seq_len=self.max_model_len,
138
+ max_seq_len=max_seq_len,
139
block_tables=block_tables,
140
slot_mappings=slot_mappings,
141
kv_cache_config=kv_cache_config,
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+6
-1
@@ -52,6 +52,7 @@ class EagleSpeculator:
52
self.max_num_reqs = self.scheduler_config.max_num_seqs
53
self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
54
self.max_model_len = vllm_config.model_config.max_model_len
55
+ self.draft_max_seq_len = self.max_model_len
56
# We need to get the hidden size from the draft model config because
57
# the draft model's hidden size can be different from the target model's
58
# hidden size (e.g., Llama 3.3 70B).
@@ -416,7 +417,7 @@ class EagleSpeculator:
417
query_start_loc_cpu=query_start_loc_cpu,
418
max_query_len=1,
419
seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],
419
- max_seq_len=self.max_model_len,
420
+ max_seq_len=self.draft_max_seq_len,
421
block_tables=block_tables,
422
slot_mappings=slot_mappings,
423
kv_cache_config=self.kv_cache_config,
@@ -494,6 +495,10 @@ class EagleSpeculator:
495
num_tokens = input_batch.num_tokens_after_padding
496
num_reqs = input_batch.num_reqs
497
max_query_len = input_batch.num_scheduled_tokens.max()
498
+ max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item()
499
+ self.draft_max_seq_len = min(
500
+ max_seq_len + self.num_speculative_steps, self.max_model_len
501
+ )
502
503
# NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the
504
# number of rejected tokens, we maintain the size of eagle's input_ids and