From 2997fcaa279c31aa5c232d7c2dec5cd33767cea9 Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Thu, 26 Feb 2026 17:19:28 +0000 Subject: [PATCH] fix Signed-off-by: yewentao256 --- vllm/v1/worker/gpu/model_runner.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 93767a52021..7b7409d4884 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -38,6 +38,7 @@ from vllm.logger import init_logger from vllm.model_executor.model_loader import get_model_loader from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.sequence import IntermediateTensors +from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput @@ -487,7 +488,23 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) # keep per-step metadata aligned with the captured full-graph shape. - attn_num_reqs = min(num_tokens_after_padding, self.max_num_reqs) + uniform_decode_query_len = self.cudagraph_manager.uniform_decode_query_len + query_lens = ( + query_start_loc_cpu[1 : num_reqs + 1] - query_start_loc_cpu[:num_reqs] + ) + if ( + self.cudagraph_manager.cudagraph_mode.separate_routine() + and self.cudagraph_manager.cudagraph_mode.decode_mode() + == CUDAGraphMode.FULL + and num_reqs > 0 + and bool(torch.all(query_lens == uniform_decode_query_len)) + ): + attn_num_reqs = min( + cdiv(num_tokens_after_padding, uniform_decode_query_len), + self.max_num_reqs, + ) + else: + attn_num_reqs = min(num_tokens_after_padding, self.max_num_reqs) attn_num_tokens = num_tokens_after_padding if attn_num_reqs > num_reqs: # fill 0 for padded block-table rows.