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 12, 12:15 UTC
[Model Runner V2] Minor code cleanup (#29570)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 26, 2025 at 19:47 UTC
0aeb698b774e2d8593b14988e3af9ebbdd773730
3 files changed
+18
-18
vllm/v1/worker/gpu/cudagraph_utils.py
+2
-9
@@ -16,6 +16,7 @@ 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
19
+from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
20
from vllm.v1.worker.gpu.input_batch import InputBuffers
21
22
@@ -127,15 +128,7 @@ class CudaGraphManager:
128
slot_mappings=slot_mappings,
129
kv_cache_config=kv_cache_config,
130
)
130
- if self.dp_size > 1:
131
- num_tokens_across_dp = torch.full(
132
- (self.dp_size,),
133
- batch_size,
134
- dtype=torch.int32,
135
- device="cpu",
136
- )
137
- else:
138
- num_tokens_across_dp = None
131
+ num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, batch_size)
132
133
# Warm up.
134
with set_forward_context(
vllm/v1/worker/gpu/dp_utils.py
+9
@@ -20,3 +20,12 @@ def get_batch_metadata_across_dp(
20
tensor[1][dp_rank] = cudagraph_size
21
dist.all_reduce(tensor, group=group)
22
return tensor[0], tensor[1]
23
+
24
+
25
+def make_num_tokens_across_dp(
26
+ dp_size: int,
27
+ num_tokens: int,
28
+) -> torch.Tensor | None:
29
+ if dp_size == 1:
30
+ return None
31
+ return torch.full((dp_size,), num_tokens, dtype=torch.int32, device="cpu")
vllm/v1/worker/gpu/model_runner.py
+7
-9
@@ -35,7 +35,10 @@ from vllm.v1.worker.gpu.attn_utils import (
35
)
36
from vllm.v1.worker.gpu.block_table import BlockTables
37
from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager
38
-from vllm.v1.worker.gpu.dp_utils import get_batch_metadata_across_dp
38
+from vllm.v1.worker.gpu.dp_utils import (
39
+ get_batch_metadata_across_dp,
40
+ make_num_tokens_across_dp,
41
+)
42
from vllm.v1.worker.gpu.input_batch import (
43
InputBatch,
44
InputBuffers,
@@ -255,12 +258,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
258
if not skip_attn:
259
self.prepare_dummy_attn_metadata(input_batch)
260
258
- if self.dp_size == 1:
259
- num_tokens_across_dp: torch.Tensor | None = None
260
- else:
261
- num_tokens_across_dp = torch.full(
262
- (self.dp_size,), num_tokens, dtype=torch.int32, device="cpu"
263
- )
261
+ num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)
262
num_sampled_tokens = np.ones(input_batch.num_reqs, dtype=np.int32)
263
with (
264
self.maybe_dummy_run_with_lora(
@@ -816,7 +814,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
814
self.req_states.last_sampled_tokens,
815
next_prefill_tokens,
816
)
819
- self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
817
return draft_tokens
818
819
def get_cudagraph_and_dp_padding(
@@ -1006,7 +1003,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1003
input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected
1004
)
1005
if self.do_spec_decode:
1009
- _ = self.propose_draft(
1006
+ draft_tokens = self.propose_draft(
1007
input_batch,
1008
sampling_metadata,
1009
hidden_states,
@@ -1014,6 +1011,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1011
num_sampled,
1012
num_rejected,
1013
)
1014
+ self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
1015
1016
if self.use_async_scheduling:
1017
return async_output