From 1edfd09ffd1ff3c15cea35d50441b2b9d78629f1 Mon Sep 17 00:00:00 2001 From: zhrrr <43847754+izhuhaoran@users.noreply.github.com> Date: Tue, 2 Jun 2026 14:07:56 +0800 Subject: [PATCH] [Model Runner V2] Use actual batch max_seq_len for attn metadata (#43991) Signed-off-by: zhuhaoran --- vllm/v1/worker/gpu/model_states/default.py | 2 +- vllm/v1/worker/gpu/model_states/mamba_hybrid.py | 8 +++++++- vllm/v1/worker/gpu/spec_decode/eagle/speculator.py | 7 ++++++- 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 7f7955a58a..9f78749e78 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -178,7 +178,7 @@ class DefaultModelState(ModelState): # Capture with worst-case max_seq_len so the graph is valid at any replay. max_seq_len = self.max_model_len else: - max_seq_len = int(seq_lens_cpu_upper_bound[:num_reqs].max().item()) + max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item() attn_metadata = build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 93115fdf64..b5abf56fbf 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -86,6 +86,12 @@ class MambaHybridModelState(DefaultModelState): num_tokens = input_batch.num_tokens query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) max_query_len = input_batch.num_scheduled_tokens.max().item() + seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound + if for_capture: + # Capture with worst-case max_seq_len so the graph is valid at any replay. + max_seq_len = self.max_model_len + else: + max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item() is_prefilling = torch.zeros(num_reqs, dtype=torch.bool, device="cpu") is_prefilling[: input_batch.num_reqs] = torch.from_numpy( @@ -129,7 +135,7 @@ class MambaHybridModelState(DefaultModelState): query_start_loc_cpu=query_start_loc_cpu, max_query_len=max_query_len, seq_lens=input_batch.seq_lens, - max_seq_len=self.max_model_len, + max_seq_len=max_seq_len, block_tables=block_tables, slot_mappings=slot_mappings, kv_cache_config=kv_cache_config, diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 8ca2882e3a..1a1ae1f63e 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -52,6 +52,7 @@ class EagleSpeculator: self.max_num_reqs = self.scheduler_config.max_num_seqs self.max_num_tokens = self.scheduler_config.max_num_batched_tokens self.max_model_len = vllm_config.model_config.max_model_len + self.draft_max_seq_len = self.max_model_len # We need to get the hidden size from the draft model config because # the draft model's hidden size can be different from the target model's # hidden size (e.g., Llama 3.3 70B). @@ -416,7 +417,7 @@ class EagleSpeculator: query_start_loc_cpu=query_start_loc_cpu, max_query_len=1, seq_lens=self.input_buffers.seq_lens[:num_reqs_padded], - max_seq_len=self.max_model_len, + max_seq_len=self.draft_max_seq_len, block_tables=block_tables, slot_mappings=slot_mappings, kv_cache_config=self.kv_cache_config, @@ -494,6 +495,10 @@ class EagleSpeculator: num_tokens = input_batch.num_tokens_after_padding num_reqs = input_batch.num_reqs max_query_len = input_batch.num_scheduled_tokens.max() + max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item() + self.draft_max_seq_len = min( + max_seq_len + self.num_speculative_steps, self.max_model_len + ) # NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the # number of rejected tokens, we maintain the size of eagle's input_ids and