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, 18:30 UTC
[Chore] Fix pre-commit error after #25266 (#29190)
Woosuk Kwon committed
Nov 21, 2025 at 10:21 UTC
1bed891f72a6cbd32c0c75dfaa29ad21d7a68b75
5 files changed
+40
-24
vllm/v1/worker/gpu/async_utils.py
+11
-9
@@ -7,6 +7,7 @@ import torch
7
8
from vllm.v1.outputs import (
9
AsyncModelRunnerOutput,
10
+ LogprobsTensors,
11
ModelRunnerOutput,
12
SamplerOutput,
13
)
@@ -46,15 +47,18 @@ class AsyncOutput(AsyncModelRunnerOutput):
47
"cpu", non_blocking=True
48
)
49
if sampler_output.logprobs_tensors is not None:
49
- self.logprobs_tensors = (
50
+ self.logprobs_tensors: LogprobsTensors | None = (
51
sampler_output.logprobs_tensors.to_cpu_nonblocking()
52
)
53
else:
54
self.logprobs_tensors = None
54
- self.prompt_logprobs_dict = {}
55
+ self.prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}
56
if self.model_runner_output.prompt_logprobs_dict:
57
for k, v in self.model_runner_output.prompt_logprobs_dict.items():
57
- self.prompt_logprobs_dict[k] = v.to_cpu_nonblocking()
58
+ if v is not None:
59
+ self.prompt_logprobs_dict[k] = v.to_cpu_nonblocking()
60
+ else:
61
+ self.prompt_logprobs_dict[k] = None
62
self.copy_event.record(self.copy_stream)
63
64
def get_output(self) -> ModelRunnerOutput:
@@ -64,12 +68,10 @@ class AsyncOutput(AsyncModelRunnerOutput):
68
# the existing model runner.
69
# Going forward, we should keep the data structures as NumPy arrays
70
# rather than Python lists.
67
- sampled_token_ids_np = self.sampled_token_ids.numpy()
68
- num_reqs = sampled_token_ids_np.shape[0]
69
- sampled_token_ids: list[np.ndarray] = [
70
- sampled_token_ids_np[i, : self.num_sampled_tokens[i]]
71
- for i in range(num_reqs)
72
- ]
71
+ sampled_token_ids: list[list[int]] = self.sampled_token_ids.tolist()
72
+ num_reqs = len(sampled_token_ids)
73
+ for i in range(num_reqs):
74
+ del sampled_token_ids[i][self.num_sampled_tokens[i] :]
75
self.model_runner_output.sampled_token_ids = sampled_token_ids
76
77
if self.logprobs_tensors is not None:
vllm/v1/worker/gpu/attn_utils.py
+8
-6
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Sequence
4
-from typing import Any
4
+from typing import Any, cast
5
6
import torch
7
@@ -13,6 +13,7 @@ from vllm.v1.attention.backends.utils import (
13
CommonAttentionMetadata,
14
)
15
from vllm.v1.kv_cache_interface import (
16
+ AttentionSpec,
17
KVCacheConfig,
18
KVCacheSpec,
19
)
@@ -22,7 +23,8 @@ from vllm.v1.worker.utils import bind_kv_cache
23
24
def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
25
kv_cache_spec: dict[str, KVCacheSpec] = {}
25
- attn_layers = get_layers_from_vllm_config(vllm_config, AttentionLayerBase)
26
+ layer_type = cast(type[Any], AttentionLayerBase)
27
+ attn_layers = get_layers_from_vllm_config(vllm_config, layer_type)
28
for layer_name, attn_module in attn_layers.items():
29
# Skip modules that don't need KV cache (eg encoder-only attention)
30
if spec := attn_module.get_kv_cache_spec(vllm_config):
@@ -35,16 +37,15 @@ def init_attn_backend(
37
vllm_config: VllmConfig,
38
device: torch.device,
39
):
38
- attn_backends: dict[str, AttentionBackend] = {}
40
+ attn_backends: dict[str, type[AttentionBackend]] = {}
41
attn_metadata_builders: list[AttentionMetadataBuilder] = []
42
flashinfer_workspace: torch.Tensor | None = None
43
for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
44
layer_names = kv_cache_group_spec.layer_names
45
any_layer_name = next(iter(layer_names))
46
45
- attn_layers = get_layers_from_vllm_config(
46
- vllm_config, AttentionLayerBase, layer_names
47
- )
47
+ layer_type = cast(type[Any], AttentionLayerBase)
48
+ attn_layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)
49
attn_backend = attn_layers[any_layer_name].get_attn_backend()
50
for layer_name in layer_names:
51
attn_backends[layer_name] = attn_backend
@@ -93,6 +94,7 @@ def _reshape_kv_cache(
94
kv_caches: dict[str, torch.Tensor] = {}
95
for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
96
kv_cache_spec = kv_cache_group_spec.kv_cache_spec
97
+ assert isinstance(kv_cache_spec, AttentionSpec)
98
for layer_name in kv_cache_group_spec.layer_names:
99
raw_tensor = kv_cache_raw_tensors[layer_name]
100
assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
vllm/v1/worker/gpu/cudagraph_utils.py
+10
-2
@@ -34,8 +34,16 @@ class CudaGraphManager:
34
self.compilation_config = vllm_config.compilation_config
35
assert self.compilation_config is not None
36
37
- self.cudagraph_mode = self.compilation_config.cudagraph_mode
38
- self.cudagraph_sizes = sorted(self.compilation_config.cudagraph_capture_sizes)
37
+ if self.compilation_config.cudagraph_mode is None:
38
+ self.cudagraph_mode = CUDAGraphMode.NONE
39
+ else:
40
+ self.cudagraph_mode = self.compilation_config.cudagraph_mode
41
+ if self.compilation_config.cudagraph_capture_sizes is not None:
42
+ self.cudagraph_sizes = sorted(
43
+ self.compilation_config.cudagraph_capture_sizes
44
+ )
45
+ else:
46
+ self.cudagraph_sizes = []
47
self.padded_sizes = self._init_padded_sizes()
48
49
self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
vllm/v1/worker/gpu/model_runner.py
+10
-6
@@ -329,8 +329,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
329
torch.cuda.synchronize()
330
331
def update_states(self, scheduler_output: SchedulerOutput) -> None:
332
- for req_id in scheduler_output.preempted_req_ids:
333
- self.req_states.remove_request(req_id)
332
+ if scheduler_output.preempted_req_ids is not None:
333
+ for req_id in scheduler_output.preempted_req_ids:
334
+ self.req_states.remove_request(req_id)
335
for req_id in scheduler_output.finished_req_ids:
336
self.req_states.remove_request(req_id)
337
@@ -346,6 +347,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
347
348
# Add new requests.
349
for new_req_data in scheduler_output.scheduled_new_reqs:
350
+ assert new_req_data.prompt_token_ids is not None
351
+ assert new_req_data.prefill_token_ids is not None
352
+ assert new_req_data.sampling_params is not None
353
req_id = new_req_data.req_id
354
self.req_states.add_request(
355
req_id=req_id,
@@ -398,8 +402,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
402
# Decode first, then prefill.
403
# batch_idx -> req_id
404
req_ids = sorted(
401
- scheduler_output.num_scheduled_tokens,
402
- key=scheduler_output.num_scheduled_tokens.get,
405
+ scheduler_output.num_scheduled_tokens.keys(),
406
+ key=lambda k: scheduler_output.num_scheduled_tokens[k],
407
)
408
num_scheduled_tokens = np.array(
409
[scheduler_output.num_scheduled_tokens[i] for i in req_ids], dtype=np.int32
@@ -637,9 +641,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
641
model_runner_output = ModelRunnerOutput(
642
req_ids=input_batch.req_ids,
643
req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
640
- sampled_token_ids=None,
644
+ sampled_token_ids=None, # type: ignore
645
logprobs=None,
642
- prompt_logprobs_dict=prompt_logprobs_dict,
646
+ prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore
647
pooler_output=[],
648
kv_connector_output=None,
649
num_nans_in_logits=None,
vllm/v1/worker/gpu/sampler.py
+1
-1
@@ -8,8 +8,8 @@ import triton.language as tl
8
9
from vllm.config.model import LogprobsMode
10
from vllm.v1.outputs import LogprobsTensors, SamplerOutput
11
-from vllm.v1.sample.metadata import SamplingMetadata
11
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
12
+from vllm.v1.worker.gpu.states import SamplingMetadata
13
14
15
class Sampler: