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, 11:15 UTC
[ModelRunner V2] Misc code simplification and cleanup (#33266)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jan 28, 2026 at 14:41 UTC
6bf3b46d7840364907405c7b02eaa66886a90839
3 files changed
+40
-84
vllm/v1/worker/gpu/async_utils.py
+16
-37
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from contextlib import contextmanager
3
4
import numpy as np
5
import torch
@@ -28,33 +27,27 @@ class AsyncOutput(AsyncModelRunnerOutput):
27
self.model_runner_output = model_runner_output
28
self.sampler_output = sampler_output
29
self.num_sampled_tokens = num_sampled_tokens
31
- self.copy_stream = copy_stream
30
self.copy_event = copy_event
31
32
default_stream = torch.cuda.current_stream()
35
- with torch.cuda.stream(self.copy_stream):
36
- self.copy_stream.wait_stream(default_stream)
33
+ with torch.cuda.stream(copy_stream):
34
+ copy_stream.wait_stream(default_stream)
35
36
self.sampled_token_ids = async_copy_to_np(sampler_output.sampled_token_ids)
37
+ self.logprobs_tensors: LogprobsTensors | None = None
38
if sampler_output.logprobs_tensors is not None:
40
- self.logprobs_tensors: LogprobsTensors | None = (
39
+ self.logprobs_tensors = (
40
sampler_output.logprobs_tensors.to_cpu_nonblocking()
41
)
43
- else:
44
- self.logprobs_tensors = None
42
+ self.num_nans: np.ndarray | None = None
43
if sampler_output.num_nans is not None:
44
self.num_nans = async_copy_to_np(sampler_output.num_nans)
47
- else:
48
- self.num_nans = None
45
self.num_sampled_tokens_np = async_copy_to_np(num_sampled_tokens)
50
- self.prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}
51
- if self.model_runner_output.prompt_logprobs_dict:
52
- for k, v in self.model_runner_output.prompt_logprobs_dict.items():
53
- if v is not None:
54
- self.prompt_logprobs_dict[k] = v.to_cpu_nonblocking()
55
- else:
56
- self.prompt_logprobs_dict[k] = None
57
- self.copy_event.record(self.copy_stream)
46
+ self.prompt_logprobs_dict = {
47
+ k: v.to_cpu_nonblocking() if v is not None else None
48
+ for k, v in self.model_runner_output.prompt_logprobs_dict.items()
49
+ }
50
+ self.copy_event.record(copy_stream)
51
52
def get_output(self) -> ModelRunnerOutput:
53
self.copy_event.synchronize()
@@ -64,18 +57,15 @@ class AsyncOutput(AsyncModelRunnerOutput):
57
# Going forward, we should keep the data structures as NumPy arrays
58
# rather than Python lists.
59
sampled_token_ids: list[list[int]] = self.sampled_token_ids.tolist()
67
- num_reqs = len(sampled_token_ids)
68
- num_sampled_tokens = self.num_sampled_tokens_np.tolist()
69
- for i in range(num_reqs):
70
- del sampled_token_ids[i][num_sampled_tokens[i] :]
60
+ num_sampled_tokens: list[int] = self.num_sampled_tokens_np.tolist()
61
+ for token_ids, num_tokens in zip(sampled_token_ids, num_sampled_tokens):
62
+ del token_ids[num_tokens:]
63
self.model_runner_output.sampled_token_ids = sampled_token_ids
64
65
if self.num_nans is not None:
74
- num_nans = self.num_nans.tolist()
75
- self.model_runner_output.num_nans_in_logits = {
76
- req_id: num_nans[i]
77
- for i, req_id in enumerate(self.model_runner_output.req_ids)
78
- }
66
+ self.model_runner_output.num_nans_in_logits = dict(
67
+ zip(self.model_runner_output.req_ids, self.num_nans.tolist())
68
+ )
69
70
if self.logprobs_tensors is not None:
71
self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
@@ -83,16 +73,5 @@ class AsyncOutput(AsyncModelRunnerOutput):
73
return self.model_runner_output
74
75
86
-@contextmanager
87
-def async_barrier(event: torch.cuda.Event | None):
88
- if event is not None:
89
- event.synchronize()
90
- try:
91
- yield
92
- finally:
93
- if event is not None:
94
- event.record()
95
-
96
-
76
def async_copy_to_np(x: torch.Tensor) -> np.ndarray:
77
return x.to("cpu", non_blocking=True).numpy()
vllm/v1/worker/gpu/model_runner.py
+21
-47
@@ -123,12 +123,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
123
self.use_async_scheduling = self.scheduler_config.async_scheduling
124
self.output_copy_stream = torch.cuda.Stream(self.device)
125
self.output_copy_event = torch.cuda.Event()
126
- if self.use_async_scheduling:
127
- self.input_prep_event = torch.cuda.Event()
128
- self.structured_outputs_event = torch.cuda.Event()
129
- else:
130
- self.input_prep_event = None
131
- self.structured_outputs_event = None
126
127
if self.speculative_config is not None:
128
self.do_spec_decode = True
@@ -179,7 +173,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
173
self.max_model_len = max_model_len
174
self.req_states.max_model_len = max_model_len
175
182
- def get_supported_tasks(self) -> tuple[str]:
176
+ @staticmethod
177
+ def get_supported_tasks() -> tuple[str]:
178
return ("generate",)
179
180
def load_model(self, *args, **kwargs) -> None:
@@ -194,9 +189,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
189
)
190
if self.lora_config:
191
self.model = self.load_lora_model(
197
- self.model,
198
- self.vllm_config,
199
- self.device,
192
+ self.model, self.vllm_config, self.device
193
)
194
if self.do_spec_decode:
195
self.speculator.load_model(self.model)
@@ -238,9 +231,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
231
)
232
233
self.attn_backends, self.attn_metadata_builders = init_attn_backend(
241
- self.kv_cache_config,
242
- self.vllm_config,
243
- self.device,
234
+ self.kv_cache_config, self.vllm_config, self.device
235
)
236
if self.do_spec_decode:
237
# HACK(woosuk)
@@ -288,11 +279,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
279
280
@torch.inference_mode()
281
def _dummy_run(
291
- self,
292
- num_tokens: int,
293
- *args,
294
- skip_attn: bool = True,
295
- **kwargs,
282
+ self, num_tokens: int, *args, skip_attn: bool = True, **kwargs
283
) -> tuple[torch.Tensor, torch.Tensor]:
284
# Create a dummy scheduler output.
285
num_reqs = min(num_tokens, self.max_num_reqs)
@@ -320,10 +307,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
307
return hidden_states, sample_hidden_states
308
309
@torch.inference_mode()
323
- def _dummy_sampler_run(
324
- self,
325
- hidden_states: torch.Tensor,
326
- ) -> None:
310
+ def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
311
num_reqs = hidden_states.shape[0]
312
logits = self.model.compute_logits(hidden_states)
313
idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=self.device)
@@ -337,8 +321,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
321
@torch.inference_mode()
322
def profile_run(self) -> None:
323
hidden_states, sample_hidden_states = self._dummy_run(
340
- self.max_num_tokens,
341
- skip_attn=True,
324
+ self.max_num_tokens, skip_attn=True
325
)
326
self._dummy_sampler_run(sample_hidden_states)
327
if self.do_spec_decode:
@@ -482,11 +465,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
465
466
def update_requests(self, scheduler_output: SchedulerOutput) -> None:
467
# Add new blocks for the existing requests.
485
- cached_reqs = scheduler_output.scheduled_cached_reqs
486
- for i, req_id in enumerate(cached_reqs.req_ids):
487
- req_index = self.req_states.req_id_to_index[req_id]
488
- req_new_block_ids = cached_reqs.new_block_ids[i]
468
+ reqs = scheduler_output.scheduled_cached_reqs
469
+ for req_new_block_ids, req_id in zip(reqs.new_block_ids, reqs.req_ids):
470
if req_new_block_ids is not None:
471
+ req_index = self.req_states.req_id_to_index[req_id]
472
self.block_tables.append_block_ids(
473
req_index, req_new_block_ids, overwrite=False
474
)
@@ -517,7 +499,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
499
idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
500
501
# Get the number of draft tokens for each request.
520
- if not scheduler_output.scheduled_spec_decode_tokens:
502
+ draft_tokens = scheduler_output.scheduled_spec_decode_tokens
503
+ if not draft_tokens:
504
# No draft token scheduled (common case).
505
total_num_draft_tokens = 0
506
total_num_logits = num_reqs
@@ -527,12 +510,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
510
)
511
expanded_idx_mapping = idx_mapping
512
else:
530
- draft_tokens = scheduler_output.scheduled_spec_decode_tokens
513
num_draft_tokens = np.array(
532
- [
533
- len(draft_tokens[req_id]) if req_id in draft_tokens else 0
534
- for req_id in req_ids
535
- ],
514
+ [len(draft_tokens.get(req_id, ())) for req_id in req_ids],
515
dtype=np.int32,
516
)
517
total_num_draft_tokens = int(num_draft_tokens.sum())
@@ -544,11 +523,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
523
np.cumsum(num_logits, out=cu_num_logits_np[1:])
524
cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)
525
526
+ max_expand_len = self.num_speculative_steps + 1
527
expanded_idx_mapping = expand_idx_mapping(
548
- idx_mapping,
549
- total_num_logits,
550
- cu_num_logits,
551
- max_expand_len=self.num_speculative_steps + 1,
528
+ idx_mapping, total_num_logits, cu_num_logits, max_expand_len
529
)
530
531
# Block tables: num_kv_cache_groups x [num_reqs, max_num_blocks]
@@ -640,9 +617,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
617
positions = self.input_buffers.positions[:num_tokens_after_padding]
618
mrope_positions = None
619
if self.uses_mrope:
643
- mrope_positions = self.mrope_states.mrope_positions[
644
- :, :num_tokens_after_padding
645
- ]
620
+ mrope_positions = self.mrope_states.mrope_positions
621
+ mrope_positions = mrope_positions[:, :num_tokens_after_padding]
622
return InputBatch(
623
req_ids=req_ids,
624
num_reqs=num_reqs,
@@ -762,10 +738,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
738
# Update the number of computed prefill tokens.
739
idx_mapping_np = input_batch.idx_mapping_np
740
computed_prefill = self.req_states.num_computed_prefill_tokens
765
- # TODO(woosuk): Simplify this.
766
- computed_prefill[idx_mapping_np] = np.minimum(
767
- computed_prefill[idx_mapping_np] + input_batch.num_scheduled_tokens,
768
- self.req_states.prefill_len.np[idx_mapping_np],
741
+ computed_prefill[idx_mapping_np] += input_batch.num_scheduled_tokens
742
+ np.minimum(
743
+ computed_prefill, self.req_states.prefill_len.np, out=computed_prefill
744
)
745
746
@torch.inference_mode()
@@ -834,8 +809,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
809
# Common case.
810
# Prepare all the inputs and copy to the input buffers.
811
input_batch = self.prepare_inputs(
837
- scheduler_output,
838
- num_tokens_after_padding,
812
+ scheduler_output, num_tokens_after_padding
813
)
814
if self.lora_config:
815
# Activate LoRA adapters.
vllm/v1/worker/gpu_worker.py
+3
@@ -107,6 +107,9 @@ class Worker(WorkerBase):
107
108
self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
109
110
+ if self.use_v2_model_runner:
111
+ logger.info_once("Using V2 Model Runner", scope="global")
112
+
113
def sleep(self, level: int = 1) -> None:
114
from vllm.device_allocator.cumem import CuMemAllocator
115