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 9, 23:00 UTC
[Model Runner V2] Refactor `dummy_run` (#32533)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 19, 2026 at 14:50 UTC
43fada536005389182c1383142a0c5d43c434bd5
3 files changed
+34
-55
vllm/v1/worker/gpu/cudagraph_utils.py
+2
-3
@@ -13,7 +13,6 @@ from vllm.config.compilation import CUDAGraphMode
13
from vllm.distributed.parallel_state import graph_capture, is_global_first_rank
14
from vllm.forward_context import set_forward_context
15
from vllm.v1.attention.backend import AttentionMetadataBuilder
16
-from vllm.v1.core.sched.output import SchedulerOutput
16
from vllm.v1.kv_cache_interface import KVCacheConfig
17
from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
18
from vllm.v1.worker.gpu.block_table import BlockTables
@@ -60,12 +59,12 @@ class CudaGraphManager:
59
60
def get_cudagraph_size(
61
self,
63
- scheduler_output: SchedulerOutput,
62
num_tokens_after_padding: int,
63
+ num_tokens_per_request: Iterable[int],
64
) -> int | None:
65
return get_cudagraph_size(
66
num_tokens_after_padding,
68
- scheduler_output.num_scheduled_tokens.values(),
67
+ num_tokens_per_request,
68
self.cudagraph_sizes,
69
self.cudagraph_mode,
70
)
vllm/v1/worker/gpu/model_runner.py
+31
-46
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import gc
4
import time
5
+from collections.abc import Iterable
6
from copy import deepcopy
7
from typing import Any
8
@@ -288,47 +289,25 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
289
skip_attn: bool = True,
290
**kwargs,
291
) -> tuple[torch.Tensor, torch.Tensor]:
292
+ # Create a dummy scheduler output.
293
num_reqs = min(num_tokens, self.max_num_reqs)
292
- input_batch = InputBatch.make_dummy(
293
- num_reqs=num_reqs,
294
- num_tokens=num_tokens,
295
- input_buffers=self.input_buffers,
296
- device=self.device,
294
+ num_tokens_per_request = [num_tokens // num_reqs] * num_reqs
295
+ num_tokens_per_request[-1] += num_tokens % num_reqs
296
+ assert sum(num_tokens_per_request) == num_tokens
297
+ num_scheduled_tokens = {
298
+ f"_dummy_req_{i}": num_tokens_per_request[i] for i in range(num_reqs)
299
+ }
300
+ dummy_scheduler_output = SchedulerOutput.make_empty()
301
+ dummy_scheduler_output.total_num_scheduled_tokens = num_tokens
302
+ dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens
303
+
304
+ # Execute the model.
305
+ self.execute_model(
306
+ dummy_scheduler_output, dummy_run=True, skip_attn_for_dummy_run=skip_attn
307
)
298
- if self.uses_mrope:
299
- input_batch.mrope_positions = self.mrope_states.mrope_positions[
300
- :, :num_tokens
301
- ]
302
- if self.supports_mm_inputs:
303
- input_batch.inputs_embeds = self.encoder_runner.inputs_embeds[:num_tokens]
304
- if not skip_attn:
305
- self.prepare_dummy_attn_metadata(input_batch)
306
-
307
- dp_size = self.parallel_config.data_parallel_size
308
- num_tokens_across_dp = make_num_tokens_across_dp(dp_size, num_tokens)
309
- num_sampled_tokens = np.ones(input_batch.num_reqs, dtype=np.int32)
310
- positions = input_batch.positions
311
- if self.uses_mrope:
312
- positions = input_batch.mrope_positions
313
- with (
314
- self.maybe_dummy_run_with_lora(
315
- self.lora_config,
316
- input_batch.num_scheduled_tokens,
317
- num_sampled_tokens,
318
- ),
319
- set_forward_context(
320
- input_batch.attn_metadata,
321
- self.vllm_config,
322
- num_tokens=num_tokens,
323
- num_tokens_across_dp=num_tokens_across_dp,
324
- ),
325
- ):
326
- hidden_states = self.model(
327
- input_ids=input_batch.input_ids,
328
- positions=positions,
329
- inputs_embeds=input_batch.inputs_embeds,
330
- )
331
- sample_hidden_states = hidden_states[input_batch.logits_indices]
308
+ assert self.execute_model_state is not None
309
+ hidden_states, input_batch = self.execute_model_state
310
+ sample_hidden_states = hidden_states[input_batch.logits_indices]
311
return hidden_states, sample_hidden_states
312
313
@torch.inference_mode()
@@ -893,9 +872,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
872
873
def get_cudagraph_and_dp_padding(
874
self,
896
- scheduler_output: SchedulerOutput,
875
+ total_num_scheduled_tokens: int,
876
+ num_tokens_per_request: Iterable[int],
877
) -> tuple[CUDAGraphMode, int, torch.Tensor | None]:
898
- total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
878
dp_size = self.parallel_config.data_parallel_size
879
if dp_size == 1:
880
# No DP. Only consider CUDA graphs.
@@ -904,7 +883,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
883
return CUDAGraphMode.NONE, 0, None
884
885
cudagraph_size = self.cudagraph_manager.get_cudagraph_size(
907
- scheduler_output, total_num_scheduled_tokens
886
+ total_num_scheduled_tokens, num_tokens_per_request
887
)
888
if cudagraph_size is not None:
889
# Use full CUDA graph.
@@ -919,7 +898,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
898
cudagraph_size_before_dp: int | None = 0
899
else:
900
cudagraph_size_before_dp = self.cudagraph_manager.get_cudagraph_size(
922
- scheduler_output, total_num_scheduled_tokens
901
+ total_num_scheduled_tokens, num_tokens_per_request
902
)
903
if cudagraph_size_before_dp is None:
904
cudagraph_size_before_dp = -1
@@ -951,6 +930,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
930
scheduler_output: SchedulerOutput,
931
intermediate_tensors: Any | None = None,
932
dummy_run: bool = False,
933
+ skip_attn_for_dummy_run: bool = False,
934
) -> ModelRunnerOutput | None:
935
assert intermediate_tensors is None
936
if not dummy_run:
@@ -965,7 +945,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
945
return EMPTY_MODEL_RUNNER_OUTPUT
946
947
cudagraph_mode, num_tokens_after_padding, num_tokens_across_dp = (
968
- self.get_cudagraph_and_dp_padding(scheduler_output)
948
+ self.get_cudagraph_and_dp_padding(
949
+ scheduler_output.total_num_scheduled_tokens,
950
+ scheduler_output.num_scheduled_tokens.values(),
951
+ )
952
)
953
if num_tokens_after_padding == 0:
954
# All DP ranks have zero tokens to run.
@@ -999,7 +982,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
982
: input_batch.num_tokens_after_padding
983
]
984
else:
1002
- # No actual tokens to run. A dummy run for DP.
985
+ # No actual tokens to run. A dummy run for DP or memory profiling.
986
num_reqs = min(num_tokens_after_padding, self.max_num_reqs)
987
input_batch = InputBatch.make_dummy(
988
num_reqs=num_reqs,
@@ -1011,7 +994,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
994
input_batch.mrope_positions = self.mrope_states.mrope_positions[
995
:, :num_tokens_after_padding
996
]
1014
- self.prepare_dummy_attn_metadata(input_batch)
997
+ if not skip_attn_for_dummy_run:
998
+ self.prepare_dummy_attn_metadata(input_batch)
999
+ # FIXME(woosuk): Fix warmup for LoRA.
1000
1001
# Run model.
1002
if cudagraph_mode == CUDAGraphMode.FULL:
vllm/v1/worker/gpu_worker.py
+1
-6
@@ -662,12 +662,7 @@ class Worker(WorkerBase):
662
self.profiler.stop()
663
664
def execute_dummy_batch(self) -> None:
665
- if self.use_v2_model_runner:
666
- self.model_runner.execute_model(
667
- SchedulerOutput.make_empty(), dummy_run=True
668
- )
669
- else:
670
- self.model_runner._dummy_run(1, uniform_decode=True)
665
+ self.model_runner._dummy_run(1, uniform_decode=True)
666
667
def add_lora(self, lora_request: LoRARequest) -> bool:
668
return self.model_runner.add_lora(lora_request)