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
GPU Model Runner V2 (#25266)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 21, 2025 at 08:20 UTC
30b44a1598ea62fd3dcfd0d72a799ca4685e829e
18 files changed
+2639
-12
.github/CODEOWNERS
+3
@@ -35,6 +35,9 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
35
/vllm/v1/kv_cache_interface.py @heheda12345
36
/vllm/v1/offloading @ApostaC
37
38
+# Model runner V2
39
+/vllm/v1/worker/gpu @WoosukKwon
40
+
41
# Test ownership
42
/.buildkite/lm-eval-harness @mgoin
43
/tests/distributed/test_multi_node_assignment.py @youkaichao
vllm/envs.py
+5
@@ -231,6 +231,7 @@ if TYPE_CHECKING:
231
VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
232
VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
233
VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
234
+ VLLM_USE_V2_MODEL_RUNNER: bool = False
235
236
237
def get_default_cache_root():
@@ -1522,6 +1523,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1523
"VLLM_COMPILE_CACHE_SAVE_FORMAT": env_with_choices(
1524
"VLLM_COMPILE_CACHE_SAVE_FORMAT", "binary", ["binary", "unpacked"]
1525
),
1526
+ # Flag to enable v2 model runner.
1527
+ "VLLM_USE_V2_MODEL_RUNNER": lambda: bool(
1528
+ int(os.getenv("VLLM_USE_V2_MODEL_RUNNER", "0"))
1529
+ ),
1530
}
1531
1532
# --8<-- [end:env-vars-definition]
vllm/v1/attention/backends/flashinfer.py
+3
@@ -593,6 +593,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
593
)
594
return self._workspace_buffer
595
596
+ def set_workspace_buffer(self, workspace_buffer: torch.Tensor):
597
+ self._workspace_buffer = workspace_buffer
598
+
599
def _get_prefill_wrapper(
600
self,
601
) -> BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper:
vllm/v1/core/sched/output.py
+24
@@ -44,11 +44,15 @@ class NewRequestData:
44
lora_request: LoRARequest | None
45
prompt_embeds: "torch.Tensor | None" = None
46
47
+ # Only used for v2 model runner.
48
+ prefill_token_ids: list[int] | None = None
49
+
50
@classmethod
51
def from_request(
52
cls,
53
request: Request,
54
block_ids: tuple[list[int], ...],
55
+ prefill_token_ids: list[int] | None = None,
56
) -> "NewRequestData":
57
return cls(
58
req_id=request.request_id,
@@ -60,6 +64,7 @@ class NewRequestData:
64
num_computed_tokens=request.num_computed_tokens,
65
lora_request=request.lora_request,
66
prompt_embeds=request.prompt_embeds,
67
+ prefill_token_ids=prefill_token_ids,
68
)
69
70
def __repr__(self) -> str:
@@ -68,6 +73,7 @@ class NewRequestData:
73
f"NewRequestData("
74
f"req_id={self.req_id},"
75
f"prompt_token_ids={self.prompt_token_ids},"
76
+ f"prefill_token_ids={self.prefill_token_ids},"
77
f"mm_features={self.mm_features},"
78
f"sampling_params={self.sampling_params},"
79
f"block_ids={self.block_ids},"
@@ -183,6 +189,10 @@ class SchedulerOutput:
189
# freed from the encoder cache.
190
free_encoder_mm_hashes: list[str]
191
192
+ # Request IDs that are preempted in this step.
193
+ # Only used for v2 model runner.
194
+ preempted_req_ids: set[str] | None = None
195
+
196
# Whether the scheduled requests have all the output tokens they
197
# need to perform grammar bitmask computation.
198
pending_structured_output_tokens: bool = False
@@ -193,6 +203,20 @@ class SchedulerOutput:
203
# EC Cache Connector metadata
204
ec_connector_metadata: ECConnectorMetadata | None = None
205
206
+ @classmethod
207
+ def make_empty(cls) -> "SchedulerOutput":
208
+ return cls(
209
+ scheduled_new_reqs=[],
210
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
211
+ num_scheduled_tokens={},
212
+ total_num_scheduled_tokens=0,
213
+ scheduled_spec_decode_tokens={},
214
+ scheduled_encoder_inputs={},
215
+ num_common_prefix_blocks=[],
216
+ finished_req_ids=set(),
217
+ free_encoder_mm_hashes=[],
218
+ )
219
+
220
221
@dataclass
222
class GrammarOutput:
vllm/v1/core/sched/scheduler.py
+22
-6
@@ -6,6 +6,7 @@ from collections import defaultdict
6
from collections.abc import Iterable
7
from typing import Any
8
9
+from vllm import envs
10
from vllm.config import VllmConfig
11
from vllm.distributed.ec_transfer.ec_connector.base import (
12
ECConnectorMetadata,
@@ -187,6 +188,7 @@ class Scheduler(SchedulerInterface):
188
pcp_world_size=self.pcp_world_size,
189
)
190
self.use_pp = self.parallel_config.pipeline_parallel_size > 1
191
+ self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
192
193
def schedule(self) -> SchedulerOutput:
194
# NOTE(woosuk) on the scheduling algorithm:
@@ -658,12 +660,25 @@ class Scheduler(SchedulerInterface):
660
)
661
662
# Construct the scheduler output.
661
- new_reqs_data = [
662
- NewRequestData.from_request(
663
- req, req_to_new_blocks[req.request_id].get_block_ids()
664
- )
665
- for req in scheduled_new_reqs
666
- ]
663
+ if self.use_v2_model_runner:
664
+ scheduled_new_reqs = scheduled_new_reqs + scheduled_resumed_reqs
665
+ scheduled_resumed_reqs = []
666
+ new_reqs_data = [
667
+ NewRequestData.from_request(
668
+ req,
669
+ req_to_new_blocks[req.request_id].get_block_ids(),
670
+ req._all_token_ids,
671
+ )
672
+ for req in scheduled_new_reqs
673
+ ]
674
+ else:
675
+ new_reqs_data = [
676
+ NewRequestData.from_request(
677
+ req, req_to_new_blocks[req.request_id].get_block_ids()
678
+ )
679
+ for req in scheduled_new_reqs
680
+ ]
681
+
682
with record_function_or_nullcontext("schedule: make_cached_request_data"):
683
cached_reqs_data = self._make_cached_request_data(
684
scheduled_running_reqs,
@@ -685,6 +700,7 @@ class Scheduler(SchedulerInterface):
700
scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
701
scheduled_encoder_inputs=scheduled_encoder_inputs,
702
num_common_prefix_blocks=num_common_prefix_blocks,
703
+ preempted_req_ids={req.request_id for req in preempted_reqs},
704
# finished_req_ids is an existing state in the scheduler,
705
# instead of being newly scheduled in this step.
706
# It contains the request IDs that are finished in between
vllm/v1/worker/gpu/README.md
new
+4
@@ -0,0 +1,4 @@
1
+# [Experimental] Model Runner V2
2
+
3
+This directory contains the new model runner which is under active development.
4
+Ping [Woosuk Kwon](https://github.com/WoosukKwon) for any changes.
vllm/v1/worker/gpu/__init__.py
vllm/v1/worker/gpu/async_utils.py
new
+89
@@ -0,0 +1,89 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from contextlib import contextmanager
4
+
5
+import numpy as np
6
+import torch
7
+
8
+from vllm.v1.outputs import (
9
+ AsyncModelRunnerOutput,
10
+ ModelRunnerOutput,
11
+ SamplerOutput,
12
+)
13
+
14
+
15
+class AsyncOutput(AsyncModelRunnerOutput):
16
+ def __init__(
17
+ self,
18
+ model_runner_output: ModelRunnerOutput,
19
+ sampler_output: SamplerOutput,
20
+ num_sampled_tokens: np.ndarray,
21
+ copy_stream: torch.cuda.Stream,
22
+ copy_event: torch.cuda.Event,
23
+ ):
24
+ self.model_runner_output = model_runner_output
25
+ self.sampler_output = sampler_output
26
+ self.num_sampled_tokens = num_sampled_tokens
27
+ self.copy_stream = copy_stream
28
+ self.copy_event = copy_event
29
+
30
+ default_stream = torch.cuda.current_stream()
31
+ with torch.cuda.stream(self.copy_stream):
32
+ self.copy_stream.wait_stream(default_stream)
33
+
34
+ # NOTE(woosuk): We must ensure that CPU tensors are not freed
35
+ # before the device-to-host copy is fully completed. For instance,
36
+ # operations like
37
+ # self.sampled_token_np = ...to("cpu", non_blocking=True).numpy()
38
+ # are unsafe because the underlying CPU tensor can be prematurely freed and
39
+ # reused by other tensors before the asynchronous copy finishes, potentially
40
+ # causing race conditions. To prevent this, we delay freeing by holding
41
+ # references until the copy event signals completion.
42
+ # Likewise, we also need to keep the reference to the GPU tensors.
43
+ # This is done by keeping the reference to sampler_output and
44
+ # model_runner_output.
45
+ self.sampled_token_ids = sampler_output.sampled_token_ids.to(
46
+ "cpu", non_blocking=True
47
+ )
48
+ if sampler_output.logprobs_tensors is not None:
49
+ self.logprobs_tensors = (
50
+ sampler_output.logprobs_tensors.to_cpu_nonblocking()
51
+ )
52
+ else:
53
+ self.logprobs_tensors = None
54
+ self.prompt_logprobs_dict = {}
55
+ if self.model_runner_output.prompt_logprobs_dict:
56
+ for k, v in self.model_runner_output.prompt_logprobs_dict.items():
57
+ self.prompt_logprobs_dict[k] = v.to_cpu_nonblocking()
58
+ self.copy_event.record(self.copy_stream)
59
+
60
+ def get_output(self) -> ModelRunnerOutput:
61
+ self.copy_event.synchronize()
62
+
63
+ # NOTE(woosuk): The following code is to ensure compatibility with
64
+ # the existing model runner.
65
+ # Going forward, we should keep the data structures as NumPy arrays
66
+ # 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
+ ]
73
+ self.model_runner_output.sampled_token_ids = sampled_token_ids
74
+
75
+ if self.logprobs_tensors is not None:
76
+ self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
77
+ self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict
78
+ return self.model_runner_output
79
+
80
+
81
+@contextmanager
82
+def async_barrier(event: torch.cuda.Event | None):
83
+ if event is not None:
84
+ event.synchronize()
85
+ try:
86
+ yield
87
+ finally:
88
+ if event is not None:
89
+ event.record()
vllm/v1/worker/gpu/attn_utils.py
new
+187
@@ -0,0 +1,187 @@
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
5
+
6
+import torch
7
+
8
+from vllm.attention.backends.abstract import AttentionBackend
9
+from vllm.config import VllmConfig, get_layers_from_vllm_config
10
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
11
+from vllm.v1.attention.backends.utils import (
12
+ AttentionMetadataBuilder,
13
+ CommonAttentionMetadata,
14
+)
15
+from vllm.v1.kv_cache_interface import (
16
+ KVCacheConfig,
17
+ KVCacheSpec,
18
+)
19
+from vllm.v1.utils import CpuGpuBuffer
20
+from vllm.v1.worker.utils import bind_kv_cache
21
+
22
+
23
+def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
24
+ kv_cache_spec: dict[str, KVCacheSpec] = {}
25
+ attn_layers = get_layers_from_vllm_config(vllm_config, AttentionLayerBase)
26
+ for layer_name, attn_module in attn_layers.items():
27
+ # Skip modules that don't need KV cache (eg encoder-only attention)
28
+ if spec := attn_module.get_kv_cache_spec(vllm_config):
29
+ kv_cache_spec[layer_name] = spec
30
+ return kv_cache_spec
31
+
32
+
33
+def init_attn_backend(
34
+ kv_cache_config: KVCacheConfig,
35
+ vllm_config: VllmConfig,
36
+ device: torch.device,
37
+):
38
+ attn_backends: dict[str, AttentionBackend] = {}
39
+ attn_metadata_builders: list[AttentionMetadataBuilder] = []
40
+ flashinfer_workspace: torch.Tensor | None = None
41
+ for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
42
+ layer_names = kv_cache_group_spec.layer_names
43
+ any_layer_name = next(iter(layer_names))
44
+
45
+ attn_layers = get_layers_from_vllm_config(
46
+ vllm_config, AttentionLayerBase, layer_names
47
+ )
48
+ attn_backend = attn_layers[any_layer_name].get_attn_backend()
49
+ for layer_name in layer_names:
50
+ attn_backends[layer_name] = attn_backend
51
+
52
+ attn_metadata_builder = attn_backend.get_builder_cls()(
53
+ kv_cache_group_spec.kv_cache_spec,
54
+ layer_names,
55
+ vllm_config,
56
+ device,
57
+ )
58
+ attn_metadata_builders.append(attn_metadata_builder) # type: ignore
59
+
60
+ if "FLASHINFER" in attn_backend.get_name():
61
+ if flashinfer_workspace is None:
62
+ flashinfer_workspace = attn_metadata_builder._get_workspace_buffer()
63
+ else:
64
+ attn_metadata_builder.set_workspace_buffer(flashinfer_workspace)
65
+ return attn_backends, attn_metadata_builders
66
+
67
+
68
+def _allocate_kv_cache(
69
+ kv_cache_config: KVCacheConfig,
70
+ device: torch.device,
71
+):
72
+ kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
73
+ for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
74
+ tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device)
75
+ for layer_name in kv_cache_tensor.shared_by:
76
+ kv_cache_raw_tensors[layer_name] = tensor
77
+
78
+ layer_names = set()
79
+ for group in kv_cache_config.kv_cache_groups:
80
+ for layer_name in group.layer_names:
81
+ layer_names.add(layer_name)
82
+ assert layer_names == set(kv_cache_raw_tensors.keys()), (
83
+ "Some layers are not correctly initialized"
84
+ )
85
+ return kv_cache_raw_tensors
86
+
87
+
88
+def _reshape_kv_cache(
89
+ kv_cache_config: KVCacheConfig,
90
+ kv_cache_raw_tensors: dict[str, torch.Tensor],
91
+ attn_backends: dict[str, AttentionBackend],
92
+) -> dict[str, torch.Tensor]:
93
+ kv_caches: dict[str, torch.Tensor] = {}
94
+ for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
95
+ kv_cache_spec = kv_cache_group_spec.kv_cache_spec
96
+ for layer_name in kv_cache_group_spec.layer_names:
97
+ raw_tensor = kv_cache_raw_tensors[layer_name]
98
+ assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
99
+ num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes
100
+
101
+ attn_backend = attn_backends[layer_name]
102
+ kv_cache_shape = attn_backend.get_kv_cache_shape(
103
+ num_blocks,
104
+ kv_cache_spec.block_size,
105
+ kv_cache_spec.num_kv_heads,
106
+ kv_cache_spec.head_size,
107
+ )
108
+
109
+ # FIXME(woosuk): Add kv_cache_stride_order to all attention backends.
110
+ try:
111
+ kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
112
+ assert len(kv_cache_stride_order) == len(kv_cache_shape)
113
+ except (AttributeError, NotImplementedError):
114
+ kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
115
+
116
+ kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
117
+ inv_order = [
118
+ kv_cache_stride_order.index(i)
119
+ for i in range(len(kv_cache_stride_order))
120
+ ]
121
+
122
+ dtype = kv_cache_spec.dtype
123
+ raw_tensor = raw_tensor.view(dtype)
124
+ raw_tensor = raw_tensor.view(kv_cache_shape)
125
+ kv_caches[layer_name] = raw_tensor.permute(*inv_order)
126
+ return kv_caches
127
+
128
+
129
+def init_kv_cache(
130
+ runner_kv_caches: list[torch.Tensor],
131
+ forward_context: dict[str, Any],
132
+ kv_cache_config: KVCacheConfig,
133
+ attn_backends: dict[str, AttentionBackend],
134
+ device: torch.device,
135
+) -> None:
136
+ kv_cache_raw_tensors = _allocate_kv_cache(kv_cache_config, device)
137
+ kv_caches = _reshape_kv_cache(kv_cache_config, kv_cache_raw_tensors, attn_backends)
138
+ bind_kv_cache(kv_caches, forward_context, runner_kv_caches)
139
+
140
+
141
+def build_attn_metadata(
142
+ attn_metadata_builders: list[AttentionMetadataBuilder],
143
+ num_reqs: int,
144
+ num_tokens: int,
145
+ query_start_loc: CpuGpuBuffer,
146
+ seq_lens: CpuGpuBuffer,
147
+ num_computed_tokens_cpu: torch.Tensor,
148
+ block_tables: Sequence[torch.Tensor],
149
+ slot_mappings: torch.Tensor,
150
+ kv_cache_config: KVCacheConfig,
151
+) -> dict[str, Any]:
152
+ query_start_loc_gpu = query_start_loc.gpu[: num_reqs + 1]
153
+ query_start_loc_cpu = query_start_loc.cpu[: num_reqs + 1]
154
+ max_query_len = int(query_start_loc.np[: num_reqs + 1].max())
155
+ seq_lens_gpu = seq_lens.gpu[:num_reqs]
156
+ seq_lens_cpu = seq_lens.cpu[:num_reqs]
157
+ max_seq_len = int(seq_lens.np[:num_reqs].max())
158
+
159
+ attn_metadata: dict[str, Any] = {}
160
+ kv_cache_groups = kv_cache_config.kv_cache_groups
161
+ for i, kv_cache_spec in enumerate(kv_cache_groups):
162
+ block_table = block_tables[i]
163
+ slot_mapping = slot_mappings[i]
164
+
165
+ common_attn_metadata = CommonAttentionMetadata(
166
+ query_start_loc=query_start_loc_gpu,
167
+ query_start_loc_cpu=query_start_loc_cpu,
168
+ seq_lens=seq_lens_gpu,
169
+ seq_lens_cpu=seq_lens_cpu,
170
+ max_seq_len=max_seq_len,
171
+ num_computed_tokens_cpu=num_computed_tokens_cpu,
172
+ num_reqs=num_reqs,
173
+ num_actual_tokens=num_tokens,
174
+ max_query_len=max_query_len,
175
+ block_table_tensor=block_table,
176
+ slot_mapping=slot_mapping,
177
+ causal=True,
178
+ )
179
+
180
+ attn_metadata_builder = attn_metadata_builders[i]
181
+ metadata = attn_metadata_builder.build(
182
+ common_prefix_len=0,
183
+ common_attn_metadata=common_attn_metadata,
184
+ )
185
+ for layer_name in kv_cache_spec.layer_names:
186
+ attn_metadata[layer_name] = metadata
187
+ return attn_metadata
vllm/v1/worker/gpu/block_table.py
new
+315
@@ -0,0 +1,315 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Iterable
4
+
5
+import torch
6
+import triton
7
+import triton.language as tl
8
+
9
+from vllm.attention.backends.utils import PAD_SLOT_ID
10
+from vllm.utils.math_utils import cdiv
11
+from vllm.v1.utils import CpuGpuBuffer
12
+
13
+
14
+class BlockTables:
15
+ def __init__(
16
+ self,
17
+ block_sizes: list[int],
18
+ max_num_reqs: int,
19
+ max_num_batched_tokens: int,
20
+ max_model_len: int,
21
+ device: torch.device,
22
+ pin_memory: bool,
23
+ ):
24
+ self.block_sizes = block_sizes
25
+ self.max_num_reqs = max_num_reqs
26
+ self.max_num_batched_tokens = max_num_batched_tokens
27
+ self.max_model_len = max_model_len
28
+ self.device = device
29
+ self.pin_memory = pin_memory
30
+
31
+ self.num_kv_cache_groups = len(self.block_sizes)
32
+ # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
33
+ self.block_tables: list[torch.Tensor] = []
34
+ for i in range(self.num_kv_cache_groups):
35
+ block_size = self.block_sizes[i]
36
+ max_num_blocks = cdiv(self.max_model_len, block_size)
37
+ block_table = torch.zeros(
38
+ self.max_num_reqs,
39
+ max_num_blocks,
40
+ dtype=torch.int32,
41
+ device=self.device,
42
+ )
43
+ self.block_tables.append(block_table)
44
+ self.block_table_ptrs = self._make_ptr_tensor(self.block_tables)
45
+
46
+ # Block tables used for model's forward pass.
47
+ # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
48
+ self.input_block_tables: list[torch.Tensor] = [
49
+ torch.zeros_like(block_table) for block_table in self.block_tables
50
+ ]
51
+ self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables)
52
+
53
+ self.block_table_strides = torch.tensor(
54
+ [b.stride(0) for b in self.block_tables],
55
+ dtype=torch.int64,
56
+ device=self.device,
57
+ )
58
+ self.block_sizes_tensor = torch.tensor(
59
+ self.block_sizes, dtype=torch.int32, device=self.device
60
+ )
61
+ self.num_blocks = torch.zeros(
62
+ self.num_kv_cache_groups,
63
+ self.max_num_reqs,
64
+ dtype=torch.int32,
65
+ device=self.device,
66
+ )
67
+ self.slot_mappings = torch.zeros(
68
+ self.num_kv_cache_groups,
69
+ self.max_num_batched_tokens,
70
+ dtype=torch.int64,
71
+ device=self.device,
72
+ )
73
+
74
+ # Misc buffers.
75
+ self.req_indices = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
76
+ self.overwrite = self._make_buffer(self.max_num_reqs, dtype=torch.bool)
77
+ self.cu_num_new_blocks = self._make_buffer(
78
+ self.num_kv_cache_groups, self.max_num_reqs + 1, dtype=torch.int32
79
+ )
80
+
81
+ def _make_buffer(self, *args, dtype: torch.dtype) -> CpuGpuBuffer:
82
+ return CpuGpuBuffer(
83
+ *args, dtype=dtype, pin_memory=self.pin_memory, device=self.device
84
+ )
85
+
86
+ def _make_ptr_tensor(self, x: Iterable[torch.Tensor]) -> torch.Tensor:
87
+ # NOTE(woosuk): Use uint64 instead of int64 to cover all possible addresses.
88
+ ptrs_tensor_cpu = torch.tensor(
89
+ [t.data_ptr() for t in x],
90
+ dtype=torch.uint64,
91
+ device="cpu",
92
+ pin_memory=self.pin_memory,
93
+ )
94
+ return ptrs_tensor_cpu.to(self.device, non_blocking=True)
95
+
96
+ def append_block_ids(
97
+ self,
98
+ # [num_reqs]
99
+ req_indices: list[int],
100
+ # [num_kv_cache_groups, num_reqs + 1]
101
+ cu_num_new_blocks: tuple[list[int], ...],
102
+ # [num_kv_cache_groups, num_new_blocks]
103
+ new_block_ids: tuple[list[int], ...],
104
+ # [num_reqs]
105
+ overwrite: list[bool],
106
+ ) -> None:
107
+ num_reqs = len(req_indices)
108
+ self.req_indices.np[:num_reqs] = req_indices
109
+ self.overwrite.np[:num_reqs] = overwrite
110
+ for i in range(self.num_kv_cache_groups):
111
+ self.cu_num_new_blocks.np[i, : num_reqs + 1] = cu_num_new_blocks[i]
112
+
113
+ # NOTE(woosuk): Here, we cannot use a fixed-size buffer because there's
114
+ # no clear upper bound to the number of new blocks in a single step.
115
+ # NOTE(woosuk): The buffer has to be cached, because otherwise we cannot
116
+ # guarantee that the buffer is not freed before the copy is completed.
117
+ self.new_block_ids_cpu = torch.empty(
118
+ self.num_kv_cache_groups,
119
+ max(len(x) for x in new_block_ids),
120
+ dtype=torch.int32,
121
+ device="cpu",
122
+ pin_memory=self.pin_memory,
123
+ )
124
+ new_block_ids_np = self.new_block_ids_cpu.numpy()
125
+ for i in range(self.num_kv_cache_groups):
126
+ new_block_ids_np[i, : len(new_block_ids[i])] = new_block_ids[i]
127
+ new_block_ids_gpu = self.new_block_ids_cpu.to(self.device, non_blocking=True)
128
+
129
+ _append_block_ids_kernel[(self.num_kv_cache_groups, num_reqs)](
130
+ self.req_indices.copy_to_gpu(num_reqs),
131
+ self.cu_num_new_blocks.copy_to_gpu(),
132
+ self.cu_num_new_blocks.gpu.stride(0),
133
+ new_block_ids_gpu,
134
+ new_block_ids_gpu.stride(0),
135
+ self.overwrite.copy_to_gpu(num_reqs),
136
+ self.block_table_strides,
137
+ self.block_table_ptrs,
138
+ self.num_blocks,
139
+ self.num_blocks.stride(0),
140
+ BLOCK_SIZE=1024, # type: ignore
141
+ )
142
+
143
+ def gather_block_tables(
144
+ self,
145
+ idx_mapping: torch.Tensor,
146
+ ) -> tuple[torch.Tensor, ...]:
147
+ num_reqs = idx_mapping.shape[0]
148
+ _gather_block_tables_kernel[(self.num_kv_cache_groups, num_reqs)](
149
+ idx_mapping,
150
+ self.block_table_ptrs,
151
+ self.input_block_table_ptrs,
152
+ self.block_table_strides,
153
+ self.num_blocks,
154
+ self.num_blocks.stride(0),
155
+ BLOCK_SIZE=1024, # type: ignore
156
+ )
157
+ return tuple(block_table[:num_reqs] for block_table in self.input_block_tables)
158
+
159
+ def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]:
160
+ return tuple(block_table[:num_reqs] for block_table in self.input_block_tables)
161
+
162
+ def compute_slot_mappings(
163
+ self,
164
+ query_start_loc: torch.Tensor,
165
+ positions: torch.Tensor,
166
+ ) -> torch.Tensor:
167
+ num_reqs = query_start_loc.shape[0] - 1
168
+ num_tokens = positions.shape[0]
169
+ num_groups = self.num_kv_cache_groups
170
+ _compute_slot_mappings_kernel[(num_groups, num_reqs + 1)](
171
+ num_tokens,
172
+ self.max_num_batched_tokens,
173
+ query_start_loc,
174
+ positions,
175
+ self.input_block_table_ptrs,
176
+ self.block_table_strides,
177
+ self.block_sizes_tensor,
178
+ self.slot_mappings,
179
+ self.slot_mappings.stride(0),
180
+ PAD_ID=PAD_SLOT_ID,
181
+ BLOCK_SIZE=1024, # type: ignore
182
+ )
183
+ return self.slot_mappings[:, :num_tokens]
184
+
185
+ def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor:
186
+ self.slot_mappings.fill_(PAD_SLOT_ID)
187
+ return self.slot_mappings[:, :num_tokens]
188
+
189
+
190
+@triton.jit
191
+def _append_block_ids_kernel(
192
+ # Inputs
193
+ req_indices, # [num_reqs]
194
+ cu_num_new_blocks_ptr, # [num_kv_cache_groups, num_reqs + 1]
195
+ cu_num_new_blocks_stride,
196
+ new_block_ids_ptr, # [num_kv_cache_groups, num_new_blocks]
197
+ new_block_ids_stride,
198
+ overwrite, # [num_reqs]
199
+ block_table_strides, # [num_kv_cache_groups]
200
+ # Outputs
201
+ block_table_ptrs, # [num_kv_cache_groups]
202
+ num_blocks_ptr, # [num_kv_cache_groups, max_num_reqs]
203
+ num_blocks_stride,
204
+ # Constants
205
+ BLOCK_SIZE: tl.constexpr,
206
+):
207
+ group_id = tl.program_id(0)
208
+ batch_idx = tl.program_id(1)
209
+ req_idx = tl.load(req_indices + batch_idx)
210
+ do_overwrite = tl.load(overwrite + batch_idx)
211
+
212
+ group_new_blocks_ptr = cu_num_new_blocks_ptr + group_id * cu_num_new_blocks_stride
213
+ start_idx = tl.load(group_new_blocks_ptr + batch_idx)
214
+ end_idx = tl.load(group_new_blocks_ptr + batch_idx + 1)
215
+ num_new_blocks = end_idx - start_idx
216
+
217
+ group_num_blocks_ptr = num_blocks_ptr + group_id * num_blocks_stride
218
+ dst_start_idx = tl.load(group_num_blocks_ptr + req_idx) if not do_overwrite else 0
219
+ dst_end_idx = dst_start_idx + num_new_blocks
220
+ tl.store(group_num_blocks_ptr + req_idx, dst_end_idx)
221
+
222
+ # Destination
223
+ block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32)
224
+ block_table_stride = tl.load(block_table_strides + group_id)
225
+ row_ptr = block_table_ptr + req_idx * block_table_stride
226
+
227
+ group_new_block_ids_ptr = new_block_ids_ptr + group_id * new_block_ids_stride
228
+ for i in range(0, num_new_blocks, BLOCK_SIZE):
229
+ offset = i + tl.arange(0, BLOCK_SIZE)
230
+ block_ids = tl.load(
231
+ group_new_block_ids_ptr + start_idx + offset, mask=offset < num_new_blocks
232
+ )
233
+ tl.store(
234
+ row_ptr + dst_start_idx + offset, block_ids, mask=offset < num_new_blocks
235
+ )
236
+
237
+
238
+@triton.jit
239
+def _gather_block_tables_kernel(
240
+ batch_idx_to_req_idx, # [batch_size]
241
+ src_block_table_ptrs, # [num_kv_cache_groups]
242
+ dst_block_table_ptrs, # [num_kv_cache_groups]
243
+ block_table_strides, # [num_kv_cache_groups]
244
+ num_blocks_ptr, # [num_kv_cache_groups, max_num_reqs]
245
+ num_blocks_stride,
246
+ BLOCK_SIZE: tl.constexpr,
247
+):
248
+ # kv cache group id
249
+ group_id = tl.program_id(0)
250
+ batch_idx = tl.program_id(1)
251
+ req_idx = tl.load(batch_idx_to_req_idx + batch_idx)
252
+
253
+ group_num_blocks_ptr = num_blocks_ptr + group_id * num_blocks_stride
254
+ num_blocks = tl.load(group_num_blocks_ptr + req_idx)
255
+
256
+ stride = tl.load(block_table_strides + group_id)
257
+ src_block_table_ptr = _load_ptr(src_block_table_ptrs + group_id, tl.int32)
258
+ src_row_ptr = src_block_table_ptr + req_idx * stride
259
+ dst_block_table_ptr = _load_ptr(dst_block_table_ptrs + group_id, tl.int32)
260
+ dst_row_ptr = dst_block_table_ptr + batch_idx * stride
261
+
262
+ for i in tl.range(0, num_blocks, BLOCK_SIZE):
263
+ offset = i + tl.arange(0, BLOCK_SIZE)
264
+ block_ids = tl.load(src_row_ptr + offset, mask=offset < num_blocks)
265
+ tl.store(dst_row_ptr + offset, block_ids, mask=offset < num_blocks)
266
+
267
+
268
+@triton.jit
269
+def _compute_slot_mappings_kernel(
270
+ num_tokens,
271
+ max_num_tokens,
272
+ cu_num_tokens, # [num_reqs + 1]
273
+ pos, # [num_tokens]
274
+ block_table_ptrs, # [num_kv_cache_groups]
275
+ block_table_strides, # [num_kv_cache_groups]
276
+ page_sizes, # [num_kv_cache_groups]
277
+ slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens]
278
+ slot_mappings_stride,
279
+ PAD_ID: tl.constexpr,
280
+ BLOCK_SIZE: tl.constexpr,
281
+):
282
+ # kv cache group id
283
+ group_id = tl.program_id(0)
284
+ req_idx = tl.program_id(1)
285
+ slot_mapping_ptr = slot_mappings_ptr + group_id * slot_mappings_stride
286
+
287
+ if req_idx == tl.num_programs(1) - 1:
288
+ # Pad remaining slots to -1. This is needed for CUDA graphs.
289
+ for i in range(num_tokens, max_num_tokens, BLOCK_SIZE):
290
+ offset = i + tl.arange(0, BLOCK_SIZE)
291
+ tl.store(slot_mapping_ptr + offset, PAD_ID, mask=offset < max_num_tokens)
292
+ return
293
+
294
+ block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32)
295
+ block_table_stride = tl.load(block_table_strides + group_id)
296
+ page_size = tl.load(page_sizes + group_id)
297
+
298
+ start_idx = tl.load(cu_num_tokens + req_idx)
299
+ end_idx = tl.load(cu_num_tokens + req_idx + 1)
300
+ for i in range(start_idx, end_idx, BLOCK_SIZE):
301
+ offset = i + tl.arange(0, BLOCK_SIZE)
302
+ positions = tl.load(pos + offset, mask=offset < end_idx, other=0)
303
+ block_indices = positions // page_size
304
+ block_numbers = tl.load(
305
+ block_table_ptr + req_idx * block_table_stride + block_indices
306
+ )
307
+ slot_ids = block_numbers * page_size + positions % page_size
308
+ tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx)
309
+
310
+
311
+@triton.jit
312
+def _load_ptr(ptr_to_ptr, elem_dtype):
313
+ ptr = tl.load(ptr_to_ptr)
314
+ ptr = tl.cast(ptr, tl.pointer_type(elem_dtype))
315
+ return tl.multiple_of(ptr, 16)
vllm/v1/worker/gpu/cudagraph_utils.py
new
+198
@@ -0,0 +1,198 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import gc
4
+from contextlib import contextmanager
5
+
6
+import numpy as np
7
+import torch
8
+import torch.nn as nn
9
+from tqdm import tqdm
10
+
11
+from vllm.config import VllmConfig
12
+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.backends.utils import AttentionMetadataBuilder
16
+from vllm.v1.core.sched.output import SchedulerOutput
17
+from vllm.v1.kv_cache_interface import KVCacheConfig
18
+from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
19
+from vllm.v1.worker.gpu.block_table import BlockTables
20
+from vllm.v1.worker.gpu.input_batch import InputBuffers
21
+
22
+
23
+class CudaGraphManager:
24
+ def __init__(
25
+ self,
26
+ vllm_config: VllmConfig,
27
+ device: torch.device,
28
+ ):
29
+ self.vllm_config = vllm_config
30
+ self.device = device
31
+
32
+ self.max_model_len = vllm_config.model_config.max_model_len
33
+ self.dp_size = vllm_config.parallel_config.data_parallel_size
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)
39
+ self.padded_sizes = self._init_padded_sizes()
40
+
41
+ self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
42
+ self.pool = torch.cuda.graph_pool_handle()
43
+ self.hidden_states: torch.Tensor | None = None
44
+
45
+ def _init_padded_sizes(self) -> dict[int, int]:
46
+ if not self.cudagraph_mode.has_full_cudagraphs():
47
+ # Full cuda graphs are not used.
48
+ return {}
49
+
50
+ padded_sizes: dict[int, int] = {}
51
+ assert len(self.cudagraph_sizes) > 0
52
+ for i in range(1, self.cudagraph_sizes[-1] + 1):
53
+ for x in self.cudagraph_sizes:
54
+ if i <= x:
55
+ padded_sizes[i] = x
56
+ break
57
+ return padded_sizes
58
+
59
+ def needs_capture(self) -> bool:
60
+ return len(self.padded_sizes) > 0
61
+
62
+ def get_cudagraph_size(
63
+ self,
64
+ scheduler_output: SchedulerOutput,
65
+ num_tokens_after_padding: int,
66
+ ) -> int | None:
67
+ if not self.cudagraph_mode.has_full_cudagraphs():
68
+ return None
69
+ if self.cudagraph_mode != CUDAGraphMode.FULL:
70
+ # TODO(woosuk): Support uniform decode with multiple tokens (spec decoding).
71
+ all_decode = all(
72
+ x == 1 for x in scheduler_output.num_scheduled_tokens.values()
73
+ )
74
+ if not all_decode:
75
+ # Prefill is included.
76
+ return None
77
+ return self.padded_sizes.get(num_tokens_after_padding)
78
+
79
+ def capture_graph(
80
+ self,
81
+ batch_size: int,
82
+ model: nn.Module,
83
+ input_buffers: InputBuffers,
84
+ block_tables: BlockTables,
85
+ attn_metadata_builders: list[AttentionMetadataBuilder],
86
+ kv_cache_config: KVCacheConfig,
87
+ ) -> None:
88
+ assert batch_size not in self.graphs
89
+
90
+ # Prepare dummy inputs.
91
+ input_ids = input_buffers.input_ids.gpu[:batch_size]
92
+ positions = input_buffers.positions.gpu[:batch_size]
93
+
94
+ input_buffers.query_start_loc.np[: batch_size + 1] = np.arange(batch_size + 1)
95
+ input_buffers.query_start_loc.np[batch_size:] = batch_size
96
+ input_buffers.query_start_loc.copy_to_gpu()
97
+ input_buffers.seq_lens.np[:batch_size] = self.max_model_len
98
+ input_buffers.seq_lens.np[batch_size:] = 0
99
+ input_buffers.seq_lens.copy_to_gpu()
100
+
101
+ input_block_tables = [x[:batch_size] for x in block_tables.input_block_tables]
102
+ slot_mappings = block_tables.slot_mappings[:, :batch_size]
103
+
104
+ attn_metadata = build_attn_metadata(
105
+ attn_metadata_builders=attn_metadata_builders,
106
+ num_reqs=batch_size,
107
+ num_tokens=batch_size,
108
+ query_start_loc=input_buffers.query_start_loc,
109
+ seq_lens=input_buffers.seq_lens,
110
+ num_computed_tokens_cpu=None, # FIXME
111
+ block_tables=input_block_tables,
112
+ slot_mappings=slot_mappings,
113
+ kv_cache_config=kv_cache_config,
114
+ )
115
+ if self.dp_size > 1:
116
+ num_tokens_across_dp = torch.full(
117
+ (self.dp_size,),
118
+ batch_size,
119
+ dtype=torch.int32,
120
+ device="cpu",
121
+ )
122
+ else:
123
+ num_tokens_across_dp = None
124
+
125
+ # Warm up.
126
+ with set_forward_context(
127
+ attn_metadata,
128
+ self.vllm_config,
129
+ num_tokens=batch_size,
130
+ num_tokens_across_dp=num_tokens_across_dp,
131
+ ):
132
+ hidden_states = model(
133
+ input_ids=input_ids,
134
+ positions=positions,
135
+ )
136
+ if self.hidden_states is None:
137
+ self.hidden_states = torch.empty_like(hidden_states)
138
+ torch.cuda.synchronize()
139
+
140
+ # Capture the graph.
141
+ graph = torch.cuda.CUDAGraph()
142
+ with (
143
+ set_forward_context(
144
+ attn_metadata,
145
+ self.vllm_config,
146
+ num_tokens=batch_size,
147
+ num_tokens_across_dp=num_tokens_across_dp,
148
+ ),
149
+ torch.cuda.graph(graph, self.pool),
150
+ ):
151
+ hidden_states = model(
152
+ input_ids=input_ids,
153
+ positions=positions,
154
+ )
155
+ self.hidden_states[:batch_size] = hidden_states
156
+ self.graphs[batch_size] = graph
157
+
158
+ @torch.inference_mode()
159
+ def capture(
160
+ self,
161
+ model: nn.Module,
162
+ input_buffers: InputBuffers,
163
+ block_tables: BlockTables,
164
+ attn_metadata_builders: list[AttentionMetadataBuilder],
165
+ kv_cache_config: KVCacheConfig,
166
+ ) -> None:
167
+ assert self.needs_capture()
168
+ # Capture larger graphs first.
169
+ sizes_to_capture = sorted(self.cudagraph_sizes, reverse=True)
170
+ if is_global_first_rank():
171
+ sizes_to_capture = tqdm(sizes_to_capture, desc="Capturing CUDA graphs")
172
+
173
+ with freeze_gc(), graph_capture(device=self.device):
174
+ for batch_size in sizes_to_capture:
175
+ self.capture_graph(
176
+ batch_size,
177
+ model,
178
+ input_buffers,
179
+ block_tables,
180
+ attn_metadata_builders,
181
+ kv_cache_config,
182
+ )
183
+
184
+ def run(self, batch_size: int) -> torch.Tensor:
185
+ assert batch_size in self.graphs
186
+ self.graphs[batch_size].replay()
187
+ assert self.hidden_states is not None
188
+ return self.hidden_states[:batch_size]
189
+
190
+
191
+@contextmanager
192
+def freeze_gc():
193
+ gc.collect()
194
+ gc.freeze()
195
+ try:
196
+ yield
197
+ finally:
198
+ gc.unfreeze()
vllm/v1/worker/gpu/dp_utils.py
new
+22
@@ -0,0 +1,22 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+import torch.distributed as dist
5
+
6
+from vllm.distributed.parallel_state import get_dp_group
7
+
8
+
9
+def get_batch_metadata_across_dp(
10
+ num_tokens: int,
11
+ cudagraph_size: int,
12
+ dp_size: int,
13
+ dp_rank: int,
14
+) -> tuple[torch.Tensor, torch.Tensor]:
15
+ assert dp_size > 1
16
+ # Use CPU group to avoid CPU-GPU synchronization.
17
+ group = get_dp_group().cpu_group
18
+ tensor = torch.zeros(2, dp_size, dtype=torch.int32, device="cpu")
19
+ tensor[0][dp_rank] = num_tokens
20
+ tensor[1][dp_rank] = cudagraph_size
21
+ dist.all_reduce(tensor, group=group)
22
+ return tensor[0], tensor[1]
vllm/v1/worker/gpu/input_batch.py
new
+265
@@ -0,0 +1,265 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+from typing import Any
5
+
6
+import numba
7
+import numba.types as types
8
+import numpy as np
9
+import torch
10
+import triton
11
+import triton.language as tl
12
+
13
+from vllm.utils import random_uuid
14
+from vllm.utils.math_utils import cdiv
15
+from vllm.v1.utils import CpuGpuBuffer
16
+
17
+
18
+class InputBuffers:
19
+ def __init__(
20
+ self,
21
+ max_num_reqs: int,
22
+ max_num_tokens: int,
23
+ hidden_size: int,
24
+ vocab_size: int,
25
+ dtype: torch.dtype,
26
+ device: torch.device,
27
+ pin_memory: bool,
28
+ ):
29
+ self.max_num_reqs = max_num_reqs
30
+ self.max_num_tokens = max_num_tokens
31
+ self.device = device
32
+ self.pin_memory = pin_memory
33
+
34
+ self.idx_mapping = self._make_buffer(max_num_reqs, dtype=torch.int32)
35
+ self.input_ids = self._make_buffer(max_num_tokens, dtype=torch.int32)
36
+ self.positions = self._make_buffer(max_num_tokens, dtype=torch.int64)
37
+ self.query_start_loc = self._make_buffer(max_num_reqs + 1, dtype=torch.int32)
38
+ self.seq_lens = self._make_buffer(max_num_reqs, dtype=torch.int32)
39
+
40
+ # Structured outputs.
41
+ self.bitmask_indices = self._make_buffer(max_num_reqs, dtype=torch.int32)
42
+ self.grammar_bitmask = self._make_buffer(
43
+ max_num_reqs, cdiv(vocab_size, 32), dtype=torch.int32
44
+ )
45
+
46
+ def _make_buffer(self, *args, dtype: torch.dtype) -> CpuGpuBuffer:
47
+ return CpuGpuBuffer(
48
+ *args, dtype=dtype, pin_memory=self.pin_memory, device=self.device
49
+ )
50
+
51
+
52
+@dataclass
53
+class InputBatch:
54
+ # batch_idx -> req_id
55
+ req_ids: list[str]
56
+ num_reqs: int
57
+
58
+ # batch_idx -> req_state_idx
59
+ idx_mapping: torch.Tensor
60
+ idx_mapping_np: np.ndarray
61
+
62
+ # [num_reqs]
63
+ # batch_idx -> num_scheduled_tokens
64
+ num_scheduled_tokens: np.ndarray
65
+ # sum(num_scheduled_tokens)
66
+ num_tokens: int
67
+ num_tokens_after_padding: int
68
+
69
+ # [num_reqs + 1]
70
+ query_start_loc: torch.Tensor
71
+ query_start_loc_np: np.ndarray
72
+ # [num_reqs]
73
+ seq_lens: torch.Tensor
74
+ seq_lens_np: np.ndarray
75
+
76
+ # [num_tokens_after_padding]
77
+ input_ids: torch.Tensor
78
+ # [num_tokens_after_padding]
79
+ positions: torch.Tensor
80
+
81
+ # layer_name -> Metadata
82
+ attn_metadata: dict[str, Any]
83
+
84
+ # [num_reqs]
85
+ logits_indices: torch.Tensor
86
+
87
+ @classmethod
88
+ def make_dummy(
89
+ cls,
90
+ num_reqs: int,
91
+ num_tokens: int,
92
+ input_buffers: InputBuffers,
93
+ device: torch.device,
94
+ ) -> "InputBatch":
95
+ assert 0 < num_reqs <= num_tokens
96
+ req_ids = [f"req_{i}_{random_uuid()}" for i in range(num_reqs)]
97
+ idx_mapping_np = np.arange(num_reqs, dtype=np.int32)
98
+ idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
99
+ num_scheduled_tokens = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)
100
+ num_scheduled_tokens[-1] += num_tokens % num_reqs
101
+ assert int(num_scheduled_tokens.sum()) == num_tokens
102
+
103
+ input_buffers.query_start_loc.np[0] = 0
104
+ input_buffers.query_start_loc.np[1 : num_reqs + 1] = np.cumsum(
105
+ num_scheduled_tokens
106
+ )
107
+ input_buffers.query_start_loc.np[num_reqs + 1 :] = num_tokens
108
+ query_start_loc_np = input_buffers.query_start_loc.np[: num_reqs + 1]
109
+ query_start_loc = input_buffers.query_start_loc.copy_to_gpu()[: num_reqs + 1]
110
+ # seq_len equals to query_len
111
+ input_buffers.seq_lens.np[:num_reqs] = num_scheduled_tokens
112
+ input_buffers.seq_lens.np[num_reqs:] = 0
113
+ seq_lens_np = input_buffers.seq_lens.np[:num_reqs]
114
+ seq_lens = input_buffers.seq_lens.copy_to_gpu()[:num_reqs]
115
+
116
+ input_ids = input_buffers.input_ids.copy_to_gpu(num_tokens)
117
+ positions = input_buffers.positions.copy_to_gpu(num_tokens)
118
+ # attn_metadata = defaultdict(lambda: None)
119
+ logits_indices = query_start_loc[1:] - 1
120
+ return cls(
121
+ req_ids=req_ids,
122
+ num_reqs=num_reqs,
123
+ idx_mapping=idx_mapping,
124
+ idx_mapping_np=idx_mapping_np,
125
+ num_scheduled_tokens=num_scheduled_tokens,
126
+ num_tokens=num_tokens,
127
+ num_tokens_after_padding=num_tokens,
128
+ query_start_loc=query_start_loc,
129
+ query_start_loc_np=query_start_loc_np,
130
+ seq_lens=seq_lens,
131
+ seq_lens_np=seq_lens_np,
132
+ input_ids=input_ids,
133
+ positions=positions,
134
+ attn_metadata=None, # type: ignore
135
+ logits_indices=logits_indices,
136
+ )
137
+
138
+
139
+# NOTE: With the type annotations, this function is pre-compiled
140
+# before the first call.
141
+@numba.jit(
142
+ [
143
+ types.none(
144
+ types.int32[:], # idx_mapping
145
+ types.int32[:, :], # token_ids
146
+ types.int32[:], # num_computed_tokens
147
+ types.int32[:], # num_scheduled_tokens
148
+ types.int32[:], # input_ids
149
+ types.int64[:], # positions
150
+ types.int32[:], # query_start_loc
151
+ types.int32[:], # seq_lens
152
+ )
153
+ ],
154
+ nopython=True,
155
+ cache=True,
156
+)
157
+def _prepare_inputs(
158
+ idx_mapping: np.ndarray, # batch_idx -> req_idx
159
+ token_ids: np.ndarray, # [N, max_model_len]
160
+ num_computed_tokens: np.ndarray, # [N]
161
+ num_scheduled_tokens: np.ndarray, # [B]
162
+ input_ids: np.ndarray, # [num_input_tokens]
163
+ positions: np.ndarray, # [num_input_tokens]
164
+ query_start_loc: np.ndarray, # [B + 1]
165
+ seq_lens: np.ndarray, # [B]
166
+) -> None:
167
+ num_reqs = num_scheduled_tokens.shape[0]
168
+ query_start_loc[0] = 0
169
+
170
+ cu_num_tokens = 0
171
+ for i in range(num_reqs):
172
+ req_idx = idx_mapping[i]
173
+ query_len = num_scheduled_tokens[i]
174
+ start = num_computed_tokens[req_idx]
175
+ end = start + query_len
176
+ seq_lens[i] = end
177
+
178
+ start_idx = cu_num_tokens
179
+ end_idx = start_idx + query_len
180
+ input_ids[start_idx:end_idx] = token_ids[req_idx, start:end]
181
+ positions[start_idx:end_idx] = np.arange(start, end, dtype=np.int64)
182
+
183
+ cu_num_tokens = end_idx
184
+ query_start_loc[i + 1] = cu_num_tokens
185
+
186
+ # Pad the inputs for CUDA graphs.
187
+ # Note: pad query_start_loc to be non-decreasing, as kernels
188
+ # like FlashAttention requires that
189
+ query_start_loc[num_reqs + 1 :].fill(cu_num_tokens)
190
+ # Fill unused with 0 for full cuda graph mode.
191
+ seq_lens[num_reqs:].fill(0)
192
+
193
+
194
+def prepare_inputs(
195
+ idx_mapping: np.ndarray,
196
+ prefill_token_ids: np.ndarray,
197
+ num_computed_tokens: np.ndarray,
198
+ num_scheduled_tokens: np.ndarray,
199
+ input_ids: CpuGpuBuffer,
200
+ positions: CpuGpuBuffer,
201
+ query_start_loc: CpuGpuBuffer,
202
+ seq_lens: CpuGpuBuffer,
203
+ num_tokens: int,
204
+) -> None:
205
+ _prepare_inputs(
206
+ idx_mapping,
207
+ prefill_token_ids,
208
+ num_computed_tokens,
209
+ num_scheduled_tokens,
210
+ input_ids.np,
211
+ positions.np,
212
+ query_start_loc.np,
213
+ seq_lens.np,
214
+ )
215
+ input_ids.copy_to_gpu(num_tokens)
216
+ positions.copy_to_gpu(num_tokens)
217
+ # NOTE(woosuk): We should copy the whole query_start_loc and seq_lens
218
+ # tensors from CPU to GPU, because they may include paddings needed
219
+ # for full CUDA graph mode.
220
+ query_start_loc.copy_to_gpu()
221
+ seq_lens.copy_to_gpu()
222
+ return
223
+
224
+
225
+@triton.jit
226
+def _combine_last_token_ids_kernel(
227
+ input_ids_ptr,
228
+ idx_mapping_ptr,
229
+ last_token_ids_ptr,
230
+ query_start_loc_ptr,
231
+ seq_lens_ptr,
232
+ prefill_len_ptr,
233
+):
234
+ batch_idx = tl.program_id(0)
235
+ req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
236
+
237
+ seq_len = tl.load(seq_lens_ptr + batch_idx)
238
+ prefill_len = tl.load(prefill_len_ptr + req_state_idx)
239
+ if seq_len <= prefill_len:
240
+ # Handling prefill tokens.
241
+ return
242
+
243
+ last_token_id = tl.load(last_token_ids_ptr + req_state_idx)
244
+ end = tl.load(query_start_loc_ptr + batch_idx + 1)
245
+ tl.store(input_ids_ptr + end - 1, last_token_id)
246
+
247
+
248
+def combine_last_token_ids(
249
+ input_ids: torch.Tensor,
250
+ idx_mapping: torch.Tensor,
251
+ last_token_ids: torch.Tensor,
252
+ query_start_loc: torch.Tensor,
253
+ seq_lens: torch.Tensor,
254
+ prefill_len: torch.Tensor,
255
+) -> torch.Tensor:
256
+ num_reqs = seq_lens.shape[0]
257
+ _combine_last_token_ids_kernel[(num_reqs,)](
258
+ input_ids,
259
+ idx_mapping,
260
+ last_token_ids,
261
+ query_start_loc,
262
+ seq_lens,
263
+ prefill_len,
264
+ )
265
+ return input_ids
vllm/v1/worker/gpu/model_runner.py
new
+814
@@ -0,0 +1,814 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import gc
4
+import time
5
+from copy import deepcopy
6
+from typing import Any
7
+
8
+import numpy as np
9
+import torch
10
+import torch.nn as nn
11
+
12
+from vllm.config import VllmConfig
13
+from vllm.config.compilation import CUDAGraphMode
14
+from vllm.forward_context import set_forward_context
15
+from vllm.logger import init_logger
16
+from vllm.model_executor.model_loader import get_model_loader
17
+from vllm.utils.mem_constants import GiB_bytes
18
+from vllm.utils.mem_utils import DeviceMemoryProfiler
19
+from vllm.utils.platform_utils import is_pin_memory_available
20
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
21
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
22
+from vllm.v1.kv_cache_interface import KVCacheConfig
23
+from vllm.v1.outputs import (
24
+ EMPTY_MODEL_RUNNER_OUTPUT,
25
+ LogprobsTensors,
26
+ ModelRunnerOutput,
27
+)
28
+from vllm.v1.sample.sampler import SamplerOutput
29
+from vllm.v1.worker.gpu.async_utils import AsyncOutput, async_barrier
30
+from vllm.v1.worker.gpu.attn_utils import (
31
+ build_attn_metadata,
32
+ get_kv_cache_spec,
33
+ init_attn_backend,
34
+ init_kv_cache,
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
39
+from vllm.v1.worker.gpu.input_batch import (
40
+ InputBatch,
41
+ InputBuffers,
42
+ combine_last_token_ids,
43
+ prepare_inputs,
44
+)
45
+from vllm.v1.worker.gpu.sampler import Sampler, compute_prompt_logprobs
46
+from vllm.v1.worker.gpu.states import RequestState, SamplingMetadata
47
+from vllm.v1.worker.gpu.structured_outputs import apply_grammar_bitmask
48
+from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
49
+from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
50
+
51
+logger = init_logger(__name__)
52
+
53
+
54
+class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
55
+ def __init__(
56
+ self,
57
+ vllm_config: VllmConfig,
58
+ device: torch.device,
59
+ ):
60
+ self.vllm_config = vllm_config
61
+ self.model_config = vllm_config.model_config
62
+ self.cache_config = vllm_config.cache_config
63
+ self.compilation_config = vllm_config.compilation_config
64
+ self.lora_config = vllm_config.lora_config
65
+ self.load_config = vllm_config.load_config
66
+ self.parallel_config = vllm_config.parallel_config
67
+ self.scheduler_config = vllm_config.scheduler_config
68
+ self.speculative_config = vllm_config.speculative_config
69
+ self.observability_config = vllm_config.observability_config
70
+
71
+ self.device = device
72
+ self.pin_memory = is_pin_memory_available()
73
+ self.dtype = self.model_config.dtype
74
+ self.kv_cache_dtype = self.dtype
75
+ if self.cache_config.cache_dtype != "auto":
76
+ # Quantized KV cache.
77
+ self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
78
+ self.cache_config.cache_dtype
79
+ ]
80
+ self.is_pooling_model = False
81
+
82
+ self.vocab_size = self.model_config.get_vocab_size()
83
+ self.max_model_len = self.model_config.max_model_len
84
+ self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
85
+ self.max_num_reqs = self.scheduler_config.max_num_seqs
86
+ self.hidden_size = self.model_config.get_hidden_size()
87
+
88
+ self.dp_size = self.parallel_config.data_parallel_size
89
+ self.dp_rank = self.parallel_config.data_parallel_rank
90
+
91
+ self.use_async_scheduling = self.scheduler_config.async_scheduling
92
+ self.output_copy_stream = torch.cuda.Stream(self.device)
93
+ self.output_copy_event = torch.cuda.Event()
94
+ if self.use_async_scheduling:
95
+ self.input_prep_event = torch.cuda.Event()
96
+ self.structured_outputs_event = torch.cuda.Event()
97
+ else:
98
+ self.input_prep_event = None
99
+ self.structured_outputs_event = None
100
+
101
+ self.req_states = RequestState(
102
+ max_num_reqs=self.max_num_reqs,
103
+ max_model_len=self.max_model_len,
104
+ max_num_batched_tokens=self.max_num_tokens,
105
+ vocab_size=self.vocab_size,
106
+ device=self.device,
107
+ pin_memory=self.pin_memory,
108
+ )
109
+ self.input_buffers = InputBuffers(
110
+ max_num_reqs=self.max_num_reqs,
111
+ max_num_tokens=self.max_num_tokens,
112
+ hidden_size=self.hidden_size,
113
+ vocab_size=self.vocab_size,
114
+ dtype=self.dtype,
115
+ device=self.device,
116
+ pin_memory=self.pin_memory,
117
+ )
118
+ self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)
119
+
120
+ # CUDA graphs.
121
+ self.cudagraph_manager = CudaGraphManager(
122
+ vllm_config=self.vllm_config,
123
+ device=self.device,
124
+ )
125
+
126
+ def get_supported_tasks(self) -> tuple[str]:
127
+ return ("generate",)
128
+
129
+ def load_model(self, *args, **kwargs) -> None:
130
+ time_before_load = time.perf_counter()
131
+ with DeviceMemoryProfiler() as m:
132
+ model_loader = get_model_loader(self.vllm_config.load_config)
133
+ logger.info("Loading model from scratch...")
134
+
135
+ self.model = model_loader.load_model(
136
+ vllm_config=self.vllm_config,
137
+ model_config=self.vllm_config.model_config,
138
+ )
139
+ if self.lora_config:
140
+ self.model = self.load_lora_model(
141
+ self.model,
142
+ self.vllm_config,
143
+ self.device,
144
+ )
145
+ time_after_load = time.perf_counter()
146
+
147
+ self.model_memory_usage = m.consumed_memory
148
+ logger.info(
149
+ "Model loading took %.4f GiB and %.6f seconds",
150
+ m.consumed_memory / GiB_bytes,
151
+ time_after_load - time_before_load,
152
+ )
153
+
154
+ def get_model(self) -> nn.Module:
155
+ return self.model
156
+
157
+ def get_kv_cache_spec(self):
158
+ return get_kv_cache_spec(self.vllm_config)
159
+
160
+ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
161
+ kv_cache_config = deepcopy(kv_cache_config)
162
+ self.kv_cache_config = kv_cache_config
163
+ block_sizes = [
164
+ kv_cache_group.kv_cache_spec.block_size
165
+ for kv_cache_group in kv_cache_config.kv_cache_groups
166
+ ]
167
+
168
+ self.block_tables = BlockTables(
169
+ block_sizes=block_sizes,
170
+ max_num_reqs=self.max_num_reqs,
171
+ max_num_batched_tokens=self.max_num_tokens,
172
+ max_model_len=self.max_model_len,
173
+ device=self.device,
174
+ pin_memory=self.pin_memory,
175
+ )
176
+
177
+ self.attn_backends, self.attn_metadata_builders = init_attn_backend(
178
+ self.kv_cache_config,
179
+ self.vllm_config,
180
+ self.device,
181
+ )
182
+
183
+ self.kv_caches: list[torch.Tensor] = []
184
+ init_kv_cache(
185
+ self.kv_caches,
186
+ self.compilation_config.static_forward_context,
187
+ self.kv_cache_config,
188
+ self.attn_backends,
189
+ self.device,
190
+ )
191
+ # Attention groups are not supported.
192
+ self.attn_groups = [] # type: ignore
193
+
194
+ def prepare_dummy_attn_metadata(self, input_batch: InputBatch) -> None:
195
+ block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs)
196
+ slot_mappings = self.block_tables.get_dummy_slot_mappings(
197
+ input_batch.num_tokens
198
+ )
199
+ num_computed_tokens_cpu = torch.zeros(
200
+ input_batch.num_reqs, dtype=torch.int32, device="cpu"
201
+ )
202
+ attn_metadata = build_attn_metadata(
203
+ attn_metadata_builders=self.attn_metadata_builders,
204
+ num_reqs=input_batch.num_reqs,
205
+ num_tokens=input_batch.num_tokens,
206
+ query_start_loc=self.input_buffers.query_start_loc,
207
+ seq_lens=self.input_buffers.seq_lens,
208
+ num_computed_tokens_cpu=num_computed_tokens_cpu,
209
+ block_tables=block_tables,
210
+ slot_mappings=slot_mappings,
211
+ kv_cache_config=self.kv_cache_config,
212
+ )
213
+ input_batch.attn_metadata = attn_metadata
214
+
215
+ @torch.inference_mode()
216
+ def _dummy_run(
217
+ self,
218
+ num_tokens: int,
219
+ *args,
220
+ skip_attn: bool = True,
221
+ **kwargs,
222
+ ) -> tuple[torch.Tensor, torch.Tensor]:
223
+ num_reqs = min(num_tokens, self.max_num_reqs)
224
+ input_batch = InputBatch.make_dummy(
225
+ num_reqs=num_reqs,
226
+ num_tokens=num_tokens,
227
+ input_buffers=self.input_buffers,
228
+ device=self.device,
229
+ )
230
+ if not skip_attn:
231
+ self.prepare_dummy_attn_metadata(input_batch)
232
+
233
+ if self.dp_size == 1:
234
+ num_tokens_across_dp: torch.Tensor | None = None
235
+ else:
236
+ num_tokens_across_dp = torch.full(
237
+ (self.dp_size,), num_tokens, dtype=torch.int32, device="cpu"
238
+ )
239
+ num_sampled_tokens = np.ones(input_batch.num_reqs, dtype=np.int32)
240
+ with (
241
+ self.maybe_dummy_run_with_lora(
242
+ self.lora_config,
243
+ input_batch.num_scheduled_tokens,
244
+ num_sampled_tokens,
245
+ ),
246
+ set_forward_context(
247
+ input_batch.attn_metadata,
248
+ self.vllm_config,
249
+ num_tokens=num_tokens,
250
+ num_tokens_across_dp=num_tokens_across_dp,
251
+ ),
252
+ ):
253
+ hidden_states = self.model(
254
+ input_ids=input_batch.input_ids,
255
+ positions=input_batch.positions,
256
+ )
257
+ sample_hidden_states = hidden_states[input_batch.logits_indices]
258
+ return hidden_states, sample_hidden_states
259
+
260
+ @torch.inference_mode()
261
+ def _dummy_sampler_run(
262
+ self,
263
+ hidden_states: torch.Tensor,
264
+ ) -> None:
265
+ num_reqs = hidden_states.shape[0]
266
+ sampling_metadata = SamplingMetadata.make_dummy(
267
+ num_reqs=num_reqs,
268
+ device=self.device,
269
+ )
270
+ logits = self.model.compute_logits(hidden_states)
271
+ self.sampler(logits, sampling_metadata)
272
+
273
+ @torch.inference_mode()
274
+ def profile_run(self) -> None:
275
+ hidden_states, sample_hidden_states = self._dummy_run(
276
+ self.max_num_tokens,
277
+ skip_attn=True,
278
+ )
279
+ self._dummy_sampler_run(sample_hidden_states)
280
+ torch.cuda.synchronize()
281
+ del hidden_states, sample_hidden_states
282
+ gc.collect()
283
+
284
+ def reset_mm_cache(self) -> None:
285
+ pass
286
+
287
+ def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int:
288
+ # SP is not supported yet.
289
+ return num_scheduled_tokens
290
+
291
+ @torch.inference_mode()
292
+ def capture_model(self) -> int:
293
+ if not self.cudagraph_manager.needs_capture():
294
+ logger.warning(
295
+ "Skipping CUDA graph capture. To turn on CUDA graph capture, "
296
+ "ensure `cudagraph_mode` was not manually set to `NONE`"
297
+ )
298
+ return 0
299
+
300
+ start_time = time.perf_counter()
301
+ start_free_gpu_memory = torch.cuda.mem_get_info()[0]
302
+
303
+ with self.maybe_setup_dummy_loras(self.lora_config):
304
+ self.cudagraph_manager.capture(
305
+ model=self.model,
306
+ input_buffers=self.input_buffers,
307
+ block_tables=self.block_tables,
308
+ attn_metadata_builders=self.attn_metadata_builders,
309
+ kv_cache_config=self.kv_cache_config,
310
+ )
311
+
312
+ end_time = time.perf_counter()
313
+ end_free_gpu_memory = torch.cuda.mem_get_info()[0]
314
+ elapsed_time = end_time - start_time
315
+ cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory
316
+ # This usually takes 5~20 seconds.
317
+ logger.info(
318
+ "Graph capturing finished in %.0f secs, took %.2f GiB",
319
+ elapsed_time,
320
+ cuda_graph_size / (1 << 30),
321
+ )
322
+ return cuda_graph_size
323
+
324
+ def warmup_for_prefill(self) -> None:
325
+ # For FlashInfer, we would like to execute a dummy prefill run
326
+ # to trigger JIT compilation.
327
+ if all("FLASHINFER" in b.get_name() for b in self.attn_backends.values()):
328
+ self._dummy_run(self.max_num_tokens, skip_attn=False)
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)
334
+ for req_id in scheduler_output.finished_req_ids:
335
+ self.req_states.remove_request(req_id)
336
+
337
+ # TODO(woosuk): Change SchedulerOutput.
338
+ req_indices: list[int] = []
339
+ cu_num_new_blocks = tuple(
340
+ [0] for _ in range(self.block_tables.num_kv_cache_groups)
341
+ )
342
+ new_block_ids: tuple[list[int], ...] = tuple(
343
+ [] for _ in range(self.block_tables.num_kv_cache_groups)
344
+ )
345
+ overwrite: list[bool] = []
346
+
347
+ # Add new requests.
348
+ for new_req_data in scheduler_output.scheduled_new_reqs:
349
+ req_id = new_req_data.req_id
350
+ self.req_states.add_request(
351
+ req_id=req_id,
352
+ prompt_len=len(new_req_data.prompt_token_ids),
353
+ prefill_token_ids=new_req_data.prefill_token_ids,
354
+ num_computed_tokens=new_req_data.num_computed_tokens,
355
+ sampling_params=new_req_data.sampling_params,
356
+ lora_request=new_req_data.lora_request,
357
+ )
358
+
359
+ req_index = self.req_states.req_id_to_index[req_id]
360
+ req_indices.append(req_index)
361
+ for i, block_ids in enumerate(new_req_data.block_ids):
362
+ x = cu_num_new_blocks[i][-1]
363
+ cu_num_new_blocks[i].append(x + len(block_ids))
364
+ new_block_ids[i].extend(block_ids)
365
+ overwrite.append(True)
366
+
367
+ # Add new blocks for the existing requests.
368
+ cached_reqs = scheduler_output.scheduled_cached_reqs
369
+ for i, req_id in enumerate(cached_reqs.req_ids):
370
+ req_index = self.req_states.req_id_to_index[req_id]
371
+
372
+ req_new_block_ids = cached_reqs.new_block_ids[i]
373
+ if req_new_block_ids is not None:
374
+ req_indices.append(req_index)
375
+ for group_id, block_ids in enumerate(req_new_block_ids):
376
+ x = cu_num_new_blocks[group_id][-1]
377
+ cu_num_new_blocks[group_id].append(x + len(block_ids))
378
+ new_block_ids[group_id].extend(block_ids)
379
+ overwrite.append(False)
380
+
381
+ if req_indices:
382
+ self.block_tables.append_block_ids(
383
+ req_indices=req_indices,
384
+ cu_num_new_blocks=cu_num_new_blocks,
385
+ new_block_ids=new_block_ids,
386
+ overwrite=overwrite,
387
+ )
388
+
389
+ def prepare_inputs(
390
+ self,
391
+ scheduler_output: SchedulerOutput,
392
+ num_tokens_after_padding: int,
393
+ ) -> InputBatch:
394
+ num_tokens = scheduler_output.total_num_scheduled_tokens
395
+ assert num_tokens > 0
396
+ num_reqs = len(scheduler_output.num_scheduled_tokens)
397
+
398
+ # Decode first, then prefill.
399
+ # batch_idx -> req_id
400
+ req_ids = sorted(
401
+ scheduler_output.num_scheduled_tokens,
402
+ key=scheduler_output.num_scheduled_tokens.get,
403
+ )
404
+ num_scheduled_tokens = np.array(
405
+ [scheduler_output.num_scheduled_tokens[i] for i in req_ids], dtype=np.int32
406
+ )
407
+
408
+ idx_mapping_list = [
409
+ self.req_states.req_id_to_index[req_id] for req_id in req_ids
410
+ ]
411
+ idx_mapping = self.input_buffers.idx_mapping
412
+ idx_mapping.np[:num_reqs] = idx_mapping_list
413
+ idx_mapping_np = idx_mapping.np[:num_reqs]
414
+ idx_mapping = idx_mapping.copy_to_gpu(num_reqs)
415
+
416
+ # Block tables: num_kv_cache_groups x [num_reqs, max_num_blocks]
417
+ block_tables = self.block_tables.gather_block_tables(idx_mapping)
418
+
419
+ prepare_inputs(
420
+ idx_mapping_np,
421
+ self.req_states.prefill_token_ids,
422
+ self.req_states.num_computed_tokens,
423
+ num_scheduled_tokens,
424
+ self.input_buffers.input_ids,
425
+ self.input_buffers.positions,
426
+ self.input_buffers.query_start_loc,
427
+ self.input_buffers.seq_lens,
428
+ num_tokens,
429
+ )
430
+
431
+ query_start_loc = self.input_buffers.query_start_loc
432
+ query_start_loc_gpu = query_start_loc.gpu[: num_reqs + 1]
433
+ query_start_loc_np = query_start_loc.np[: num_reqs + 1]
434
+ seq_lens_gpu = self.input_buffers.seq_lens.gpu[:num_reqs]
435
+ seq_lens_np = self.input_buffers.seq_lens.np[:num_reqs]
436
+
437
+ # Some input token ids are directly read from the last sampled tokens.
438
+ combine_last_token_ids(
439
+ self.input_buffers.input_ids.gpu,
440
+ idx_mapping,
441
+ self.req_states.last_sampled_tokens,
442
+ query_start_loc_gpu,
443
+ seq_lens_gpu,
444
+ self.req_states.prefill_len.copy_to_gpu(),
445
+ )
446
+
447
+ # Compute slot mappings: [num_kv_cache_groups, num_tokens]
448
+ slot_mappings = self.block_tables.compute_slot_mappings(
449
+ query_start_loc_gpu, self.input_buffers.positions.gpu[:num_tokens]
450
+ )
451
+
452
+ num_computed_tokens_cpu = torch.from_numpy(
453
+ self.req_states.num_computed_tokens[idx_mapping_np]
454
+ )
455
+
456
+ # Logits indices to sample next token from.
457
+ logits_indices = query_start_loc_gpu[1:] - 1
458
+
459
+ # Layer name -> attention metadata.
460
+ attn_metadata = build_attn_metadata(
461
+ attn_metadata_builders=self.attn_metadata_builders,
462
+ num_reqs=num_reqs,
463
+ num_tokens=num_tokens,
464
+ query_start_loc=self.input_buffers.query_start_loc,
465
+ seq_lens=self.input_buffers.seq_lens,
466
+ num_computed_tokens_cpu=num_computed_tokens_cpu,
467
+ block_tables=block_tables,
468
+ slot_mappings=slot_mappings,
469
+ kv_cache_config=self.kv_cache_config,
470
+ )
471
+
472
+ input_ids = self.input_buffers.input_ids.gpu[:num_tokens_after_padding]
473
+ positions = self.input_buffers.positions.gpu[:num_tokens_after_padding]
474
+ return InputBatch(
475
+ req_ids=req_ids,
476
+ num_reqs=num_reqs,
477
+ idx_mapping=idx_mapping,
478
+ idx_mapping_np=idx_mapping_np,
479
+ num_scheduled_tokens=num_scheduled_tokens,
480
+ num_tokens=num_tokens,
481
+ num_tokens_after_padding=num_tokens_after_padding,
482
+ query_start_loc=query_start_loc_gpu,
483
+ query_start_loc_np=query_start_loc_np,
484
+ seq_lens=seq_lens_gpu,
485
+ seq_lens_np=seq_lens_np,
486
+ input_ids=input_ids,
487
+ positions=positions,
488
+ attn_metadata=attn_metadata,
489
+ logits_indices=logits_indices,
490
+ )
491
+
492
+ def sample(
493
+ self,
494
+ hidden_states: torch.Tensor,
495
+ input_batch: InputBatch,
496
+ sampling_metadata: SamplingMetadata,
497
+ grammar_output: GrammarOutput | None,
498
+ ) -> SamplerOutput:
499
+ sample_hidden_states = hidden_states[input_batch.logits_indices]
500
+ logits = self.model.compute_logits(sample_hidden_states)
501
+ if grammar_output is not None:
502
+ # Apply grammar bitmask to the logits in-place.
503
+ with async_barrier(self.structured_outputs_event):
504
+ apply_grammar_bitmask(
505
+ logits,
506
+ input_batch.req_ids,
507
+ grammar_output.structured_output_request_ids,
508
+ grammar_output.grammar_bitmask,
509
+ self.input_buffers,
510
+ )
511
+ sampler_output = self.sampler(logits, sampling_metadata)
512
+ return sampler_output
513
+
514
+ def compute_prompt_logprobs(
515
+ self,
516
+ hidden_states: torch.Tensor,
517
+ input_batch: InputBatch,
518
+ ) -> dict[str, LogprobsTensors]:
519
+ idx_mapping_np = input_batch.idx_mapping_np
520
+ needs_prompt_logprobs = self.req_states.needs_prompt_logprobs[idx_mapping_np]
521
+ if not np.any(needs_prompt_logprobs):
522
+ # No request asks for prompt logprobs.
523
+ return {}
524
+
525
+ num_computed_tokens = self.req_states.num_computed_tokens[idx_mapping_np]
526
+ prompt_lens = self.req_states.prompt_len[idx_mapping_np]
527
+ # NOTE(woosuk): -1 because the last prompt token's hidden state is not
528
+ # needed for prompt logprobs.
529
+ includes_prompt = num_computed_tokens < prompt_lens - 1
530
+ # NOTE(woosuk): If the request was resumed after preemption, its prompt
531
+ # logprobs must have been computed before preemption. Skip.
532
+ resumed_after_prompt = (
533
+ prompt_lens < self.req_states.prefill_len.np[idx_mapping_np]
534
+ )
535
+ needs_prompt_logprobs &= includes_prompt & ~resumed_after_prompt
536
+ if not np.any(needs_prompt_logprobs):
537
+ return {}
538
+
539
+ # Just to be safe, clone the input ids.
540
+ n = input_batch.num_tokens
541
+ # Shift the input ids by one.
542
+ token_ids = torch.empty_like(input_batch.input_ids[:n])
543
+ token_ids[: n - 1] = input_batch.input_ids[1:n]
544
+ # To avoid out-of-bound access, set the last token id to 0.
545
+ token_ids[n - 1] = 0
546
+
547
+ # Handle chunked prompts.
548
+ seq_lens = self.input_buffers.seq_lens.np[: input_batch.num_reqs]
549
+ is_prompt_chunked = seq_lens < prompt_lens
550
+ prefill_token_ids = self.req_states.prefill_token_ids
551
+ query_start_loc = self.input_buffers.query_start_loc.np
552
+ for i, req_id in enumerate(input_batch.req_ids):
553
+ if not needs_prompt_logprobs[i]:
554
+ continue
555
+ if not is_prompt_chunked[i]:
556
+ continue
557
+ # The prompt is chunked. Get the next prompt token.
558
+ req_idx = input_batch.idx_mapping_np[i]
559
+ next_prompt_token = int(prefill_token_ids[req_idx, seq_lens[i]])
560
+ idx = int(query_start_loc[i + 1] - 1)
561
+ # Set the next prompt token.
562
+ # NOTE(woosuk): This triggers a GPU operation.
563
+ token_ids[idx] = next_prompt_token
564
+
565
+ # NOTE(woosuk): We mask out logprobs for negative tokens.
566
+ prompt_logprobs, prompt_ranks = compute_prompt_logprobs(
567
+ token_ids,
568
+ hidden_states[:n],
569
+ self.model.compute_logits,
570
+ )
571
+
572
+ prompt_token_ids = token_ids.unsqueeze(-1)
573
+ prompt_logprobs_dict: dict[str, LogprobsTensors] = {}
574
+ for i, req_id in enumerate(input_batch.req_ids):
575
+ if not needs_prompt_logprobs[i]:
576
+ continue
577
+
578
+ start_idx = query_start_loc[i]
579
+ end_idx = query_start_loc[i + 1]
580
+ assert start_idx < end_idx, (
581
+ f"start_idx ({start_idx}) >= end_idx ({end_idx})"
582
+ )
583
+ logprobs = LogprobsTensors(
584
+ logprob_token_ids=prompt_token_ids[start_idx:end_idx],
585
+ logprobs=prompt_logprobs[start_idx:end_idx],
586
+ selected_token_ranks=prompt_ranks[start_idx:end_idx],
587
+ )
588
+
589
+ req_extra_data = self.req_states.extra_data[req_id]
590
+ prompt_logprobs_list = req_extra_data.in_progress_prompt_logprobs
591
+ if is_prompt_chunked[i]:
592
+ # Prompt is chunked. Do not return the logprobs yet.
593
+ prompt_logprobs_list.append(logprobs)
594
+ continue
595
+
596
+ if prompt_logprobs_list:
597
+ # Merge the in-progress logprobs.
598
+ prompt_logprobs_list.append(logprobs)
599
+ logprobs = LogprobsTensors(
600
+ logprob_token_ids=torch.cat(
601
+ [x.logprob_token_ids for x in prompt_logprobs_list]
602
+ ),
603
+ logprobs=torch.cat([x.logprobs for x in prompt_logprobs_list]),
604
+ selected_token_ranks=torch.cat(
605
+ [x.selected_token_ranks for x in prompt_logprobs_list]
606
+ ),
607
+ )
608
+ prompt_logprobs_list.clear()
609
+
610
+ prompt_logprobs_dict[req_id] = logprobs
611
+ return prompt_logprobs_dict
612
+
613
+ def postprocess(
614
+ self,
615
+ sampler_output: SamplerOutput,
616
+ prompt_logprobs_dict: dict[str, LogprobsTensors],
617
+ input_batch: InputBatch,
618
+ ) -> AsyncOutput | ModelRunnerOutput:
619
+ # Store the last sampled token ids.
620
+ self.req_states.last_sampled_tokens[input_batch.idx_mapping] = (
621
+ sampler_output.sampled_token_ids
622
+ )
623
+ # Get the number of sampled tokens.
624
+ # 0 if chunked-prefilling, 1 if not.
625
+ idx_mapping_np = input_batch.idx_mapping_np
626
+ is_chunked_prefilling = (
627
+ input_batch.seq_lens_np < self.req_states.num_tokens[idx_mapping_np]
628
+ )
629
+ num_sampled_tokens = (~is_chunked_prefilling).astype(np.int32)
630
+ # Increment the number of tokens.
631
+ self.req_states.num_tokens[idx_mapping_np] += num_sampled_tokens
632
+ # Increment the number of computed tokens.
633
+ self.req_states.num_computed_tokens[idx_mapping_np] += (
634
+ input_batch.num_scheduled_tokens
635
+ )
636
+
637
+ model_runner_output = ModelRunnerOutput(
638
+ req_ids=input_batch.req_ids,
639
+ req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
640
+ sampled_token_ids=None,
641
+ logprobs=None,
642
+ prompt_logprobs_dict=prompt_logprobs_dict,
643
+ pooler_output=[],
644
+ kv_connector_output=None,
645
+ num_nans_in_logits=None,
646
+ )
647
+ async_output = AsyncOutput(
648
+ model_runner_output=model_runner_output,
649
+ sampler_output=sampler_output,
650
+ num_sampled_tokens=num_sampled_tokens,
651
+ copy_stream=self.output_copy_stream,
652
+ copy_event=self.output_copy_event,
653
+ )
654
+ if self.use_async_scheduling:
655
+ return async_output
656
+ return async_output.get_output()
657
+
658
+ def get_cudagraph_and_dp_padding(
659
+ self,
660
+ scheduler_output: SchedulerOutput,
661
+ ) -> tuple[CUDAGraphMode, int, torch.Tensor | None]:
662
+ total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
663
+ if self.dp_size == 1:
664
+ # No DP. Only consider CUDA graphs.
665
+ if total_num_scheduled_tokens == 0:
666
+ # Special case: no tokens to run.
667
+ return CUDAGraphMode.NONE, 0, None
668
+
669
+ cudagraph_size = self.cudagraph_manager.get_cudagraph_size(
670
+ scheduler_output, total_num_scheduled_tokens
671
+ )
672
+ if cudagraph_size is not None:
673
+ # Use full CUDA graph.
674
+ return CUDAGraphMode.FULL, cudagraph_size, None
675
+ # Fall back to eager mode.
676
+ # TODO(woosuk): Support piecewise CUDA graphs.
677
+ return CUDAGraphMode.NONE, total_num_scheduled_tokens, None
678
+
679
+ # Consider DP padding and CUDA graph.
680
+ if total_num_scheduled_tokens == 0:
681
+ # Special handling is needed for 0.
682
+ cudagraph_size_before_dp: int | None = 0
683
+ else:
684
+ cudagraph_size_before_dp = self.cudagraph_manager.get_cudagraph_size(
685
+ scheduler_output, total_num_scheduled_tokens
686
+ )
687
+ if cudagraph_size_before_dp is None:
688
+ cudagraph_size_before_dp = -1
689
+
690
+ assert cudagraph_size_before_dp is not None
691
+ num_tokens_across_dp, cudagraph_size_across_dp = get_batch_metadata_across_dp(
692
+ total_num_scheduled_tokens,
693
+ cudagraph_size_before_dp,
694
+ self.dp_size,
695
+ self.dp_rank,
696
+ )
697
+ if all(cudagraph_size_across_dp >= 0):
698
+ # If all ranks can use CUDA graph, pad to the maximum number of tokens
699
+ # across DP and use CUDA graph.
700
+ num_tokens_after_padding = int(cudagraph_size_across_dp.max().item())
701
+ cudagraph_mode = CUDAGraphMode.FULL
702
+ else:
703
+ # If any of the ranks cannot use CUDA graph, use eager mode for all ranks.
704
+ # No padding is needed except for ranks that have no tokens to run.
705
+ num_tokens_across_dp = torch.clamp(num_tokens_across_dp, min=1)
706
+ num_tokens_after_padding = num_tokens_across_dp[self.dp_rank]
707
+ cudagraph_mode = CUDAGraphMode.NONE
708
+ return cudagraph_mode, num_tokens_after_padding, num_tokens_across_dp
709
+
710
+ @torch.inference_mode()
711
+ def execute_model(
712
+ self,
713
+ scheduler_output: SchedulerOutput,
714
+ intermediate_tensors: Any | None = None,
715
+ dummy_run: bool = False,
716
+ ) -> ModelRunnerOutput | None:
717
+ assert intermediate_tensors is None
718
+ if scheduler_output.total_num_scheduled_tokens == 0 and not dummy_run:
719
+ # No need to run the model.
720
+ with async_barrier(self.input_prep_event):
721
+ self.update_states(scheduler_output)
722
+ return EMPTY_MODEL_RUNNER_OUTPUT
723
+
724
+ # NOTE: Call this before the async barrier so CPU all-reduce and
725
+ # GPU execution can overlap.
726
+ cudagraph_mode, num_tokens_after_padding, num_tokens_across_dp = (
727
+ self.get_cudagraph_and_dp_padding(scheduler_output)
728
+ )
729
+ with async_barrier(self.input_prep_event):
730
+ self.update_states(scheduler_output)
731
+ if num_tokens_after_padding == 0:
732
+ # All DP ranks have zero tokens to run.
733
+ return EMPTY_MODEL_RUNNER_OUTPUT
734
+
735
+ if not dummy_run:
736
+ # Common case.
737
+ # Prepare all the inputs and copy to the input buffers.
738
+ input_batch = self.prepare_inputs(
739
+ scheduler_output,
740
+ num_tokens_after_padding,
741
+ )
742
+
743
+ # NOTE(woosuk): Sampling metadata should be built under the async
744
+ # barrier to avoid race conditions.
745
+ pos = input_batch.positions[input_batch.logits_indices]
746
+ sampling_metadata = self.req_states.make_sampling_metadata(
747
+ input_batch.idx_mapping_np, pos
748
+ )
749
+
750
+ if self.lora_config:
751
+ # Activate LoRA adapters.
752
+ lora_inputs = self.req_states.make_lora_inputs(
753
+ input_batch.req_ids,
754
+ input_batch.idx_mapping_np,
755
+ input_batch.num_scheduled_tokens,
756
+ )
757
+ self._set_active_loras(*lora_inputs)
758
+ else:
759
+ # No actual tokens to run. A dummy run for DP.
760
+ num_reqs = min(num_tokens_after_padding, self.max_num_reqs)
761
+ input_batch = InputBatch.make_dummy(
762
+ num_reqs=num_reqs,
763
+ num_tokens=num_tokens_after_padding,
764
+ input_buffers=self.input_buffers,
765
+ device=self.device,
766
+ )
767
+ self.prepare_dummy_attn_metadata(input_batch)
768
+ sampling_metadata = None
769
+
770
+ # Run model.
771
+ if cudagraph_mode == CUDAGraphMode.FULL:
772
+ # Run CUDA graph.
773
+ # NOTE(woosuk): Here, we don't need to pass the input tensors,
774
+ # because they are already copied to the CUDA graph input buffers.
775
+ hidden_states = self.cudagraph_manager.run(
776
+ input_batch.num_tokens_after_padding
777
+ )
778
+ else:
779
+ # Run PyTorch model in eager mode.
780
+ with set_forward_context(
781
+ input_batch.attn_metadata,
782
+ self.vllm_config,
783
+ num_tokens=input_batch.num_tokens_after_padding,
784
+ cudagraph_runtime_mode=cudagraph_mode,
785
+ num_tokens_across_dp=num_tokens_across_dp,
786
+ ):
787
+ hidden_states = self.model(
788
+ input_ids=input_batch.input_ids,
789
+ positions=input_batch.positions,
790
+ )
791
+
792
+ self.execute_model_state = hidden_states, input_batch, sampling_metadata
793
+ return None
794
+
795
+ @torch.inference_mode()
796
+ def sample_tokens(
797
+ self,
798
+ grammar_output: GrammarOutput | None,
799
+ ) -> AsyncOutput | ModelRunnerOutput:
800
+ assert self.execute_model_state is not None
801
+ hidden_states, input_batch, sampling_metadata = self.execute_model_state
802
+ self.execute_model_state = None # type: ignore
803
+ assert sampling_metadata is not None
804
+
805
+ sampler_output = self.sample(
806
+ hidden_states, input_batch, sampling_metadata, grammar_output
807
+ )
808
+ prompt_logprobs_dict = self.compute_prompt_logprobs(hidden_states, input_batch)
809
+ output = self.postprocess(
810
+ sampler_output,
811
+ prompt_logprobs_dict,
812
+ input_batch,
813
+ )
814
+ return output
vllm/v1/worker/gpu/sampler.py
new
+327
@@ -0,0 +1,327 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+
5
+import torch
6
+import triton
7
+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
12
+from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
13
+
14
+
15
+class Sampler:
16
+ def __init__(
17
+ self,
18
+ logprobs_mode: LogprobsMode = "raw_logprobs",
19
+ ):
20
+ if logprobs_mode not in ["processed_logprobs", "raw_logprobs"]:
21
+ raise NotImplementedError(f"Unsupported logprobs_mode: {logprobs_mode}")
22
+ self.logprobs_mode = logprobs_mode
23
+
24
+ def __call__(
25
+ self,
26
+ logits: torch.Tensor,
27
+ sampling_metadata: SamplingMetadata,
28
+ ) -> SamplerOutput:
29
+ if sampling_metadata.max_num_logprobs is not None:
30
+ if self.logprobs_mode == "processed_logprobs":
31
+ sampled, logits = self.sample(
32
+ logits, sampling_metadata, return_logits=True
33
+ )
34
+ else:
35
+ assert self.logprobs_mode == "raw_logprobs"
36
+ sampled, _ = self.sample(logits, sampling_metadata, return_logits=False)
37
+
38
+ logprobs_tensors = compute_topk_logprobs(
39
+ logits,
40
+ sampling_metadata.max_num_logprobs,
41
+ sampled,
42
+ )
43
+ else:
44
+ sampled, _ = self.sample(logits, sampling_metadata, return_logits=False)
45
+ logprobs_tensors = None
46
+
47
+ # These are GPU tensors.
48
+ sampler_output = SamplerOutput(
49
+ # The sampled tokens are expanded to 2D tensor with shape
50
+ # [num_requests, 1], where each row represents one generated
51
+ # token per request.
52
+ sampled_token_ids=sampled.view(-1, 1),
53
+ logprobs_tensors=logprobs_tensors,
54
+ )
55
+ return sampler_output
56
+
57
+ def sample(
58
+ self,
59
+ logits: torch.Tensor,
60
+ sampling_metadata: SamplingMetadata,
61
+ return_logits: bool = False,
62
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
63
+ is_greedy = sampling_metadata.temperature == 0
64
+ temp = torch.where(is_greedy, 1.0, sampling_metadata.temperature)
65
+ logits = logits / temp.view(-1, 1)
66
+ logits = apply_top_k_top_p(
67
+ logits, sampling_metadata.top_k, sampling_metadata.top_p
68
+ )
69
+
70
+ sampled = gumbel_sample(
71
+ logits,
72
+ is_greedy,
73
+ sampling_metadata.seeds,
74
+ sampling_metadata.pos,
75
+ )
76
+ return sampled, logits if return_logits else None
77
+
78
+
79
+@triton.jit
80
+def _gumbel_sample_kernel(
81
+ sampled_ptr,
82
+ logits_ptr,
83
+ logits_stride,
84
+ seeds_ptr,
85
+ pos_ptr,
86
+ is_greedy_ptr,
87
+ vocab_size,
88
+ BLOCK_SIZE: tl.constexpr,
89
+):
90
+ req_idx = tl.program_id(0)
91
+ is_greedy = tl.load(is_greedy_ptr + req_idx)
92
+
93
+ if is_greedy:
94
+ # Greedy sampling. Don't apply gumbel noise.
95
+ max_val = float("-inf")
96
+ max_idx = 0
97
+ for i in range(0, vocab_size, BLOCK_SIZE):
98
+ block = i + tl.arange(0, BLOCK_SIZE)
99
+ mask = block < vocab_size
100
+ logits = tl.load(
101
+ logits_ptr + req_idx * logits_stride + block,
102
+ mask=mask,
103
+ other=float("-inf"),
104
+ )
105
+
106
+ idx = tl.argmax(logits, axis=0)
107
+ value = tl.max(logits, axis=0)
108
+ is_greater = value > max_val
109
+ max_val = tl.where(is_greater, value, max_val)
110
+ max_idx = tl.where(is_greater, i + idx, max_idx)
111
+ tl.store(sampled_ptr + req_idx, max_idx)
112
+ return
113
+
114
+ # Random sampling.
115
+ # Calculate gumbel seed.
116
+ seed = tl.load(seeds_ptr + req_idx)
117
+ pos = tl.load(pos_ptr + req_idx)
118
+ gumbel_seed = tl.randint(seed, pos)
119
+
120
+ max_val = float("-inf")
121
+ max_idx = 0
122
+ for i in range(0, vocab_size, BLOCK_SIZE):
123
+ block = i + tl.arange(0, BLOCK_SIZE)
124
+ mask = block < vocab_size
125
+
126
+ # Generate gumbel noise.
127
+ r = tl.rand(gumbel_seed, block).to(tl.float64)
128
+ gumbel_noise = -tl.log(-tl.log(r + 1e-20) + 1e-20)
129
+ gumbel_noise = gumbel_noise.to(tl.float32)
130
+
131
+ # Apply gumbel noise.
132
+ logits = tl.load(logits_ptr + req_idx * logits_stride + block, mask=mask)
133
+ logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
134
+
135
+ # Argmax to get the sampled token.
136
+ idx = tl.argmax(logits, axis=0)
137
+ value = tl.max(logits, axis=0)
138
+ is_greater = value > max_val
139
+ max_val = tl.where(is_greater, value, max_val)
140
+ max_idx = tl.where(is_greater, i + idx, max_idx)
141
+ tl.store(sampled_ptr + req_idx, max_idx)
142
+
143
+
144
+def gumbel_sample(
145
+ logits: torch.Tensor, # [num_reqs, vocab_size]
146
+ is_greedy: torch.Tensor, # [num_reqs]
147
+ seed: torch.Tensor, # [num_reqs]
148
+ pos: torch.Tensor, # [num_reqs]
149
+) -> torch.Tensor:
150
+ num_reqs, vocab_size = logits.shape
151
+ # NOTE(woosuk): Use int64 for later indexing.
152
+ sampled = torch.empty(
153
+ num_reqs,
154
+ dtype=torch.int64,
155
+ device=logits.device,
156
+ )
157
+ _gumbel_sample_kernel[(num_reqs,)](
158
+ sampled,
159
+ logits,
160
+ logits.stride(0),
161
+ seed,
162
+ pos,
163
+ is_greedy,
164
+ vocab_size,
165
+ num_warps=8,
166
+ BLOCK_SIZE=16384, # type: ignore
167
+ )
168
+ return sampled
169
+
170
+
171
+@triton.jit
172
+def _topk_log_softmax_kernel(
173
+ output_ptr,
174
+ logits_ptr,
175
+ logits_stride,
176
+ topk_ids_ptr,
177
+ topk,
178
+ vocab_size,
179
+ BLOCK_SIZE: tl.constexpr,
180
+ PADDED_TOPK: tl.constexpr,
181
+):
182
+ req_idx = tl.program_id(0)
183
+ row_ptr = logits_ptr + req_idx * logits_stride
184
+
185
+ max_val = float("-inf")
186
+ for i in range(0, vocab_size, BLOCK_SIZE):
187
+ block = i + tl.arange(0, BLOCK_SIZE)
188
+ logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
189
+ max_val = tl.max(tl.maximum(logits, max_val))
190
+ max_val = max_val.to(tl.float32) # type: ignore
191
+
192
+ se = 0.0
193
+ for i in range(0, vocab_size, BLOCK_SIZE):
194
+ block = i + tl.arange(0, BLOCK_SIZE)
195
+ logits = tl.load(row_ptr + block, mask=block < vocab_size, other=0.0)
196
+ # NOTE(woosuk): Make sure that logits and all following operations use FP32.
197
+ logits = logits.to(tl.float32)
198
+ e = tl.exp(logits - max_val)
199
+ e = tl.where(block < vocab_size, e, 0.0)
200
+ se += tl.sum(e)
201
+ lse = tl.log(se)
202
+
203
+ k_offset = tl.arange(0, PADDED_TOPK)
204
+ k_mask = k_offset < topk
205
+ topk_ids = tl.load(topk_ids_ptr + req_idx * topk + k_offset, mask=k_mask, other=0)
206
+
207
+ logits = tl.load(row_ptr + topk_ids, mask=k_mask)
208
+ logits = logits.to(tl.float32)
209
+ o = logits - max_val - lse
210
+ tl.store(output_ptr + req_idx * topk + k_offset, o, mask=k_mask)
211
+
212
+
213
+@triton.jit
214
+def _ranks_kernel(
215
+ output_ptr,
216
+ logits_ptr,
217
+ logits_stride,
218
+ token_ids_ptr,
219
+ vocab_size,
220
+ BLOCK_SIZE: tl.constexpr,
221
+):
222
+ req_idx = tl.program_id(0)
223
+ row_ptr = logits_ptr + req_idx * logits_stride
224
+
225
+ token_id = tl.load(token_ids_ptr + req_idx)
226
+ x = tl.load(row_ptr + token_id)
227
+
228
+ n = 0
229
+ for i in range(0, vocab_size, BLOCK_SIZE):
230
+ block = i + tl.arange(0, BLOCK_SIZE)
231
+ logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
232
+ n += tl.sum((logits > x).to(tl.int32))
233
+ tl.store(output_ptr + req_idx, n)
234
+
235
+
236
+def compute_token_logprobs(
237
+ logits: torch.Tensor,
238
+ token_ids: torch.Tensor,
239
+) -> torch.Tensor:
240
+ batch_size = logits.shape[0]
241
+ vocab_size = logits.shape[1]
242
+ token_ids = token_ids.to(torch.int64)
243
+ num_logprobs = token_ids.shape[1]
244
+ logprobs = torch.empty(
245
+ batch_size,
246
+ num_logprobs,
247
+ dtype=torch.float32,
248
+ device=logits.device,
249
+ )
250
+ _topk_log_softmax_kernel[(batch_size,)](
251
+ logprobs,
252
+ logits,
253
+ logits.stride(0),
254
+ token_ids,
255
+ num_logprobs,
256
+ vocab_size,
257
+ BLOCK_SIZE=1024, # type: ignore
258
+ PADDED_TOPK=triton.next_power_of_2(num_logprobs),
259
+ )
260
+ return logprobs
261
+
262
+
263
+def compute_topk_logprobs(
264
+ logits: torch.Tensor,
265
+ num_logprobs: int,
266
+ sampled_token_ids: torch.Tensor,
267
+) -> LogprobsTensors:
268
+ assert num_logprobs >= 0
269
+ batch_size, vocab_size = logits.shape
270
+ if num_logprobs == 0:
271
+ logprob_token_ids = sampled_token_ids.unsqueeze(-1)
272
+ else:
273
+ topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices
274
+ logprob_token_ids = torch.cat(
275
+ (sampled_token_ids.unsqueeze(-1), topk_indices), dim=1
276
+ )
277
+
278
+ # NOTE(woosuk): Here, to save GPU memory, we do not materialize the full
279
+ # logprobs tensor. Instead, we only compute and return the logprobs of
280
+ # the topk + 1 tokens.
281
+ logprobs = compute_token_logprobs(logits, logprob_token_ids)
282
+ token_ranks = torch.empty(
283
+ batch_size,
284
+ dtype=torch.int64,
285
+ device=logits.device,
286
+ )
287
+ _ranks_kernel[(batch_size,)](
288
+ token_ranks,
289
+ logits,
290
+ logits.stride(0),
291
+ sampled_token_ids,
292
+ vocab_size,
293
+ BLOCK_SIZE=8192, # type: ignore
294
+ )
295
+ return LogprobsTensors(
296
+ logprob_token_ids=logprob_token_ids,
297
+ logprobs=logprobs,
298
+ selected_token_ranks=token_ranks,
299
+ )
300
+
301
+
302
+def compute_prompt_logprobs(
303
+ prompt_token_ids: torch.Tensor,
304
+ prompt_hidden_states: torch.Tensor,
305
+ logits_fn: Callable[[torch.Tensor], torch.Tensor],
306
+) -> tuple[torch.Tensor, torch.Tensor]:
307
+ # Since materializing the full prompt logits can take too much memory,
308
+ # we compute it in chunks.
309
+ CHUNK_SIZE = 1024
310
+ logprobs = []
311
+ ranks = []
312
+ prompt_token_ids = prompt_token_ids.to(torch.int64)
313
+ for start_idx in range(0, prompt_token_ids.shape[0], CHUNK_SIZE):
314
+ end_idx = start_idx + CHUNK_SIZE
315
+ # NOTE(woosuk): logits_fn can be slow because it involves all-gather.
316
+ prompt_logits = logits_fn(prompt_hidden_states[start_idx:end_idx])
317
+ prompt_logprobs = compute_topk_logprobs(
318
+ prompt_logits,
319
+ 0, # num_logprobs
320
+ prompt_token_ids[start_idx:end_idx],
321
+ )
322
+ logprobs.append(prompt_logprobs.logprobs)
323
+ ranks.append(prompt_logprobs.selected_token_ranks)
324
+
325
+ logprobs = torch.cat(logprobs, dim=0) if len(logprobs) > 1 else logprobs[0]
326
+ ranks = torch.cat(ranks, dim=0) if len(ranks) > 1 else ranks[0]
327
+ return logprobs, ranks
vllm/v1/worker/gpu/states.py
new
+265
@@ -0,0 +1,265 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass, field
4
+
5
+import numpy as np
6
+import torch
7
+
8
+from vllm.lora.request import LoRARequest
9
+from vllm.sampling_params import SamplingParams
10
+from vllm.v1.outputs import LogprobsTensors
11
+from vllm.v1.utils import CpuGpuBuffer
12
+
13
+_NP_INT64_MIN = np.iinfo(np.int64).min
14
+_NP_INT64_MAX = np.iinfo(np.int64).max
15
+NO_LORA_ID = 0
16
+
17
+
18
+@dataclass
19
+class SamplingMetadata:
20
+ temperature: torch.Tensor
21
+
22
+ top_p: torch.Tensor | None
23
+ top_k: torch.Tensor | None
24
+
25
+ seeds: torch.Tensor
26
+ pos: torch.Tensor
27
+
28
+ # None means no logprobs, 0 means sampled token logprobs only
29
+ max_num_logprobs: int | None
30
+
31
+ @classmethod
32
+ def make_dummy(
33
+ cls,
34
+ num_reqs: int,
35
+ device: torch.device,
36
+ ) -> "SamplingMetadata":
37
+ assert num_reqs > 0
38
+ temperature = torch.zeros(num_reqs, dtype=torch.float32, device=device)
39
+ temperature[0] = 0.5
40
+ # TODO(woosuk): Use top-p and top-k for dummy sampler.
41
+ # Currently, they are disabled because of memory usage.
42
+ # top_p = torch.full((num_reqs,), 0.95, dtype=torch.float32, device=device)
43
+ # top_k = torch.full((num_reqs,), 20, dtype=torch.int32, device=device)
44
+ top_p = None
45
+ top_k = None
46
+ seeds = torch.zeros(num_reqs, dtype=torch.int64, device=device)
47
+ pos = torch.zeros(num_reqs, dtype=torch.int64, device=device)
48
+ max_num_logprobs = 20
49
+
50
+ return cls(
51
+ temperature=temperature,
52
+ top_p=top_p,
53
+ top_k=top_k,
54
+ seeds=seeds,
55
+ pos=pos,
56
+ max_num_logprobs=max_num_logprobs,
57
+ )
58
+
59
+
60
+class RequestState:
61
+ def __init__(
62
+ self,
63
+ max_num_reqs: int,
64
+ max_model_len: int,
65
+ max_num_batched_tokens: int,
66
+ vocab_size: int,
67
+ device: torch.device,
68
+ pin_memory: bool,
69
+ ):
70
+ self.max_num_reqs = max_num_reqs
71
+ self.max_model_len = max_model_len
72
+ self.max_num_batched_tokens = max_num_batched_tokens
73
+ self.vocab_size = vocab_size
74
+ self.device = device
75
+ self.pin_memory = pin_memory
76
+
77
+ self.req_id_to_index: dict[str, int] = {}
78
+ self.index_to_req_id: dict[int, str] = {}
79
+ self.free_indices = list(range(max_num_reqs))
80
+ self.extra_data: dict[str, ExtraData] = {}
81
+
82
+ self.prompt_len = np.zeros(self.max_num_reqs, dtype=np.int32)
83
+ self.prefill_token_ids = np.zeros(
84
+ (self.max_num_reqs, self.max_model_len),
85
+ dtype=np.int32,
86
+ )
87
+ self.prefill_len = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
88
+ self.num_tokens = np.zeros(self.max_num_reqs, dtype=np.int32)
89
+ self.num_computed_tokens = np.zeros(self.max_num_reqs, dtype=np.int32)
90
+
91
+ # Last sampled tokens.
92
+ self.last_sampled_tokens = torch.zeros(
93
+ self.max_num_reqs,
94
+ 1,
95
+ dtype=torch.int64,
96
+ device=device,
97
+ )
98
+
99
+ # LoRA.
100
+ self.lora_ids = np.zeros(self.max_num_reqs, dtype=np.int32)
101
+ self.lora_ids.fill(NO_LORA_ID)
102
+
103
+ # Sampling parameters.
104
+ self.temperature = self._make_param(self.max_num_reqs, torch.float32)
105
+ self.top_p = self._make_param(self.max_num_reqs, torch.float32)
106
+ self.top_k = self._make_param(self.max_num_reqs, torch.int32)
107
+ self.seeds = self._make_param(self.max_num_reqs, torch.int64)
108
+
109
+ self.num_logprobs = np.empty(self.max_num_reqs, dtype=np.int32)
110
+ # -1 means no logprobs are requested.
111
+ self.num_logprobs.fill(-1)
112
+ self.needs_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=bool)
113
+
114
+ def _make_param(self, size: int, dtype: torch.dtype) -> "Param":
115
+ return Param(size, dtype=dtype, device=self.device, pin_memory=self.pin_memory)
116
+
117
+ def _make_buffer(self, size: int, dtype: torch.dtype) -> CpuGpuBuffer:
118
+ return CpuGpuBuffer(
119
+ size, dtype=dtype, device=self.device, pin_memory=self.pin_memory
120
+ )
121
+
122
+ @property
123
+ def num_reqs(self) -> int:
124
+ return len(self.req_id_to_index)
125
+
126
+ def add_request(
127
+ self,
128
+ req_id: str,
129
+ prompt_len: int,
130
+ prefill_token_ids: list[int],
131
+ num_computed_tokens: int,
132
+ sampling_params: SamplingParams,
133
+ lora_request: LoRARequest | None,
134
+ ) -> None:
135
+ assert len(self.free_indices) > 0, "No free indices"
136
+ req_idx = self.free_indices.pop()
137
+ self.req_id_to_index[req_id] = req_idx
138
+ self.index_to_req_id[req_idx] = req_id
139
+ self.extra_data[req_id] = ExtraData(lora_request)
140
+
141
+ self.prompt_len[req_idx] = prompt_len
142
+ prefill_len = len(prefill_token_ids)
143
+ assert prefill_len >= prompt_len, (
144
+ f"prefill_len {prefill_len} < prompt_len {prompt_len}"
145
+ )
146
+ self.prefill_len.np[req_idx] = prefill_len
147
+ self.prefill_token_ids[req_idx, :prefill_len] = prefill_token_ids
148
+ self.num_tokens[req_idx] = prefill_len
149
+ self.num_computed_tokens[req_idx] = num_computed_tokens
150
+
151
+ if lora_request is not None:
152
+ self.lora_ids[req_idx] = lora_request.lora_int_id
153
+ else:
154
+ self.lora_ids[req_idx] = NO_LORA_ID
155
+
156
+ self.temperature.np[req_idx] = sampling_params.temperature
157
+ self.top_p.np[req_idx] = sampling_params.top_p
158
+ if 0 < sampling_params.top_k < self.vocab_size:
159
+ top_k = sampling_params.top_k
160
+ else:
161
+ top_k = self.vocab_size
162
+ self.top_k.np[req_idx] = top_k
163
+
164
+ if sampling_params.seed is not None:
165
+ seed = sampling_params.seed
166
+ else:
167
+ seed = np.random.randint(_NP_INT64_MIN, _NP_INT64_MAX)
168
+ self.seeds.np[req_idx] = seed
169
+
170
+ if sampling_params.logprobs is not None:
171
+ num_logprobs = sampling_params.logprobs
172
+ else:
173
+ num_logprobs = -1
174
+ self.num_logprobs[req_idx] = num_logprobs
175
+
176
+ # For now, only support prompt logprobs for the prompt tokens.
177
+ needs_prompt_logprobs = sampling_params.prompt_logprobs is not None
178
+ self.needs_prompt_logprobs[req_idx] = needs_prompt_logprobs
179
+
180
+ def remove_request(self, req_id: str) -> None:
181
+ self.extra_data.pop(req_id, None)
182
+ req_idx = self.req_id_to_index.pop(req_id, None)
183
+ if req_idx is None:
184
+ # Request not found.
185
+ return
186
+ self.index_to_req_id.pop(req_idx, None)
187
+ self.free_indices.append(req_idx)
188
+
189
+ def make_sampling_metadata(
190
+ self,
191
+ idx_mapping: np.ndarray,
192
+ pos: torch.Tensor,
193
+ ) -> SamplingMetadata:
194
+ temperature = self.temperature.np[idx_mapping]
195
+ temperature = self.temperature.copy_np_to_gpu(temperature)
196
+
197
+ top_p = self.top_p.np[idx_mapping]
198
+ no_top_p = np.all(top_p == 1.0)
199
+ top_p = self.top_p.copy_np_to_gpu(top_p) if not no_top_p else None
200
+
201
+ top_k = self.top_k.np[idx_mapping]
202
+ no_top_k = np.all(top_k == self.vocab_size)
203
+ top_k = self.top_k.copy_np_to_gpu(top_k) if not no_top_k else None
204
+
205
+ seeds = self.seeds.np[idx_mapping]
206
+ seeds = self.seeds.copy_np_to_gpu(seeds)
207
+
208
+ num_logprobs = self.num_logprobs[idx_mapping]
209
+ max_num_logprobs: int | None = int(np.max(num_logprobs))
210
+ if max_num_logprobs == -1:
211
+ max_num_logprobs = None
212
+
213
+ return SamplingMetadata(
214
+ temperature=temperature,
215
+ top_p=top_p,
216
+ top_k=top_k,
217
+ seeds=seeds,
218
+ pos=pos,
219
+ max_num_logprobs=max_num_logprobs,
220
+ )
221
+
222
+ def make_lora_inputs(
223
+ self,
224
+ req_ids: list[str],
225
+ idx_mapping: np.ndarray,
226
+ num_scheduled_tokens: np.ndarray,
227
+ ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
228
+ lora_ids = self.lora_ids[idx_mapping]
229
+ prompt_lora_mapping = tuple(lora_ids)
230
+ token_lora_mapping = tuple(lora_ids.repeat(num_scheduled_tokens))
231
+
232
+ active_lora_requests: set[LoRARequest] = set()
233
+ for req_id in req_ids:
234
+ lora_request = self.extra_data[req_id].lora_request
235
+ if lora_request is not None:
236
+ active_lora_requests.add(lora_request)
237
+ return prompt_lora_mapping, token_lora_mapping, active_lora_requests
238
+
239
+
240
+class Param:
241
+ def __init__(
242
+ self,
243
+ size: int,
244
+ dtype: torch.dtype,
245
+ device: torch.device,
246
+ pin_memory: bool,
247
+ ):
248
+ self.buffer = CpuGpuBuffer(
249
+ size,
250
+ dtype=dtype,
251
+ device=device,
252
+ pin_memory=pin_memory,
253
+ )
254
+ self.np = np.zeros_like(self.buffer.np)
255
+
256
+ def copy_np_to_gpu(self, x: np.ndarray) -> torch.Tensor:
257
+ n = x.shape[0]
258
+ self.buffer.np[:n] = x
259
+ return self.buffer.copy_to_gpu(n)
260
+
261
+
262
+@dataclass
263
+class ExtraData:
264
+ lora_request: LoRARequest | None
265
+ in_progress_prompt_logprobs: list[LogprobsTensors] = field(default_factory=list)
vllm/v1/worker/gpu/structured_outputs.py
new
+76
@@ -0,0 +1,76 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import numpy as np
4
+import torch
5
+
6
+from vllm.triton_utils import tl, triton
7
+from vllm.v1.worker.gpu.input_batch import InputBuffers
8
+
9
+
10
+def apply_grammar_bitmask(
11
+ logits: torch.Tensor,
12
+ req_ids: list[str],
13
+ grammar_req_ids: list[str],
14
+ grammar_bitmask: np.ndarray,
15
+ input_buffers: InputBuffers,
16
+) -> None:
17
+ input_buffers.grammar_bitmask.np[: grammar_bitmask.shape[0]] = grammar_bitmask
18
+ input_buffers.grammar_bitmask.copy_to_gpu(grammar_bitmask.shape[0])
19
+
20
+ batch_size = logits.shape[0]
21
+ grammar_req_id_to_idx = {req_id: i for i, req_id in enumerate(grammar_req_ids)}
22
+ # logits -> bitmask mapping
23
+ mapping = [grammar_req_id_to_idx.get(req_id, -1) for req_id in req_ids]
24
+ input_buffers.bitmask_indices.np[:batch_size] = mapping
25
+ input_buffers.bitmask_indices.copy_to_gpu(batch_size)
26
+
27
+ vocab_size = logits.shape[-1]
28
+ BLOCK_SIZE = 8192
29
+ grid = (batch_size, triton.cdiv(vocab_size, BLOCK_SIZE))
30
+ _apply_grammar_bitmask_kernel[grid](
31
+ logits,
32
+ logits.stride(0),
33
+ input_buffers.grammar_bitmask.gpu,
34
+ input_buffers.grammar_bitmask.gpu.stride(0),
35
+ input_buffers.bitmask_indices.gpu,
36
+ vocab_size,
37
+ BLOCK_SIZE=BLOCK_SIZE,
38
+ )
39
+
40
+
41
+# Adapted from
42
+# https://github.com/mlc-ai/xgrammar/blob/main/python/xgrammar/kernels/apply_token_bitmask_inplace_triton.py
43
+@triton.jit
44
+def _apply_grammar_bitmask_kernel(
45
+ logits_ptr,
46
+ logits_stride,
47
+ bitmask_ptr,
48
+ bitmask_stride,
49
+ bitmask_indices_ptr,
50
+ vocab_size,
51
+ BLOCK_SIZE: tl.constexpr,
52
+):
53
+ logits_idx = tl.program_id(0)
54
+ bitmask_idx = tl.load(bitmask_indices_ptr + logits_idx)
55
+ if bitmask_idx == -1:
56
+ # No bitmask to apply.
57
+ return
58
+
59
+ # Load the bitmask.
60
+ block_id = tl.program_id(1)
61
+ bitmask_offset = (block_id * BLOCK_SIZE) // 32 + tl.arange(0, BLOCK_SIZE // 32)
62
+ packed_bitmask = tl.load(
63
+ bitmask_ptr + bitmask_idx * bitmask_stride + bitmask_offset,
64
+ mask=bitmask_offset < bitmask_stride,
65
+ )
66
+ # Unpack the bitmask.
67
+ bitmask = ((packed_bitmask[:, None] >> (tl.arange(0, 32)[None, :])) & 1) == 0
68
+ bitmask = bitmask.reshape(BLOCK_SIZE)
69
+
70
+ # Apply the bitmask to the logits.
71
+ block_offset = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
72
+ tl.store(
73
+ logits_ptr + logits_idx * logits_stride + block_offset,
74
+ -float("inf"),
75
+ mask=bitmask & (block_offset < vocab_size),
76
+ )
vllm/v1/worker/gpu_worker.py
+20
-6
@@ -41,7 +41,7 @@ from vllm.sequence import IntermediateTensors
41
from vllm.tasks import SupportedTask
42
from vllm.utils.mem_constants import GiB_bytes
43
from vllm.utils.mem_utils import MemorySnapshot, memory_profiling
44
-from vllm.v1.core.sched.output import GrammarOutput
44
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
45
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
46
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
47
from vllm.v1.outputs import (
@@ -58,7 +58,6 @@ logger = init_logger(__name__)
58
59
if TYPE_CHECKING:
60
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
61
- from vllm.v1.core.sched.output import SchedulerOutput
61
62
63
class Worker(WorkerBase):
@@ -101,6 +100,8 @@ class Worker(WorkerBase):
100
else:
101
self.profiler = None
102
103
+ self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
104
+
105
def sleep(self, level: int = 1) -> None:
106
from vllm.device_allocator.cumem import CuMemAllocator
107
@@ -237,9 +238,17 @@ class Worker(WorkerBase):
238
raise RuntimeError(f"Not support device type: {self.device_config.device}")
239
240
# Construct the model runner
240
- self.model_runner: GPUModelRunner = GPUModelRunner(
241
- self.vllm_config, self.device
242
- )
241
+ if self.use_v2_model_runner:
242
+ from vllm.v1.worker.gpu.model_runner import (
243
+ GPUModelRunner as GPUModelRunnerV2,
244
+ )
245
+
246
+ # HACK(woosuk): This is a temporary fix to avoid type errors.
247
+ self.model_runner: GPUModelRunner = GPUModelRunnerV2( # type: ignore
248
+ self.vllm_config, self.device
249
+ )
250
+ else:
251
+ self.model_runner = GPUModelRunner(self.vllm_config, self.device)
252
253
if self.rank == 0:
254
# If usage stat is enabled, collect relevant info.
@@ -573,7 +582,12 @@ class Worker(WorkerBase):
582
self.profiler.stop()
583
584
def execute_dummy_batch(self) -> None:
576
- self.model_runner._dummy_run(1, uniform_decode=True)
585
+ if self.use_v2_model_runner:
586
+ self.model_runner.execute_model(
587
+ SchedulerOutput.make_empty(), dummy_run=True
588
+ )
589
+ else:
590
+ self.model_runner._dummy_run(1, uniform_decode=True)
591
592
def add_lora(self, lora_request: LoRARequest) -> bool:
593
return self.model_runner.add_lora(lora_request)