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 7, 10:15 UTC
[Model Runner V2][Perf] align dummy_run tokens to uniform decode for dp cudagraph (#35376)
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>
zhrrr committed
Mar 3, 2026 at 09:10 UTC
168ee03e1cbba2b962adbc704b16762b266be184
1 file changed
+18
-4
vllm/v1/worker/gpu/model_runner.py
+18
-4
@@ -39,6 +39,7 @@ from vllm.model_executor.model_loader import get_model_loader
39
from vllm.multimodal import MULTIMODAL_REGISTRY
40
from vllm.sequence import IntermediateTensors
41
from vllm.tasks import SupportedTask
42
+from vllm.utils.math_utils import cdiv
43
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
44
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
45
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -327,12 +328,25 @@ class GPUModelRunner(LoRAModelRunnerMixin):
328
329
@torch.inference_mode()
330
def _dummy_run(
330
- self, num_tokens: int, *args, skip_attn: bool = True, **kwargs
331
+ self,
332
+ num_tokens: int,
333
+ *args,
334
+ skip_attn: bool = True,
335
+ uniform_decode: bool = False,
336
+ **kwargs,
337
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
338
# Create a dummy scheduler output.
333
- num_reqs = min(num_tokens, self.max_num_reqs)
334
- num_tokens_per_request = [num_tokens // num_reqs] * num_reqs
335
- num_tokens_per_request[-1] += num_tokens % num_reqs
339
+ if uniform_decode:
340
+ # Align tokens to uniform_decode_query_len for cudagraph
341
+ # compatibility across DP ranks.
342
+ query_len = self.cudagraph_manager.uniform_decode_query_len
343
+ num_reqs = min(cdiv(num_tokens, query_len), self.max_num_reqs)
344
+ num_tokens = num_reqs * query_len
345
+ num_tokens_per_request = [query_len] * num_reqs
346
+ else:
347
+ num_reqs = min(num_tokens, self.max_num_reqs)
348
+ num_tokens_per_request = [num_tokens // num_reqs] * num_reqs
349
+ num_tokens_per_request[-1] += num_tokens % num_reqs
350
assert sum(num_tokens_per_request) == num_tokens
351
num_scheduled_tokens = {
352
f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request)