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, 16:00 UTC
[Model Runner V2] Change bookkeeping logic in preparation for spec decoding (#29194)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 23, 2025 at 09:42 UTC
7f12c82fa62cee7e815332c1ab2714bdee2a075e
6 files changed
+268
-139
vllm/v1/worker/gpu/async_utils.py
+4
-3
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from contextlib import contextmanager
4
5
-import numpy as np
5
import torch
6
7
from vllm.v1.outputs import (
@@ -18,7 +17,7 @@ class AsyncOutput(AsyncModelRunnerOutput):
17
self,
18
model_runner_output: ModelRunnerOutput,
19
sampler_output: SamplerOutput,
21
- num_sampled_tokens: np.ndarray,
20
+ num_sampled_tokens: torch.Tensor,
21
copy_stream: torch.cuda.Stream,
22
copy_event: torch.cuda.Event,
23
):
@@ -52,6 +51,7 @@ class AsyncOutput(AsyncModelRunnerOutput):
51
)
52
else:
53
self.logprobs_tensors = None
54
+ self.num_sampled_tokens = num_sampled_tokens.to("cpu", non_blocking=True)
55
self.prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}
56
if self.model_runner_output.prompt_logprobs_dict:
57
for k, v in self.model_runner_output.prompt_logprobs_dict.items():
@@ -63,6 +63,7 @@ class AsyncOutput(AsyncModelRunnerOutput):
63
64
def get_output(self) -> ModelRunnerOutput:
65
self.copy_event.synchronize()
66
+ num_sampled_tokens_np = self.num_sampled_tokens.numpy()
67
68
# NOTE(woosuk): The following code is to ensure compatibility with
69
# the existing model runner.
@@ -71,7 +72,7 @@ class AsyncOutput(AsyncModelRunnerOutput):
72
sampled_token_ids: list[list[int]] = self.sampled_token_ids.tolist()
73
num_reqs = len(sampled_token_ids)
74
for i in range(num_reqs):
74
- del sampled_token_ids[i][self.num_sampled_tokens[i] :]
75
+ del sampled_token_ids[i][num_sampled_tokens_np[i] :]
76
self.model_runner_output.sampled_token_ids = sampled_token_ids
77
78
if self.logprobs_tensors is not None:
vllm/v1/worker/gpu/attn_utils.py
+8
-6
@@ -3,6 +3,7 @@
3
from collections.abc import Sequence
4
from typing import Any, cast
5
6
+import numpy as np
7
import torch
8
9
from vllm.attention.backends.abstract import AttentionBackend
@@ -145,8 +146,9 @@ def build_attn_metadata(
146
num_reqs: int,
147
num_tokens: int,
148
query_start_loc: CpuGpuBuffer,
148
- seq_lens: CpuGpuBuffer,
149
- num_computed_tokens_cpu: torch.Tensor,
149
+ seq_lens: torch.Tensor,
150
+ seq_lens_np: np.ndarray,
151
+ num_computed_tokens_cpu: torch.Tensor | None,
152
block_tables: Sequence[torch.Tensor],
153
slot_mappings: torch.Tensor,
154
kv_cache_config: KVCacheConfig,
@@ -154,9 +156,9 @@ def build_attn_metadata(
156
query_start_loc_gpu = query_start_loc.gpu[: num_reqs + 1]
157
query_start_loc_cpu = query_start_loc.cpu[: num_reqs + 1]
158
max_query_len = int(query_start_loc.np[: num_reqs + 1].max())
157
- seq_lens_gpu = seq_lens.gpu[:num_reqs]
158
- seq_lens_cpu = seq_lens.cpu[:num_reqs]
159
- max_seq_len = int(seq_lens.np[:num_reqs].max())
159
+ seq_lens = seq_lens[:num_reqs]
160
+ seq_lens_cpu = torch.from_numpy(seq_lens_np)
161
+ max_seq_len = int(seq_lens_np.max())
162
163
attn_metadata: dict[str, Any] = {}
164
kv_cache_groups = kv_cache_config.kv_cache_groups
@@ -167,7 +169,7 @@ def build_attn_metadata(
169
common_attn_metadata = CommonAttentionMetadata(
170
query_start_loc=query_start_loc_gpu,
171
query_start_loc_cpu=query_start_loc_cpu,
170
- seq_lens=seq_lens_gpu,
172
+ seq_lens=seq_lens,
173
seq_lens_cpu=seq_lens_cpu,
174
max_seq_len=max_seq_len,
175
num_computed_tokens_cpu=num_computed_tokens_cpu,
vllm/v1/worker/gpu/cudagraph_utils.py
+4
-4
@@ -101,14 +101,13 @@ class CudaGraphManager:
101
102
# Prepare dummy inputs.
103
input_ids = input_buffers.input_ids.gpu[:batch_size]
104
- positions = input_buffers.positions.gpu[:batch_size]
104
+ positions = input_buffers.positions[:batch_size]
105
106
input_buffers.query_start_loc.np[: batch_size + 1] = np.arange(batch_size + 1)
107
input_buffers.query_start_loc.np[batch_size:] = batch_size
108
input_buffers.query_start_loc.copy_to_gpu()
109
- input_buffers.seq_lens.np[:batch_size] = self.max_model_len
110
- input_buffers.seq_lens.np[batch_size:] = 0
111
- input_buffers.seq_lens.copy_to_gpu()
109
+ input_buffers.seq_lens[:batch_size] = self.max_model_len
110
+ input_buffers.seq_lens[batch_size:] = 0
111
112
input_block_tables = [x[:batch_size] for x in block_tables.input_block_tables]
113
slot_mappings = block_tables.slot_mappings[:, :batch_size]
@@ -119,6 +118,7 @@ class CudaGraphManager:
118
num_tokens=batch_size,
119
query_start_loc=input_buffers.query_start_loc,
120
seq_lens=input_buffers.seq_lens,
121
+ seq_lens_np=np.full(batch_size, self.max_model_len, dtype=np.int32),
122
num_computed_tokens_cpu=None, # FIXME
123
block_tables=input_block_tables,
124
slot_mappings=slot_mappings,
vllm/v1/worker/gpu/input_batch.py
+127
-47
@@ -32,9 +32,9 @@ class InputBuffers:
32
33
self.idx_mapping = self._make_buffer(max_num_reqs, dtype=torch.int32)
34
self.input_ids = self._make_buffer(max_num_tokens, dtype=torch.int32)
35
- self.positions = self._make_buffer(max_num_tokens, dtype=torch.int64)
35
+ self.positions = torch.zeros(max_num_tokens, dtype=torch.int64, device=device)
36
self.query_start_loc = self._make_buffer(max_num_reqs + 1, dtype=torch.int32)
37
- self.seq_lens = self._make_buffer(max_num_reqs, dtype=torch.int32)
37
+ self.seq_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
38
39
# Structured outputs.
40
self.bitmask_indices = self._make_buffer(max_num_reqs, dtype=torch.int32)
@@ -107,13 +107,15 @@ class InputBatch:
107
query_start_loc_np = input_buffers.query_start_loc.np[: num_reqs + 1]
108
query_start_loc = input_buffers.query_start_loc.copy_to_gpu()[: num_reqs + 1]
109
# seq_len equals to query_len
110
- input_buffers.seq_lens.np[:num_reqs] = num_scheduled_tokens
111
- input_buffers.seq_lens.np[num_reqs:] = 0
112
- seq_lens_np = input_buffers.seq_lens.np[:num_reqs]
113
- seq_lens = input_buffers.seq_lens.copy_to_gpu()[:num_reqs]
110
+ seq_lens_np = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)
111
+ seq_lens_np[-1] += num_tokens % num_reqs
112
+ input_buffers.seq_lens[:num_reqs] = num_tokens // num_reqs
113
+ input_buffers.seq_lens[num_reqs - 1] += num_tokens % num_reqs
114
+ input_buffers.seq_lens[num_reqs:] = 0
115
+ seq_lens = input_buffers.seq_lens[:num_reqs]
116
117
input_ids = input_buffers.input_ids.copy_to_gpu(num_tokens)
116
- positions = input_buffers.positions.copy_to_gpu(num_tokens)
118
+ positions = input_buffers.positions[:num_tokens]
119
# attn_metadata = defaultdict(lambda: None)
120
logits_indices = query_start_loc[1:] - 1
121
return cls(
@@ -141,27 +143,25 @@ class InputBatch:
143
[
144
types.none(
145
types.int32[:], # idx_mapping
144
- types.int32[:, :], # token_ids
145
- types.int32[:], # num_computed_tokens
146
types.int32[:], # num_scheduled_tokens
147
+ types.int32[:, :], # prefill_token_ids
148
+ types.int32[:], # num_computed_prefill_tokens
149
+ types.int32[:], # prefill_len
150
types.int32[:], # input_ids
148
- types.int64[:], # positions
151
types.int32[:], # query_start_loc
150
- types.int32[:], # seq_lens
152
)
153
],
154
nopython=True,
155
cache=True,
156
)
156
-def _prepare_inputs(
157
+def _prepare_prefill_inputs(
158
idx_mapping: np.ndarray, # batch_idx -> req_idx
158
- token_ids: np.ndarray, # [N, max_model_len]
159
- num_computed_tokens: np.ndarray, # [N]
159
num_scheduled_tokens: np.ndarray, # [B]
160
+ prefill_token_ids: np.ndarray, # [N, max_model_len]
161
+ num_computed_prefill_tokens: np.ndarray, # [N]
162
+ prefill_len: np.ndarray, # [N]
163
input_ids: np.ndarray, # [num_input_tokens]
162
- positions: np.ndarray, # [num_input_tokens]
164
query_start_loc: np.ndarray, # [B + 1]
164
- seq_lens: np.ndarray, # [B]
165
) -> None:
166
num_reqs = num_scheduled_tokens.shape[0]
167
query_start_loc[0] = 0
@@ -170,62 +170,112 @@ def _prepare_inputs(
170
for i in range(num_reqs):
171
req_idx = idx_mapping[i]
172
query_len = num_scheduled_tokens[i]
173
- start = num_computed_tokens[req_idx]
174
- end = start + query_len
175
- seq_lens[i] = end
173
+
174
+ start = num_computed_prefill_tokens[req_idx]
175
+ end = min(start + query_len, prefill_len[req_idx])
176
+ n = end - start
177
178
start_idx = cu_num_tokens
178
- end_idx = start_idx + query_len
179
- input_ids[start_idx:end_idx] = token_ids[req_idx, start:end]
180
- positions[start_idx:end_idx] = np.arange(start, end, dtype=np.int64)
179
+ input_ids[start_idx : start_idx + n] = prefill_token_ids[req_idx, start:end]
180
182
- cu_num_tokens = end_idx
181
+ cu_num_tokens = start_idx + query_len
182
query_start_loc[i + 1] = cu_num_tokens
183
184
# Pad the inputs for CUDA graphs.
185
# Note: pad query_start_loc to be non-decreasing, as kernels
186
# like FlashAttention requires that
187
query_start_loc[num_reqs + 1 :].fill(cu_num_tokens)
189
- # Fill unused with 0 for full cuda graph mode.
190
- seq_lens[num_reqs:].fill(0)
188
189
193
-def prepare_inputs(
190
+def prepare_prefill_inputs(
191
idx_mapping: np.ndarray,
195
- prefill_token_ids: np.ndarray,
196
- num_computed_tokens: np.ndarray,
192
num_scheduled_tokens: np.ndarray,
193
+ total_num_tokens: int,
194
+ prefill_token_ids: np.ndarray,
195
+ num_computed_prefill_tokens: np.ndarray,
196
+ prefill_len: np.ndarray,
197
input_ids: CpuGpuBuffer,
199
- positions: CpuGpuBuffer,
198
query_start_loc: CpuGpuBuffer,
201
- seq_lens: CpuGpuBuffer,
202
- num_tokens: int,
199
) -> None:
204
- _prepare_inputs(
200
+ _prepare_prefill_inputs(
201
idx_mapping,
206
- prefill_token_ids,
207
- num_computed_tokens,
202
num_scheduled_tokens,
203
+ prefill_token_ids,
204
+ num_computed_prefill_tokens,
205
+ prefill_len,
206
input_ids.np,
210
- positions.np,
207
query_start_loc.np,
212
- seq_lens.np,
208
)
214
- input_ids.copy_to_gpu(num_tokens)
215
- positions.copy_to_gpu(num_tokens)
209
+ input_ids.copy_to_gpu(total_num_tokens)
210
# NOTE(woosuk): We should copy the whole query_start_loc and seq_lens
211
# tensors from CPU to GPU, because they may include paddings needed
212
# for full CUDA graph mode.
213
query_start_loc.copy_to_gpu()
220
- seq_lens.copy_to_gpu()
221
- return
214
215
216
@triton.jit
225
-def _combine_last_token_ids_kernel(
217
+def _prepare_pos_seq_lens_kernel(
218
+ pos_ptr,
219
+ seq_lens_ptr,
220
+ idx_mapping_ptr,
221
+ query_start_loc_ptr,
222
+ num_computed_tokens_ptr,
223
+ max_num_reqs,
224
+ BLOCK_SIZE: tl.constexpr,
225
+):
226
+ req_id = tl.program_id(0)
227
+ num_reqs = tl.num_programs(0) - 1
228
+ if req_id == num_reqs:
229
+ # Pad unused seq_lens as 0 for full CUDA graphs.
230
+ for i in tl.range(num_reqs, max_num_reqs, BLOCK_SIZE):
231
+ block = i + tl.arange(0, BLOCK_SIZE)
232
+ mask = block < max_num_reqs
233
+ tl.store(seq_lens_ptr + block, 0, mask=mask)
234
+ return
235
+
236
+ req_state_idx = tl.load(idx_mapping_ptr + req_id)
237
+ num_computed_tokens = tl.load(num_computed_tokens_ptr + req_state_idx)
238
+
239
+ start = tl.load(query_start_loc_ptr + req_id)
240
+ end = tl.load(query_start_loc_ptr + req_id + 1)
241
+ query_len = end - start
242
+
243
+ seq_len = num_computed_tokens + query_len
244
+ tl.store(seq_lens_ptr + req_id, seq_len)
245
+
246
+ for i in tl.range(0, query_len, BLOCK_SIZE):
247
+ block = i + tl.arange(0, BLOCK_SIZE)
248
+ mask = block < query_len
249
+ pos = num_computed_tokens + block
250
+ tl.store(pos_ptr + start + block, pos, mask=mask)
251
+
252
+
253
+def prepare_pos_seq_lens(
254
+ idx_mapping: torch.Tensor,
255
+ query_start_loc: torch.Tensor,
256
+ num_computed_tokens: torch.Tensor,
257
+ pos: torch.Tensor,
258
+ seq_lens: torch.Tensor,
259
+) -> None:
260
+ num_reqs = idx_mapping.shape[0]
261
+ # NOTE(woosuk): We do +1 because the last thread block is used
262
+ # to pad unused seq_lens as 0 for full CUDA graphs.
263
+ _prepare_pos_seq_lens_kernel[(num_reqs + 1,)](
264
+ pos,
265
+ seq_lens,
266
+ idx_mapping,
267
+ query_start_loc,
268
+ num_computed_tokens,
269
+ seq_lens.shape[0],
270
+ BLOCK_SIZE=1024,
271
+ )
272
+
273
+
274
+@triton.jit
275
+def _combine_sampled_and_draft_tokens_kernel(
276
input_ids_ptr,
277
idx_mapping_ptr,
228
- last_token_ids_ptr,
278
+ last_sampled_tokens_ptr,
279
query_start_loc_ptr,
280
seq_lens_ptr,
281
prefill_len_ptr,
@@ -239,26 +289,56 @@ def _combine_last_token_ids_kernel(
289
# Handling prefill tokens.
290
return
291
242
- last_token_id = tl.load(last_token_ids_ptr + req_state_idx)
292
+ last_token_id = tl.load(last_sampled_tokens_ptr + req_state_idx)
293
end = tl.load(query_start_loc_ptr + batch_idx + 1)
294
tl.store(input_ids_ptr + end - 1, last_token_id)
295
296
247
-def combine_last_token_ids(
297
+def combine_sampled_and_draft_tokens(
298
input_ids: torch.Tensor,
299
idx_mapping: torch.Tensor,
250
- last_token_ids: torch.Tensor,
300
+ last_sampled_tokens: torch.Tensor,
301
query_start_loc: torch.Tensor,
302
seq_lens: torch.Tensor,
303
prefill_len: torch.Tensor,
304
) -> torch.Tensor:
305
num_reqs = seq_lens.shape[0]
256
- _combine_last_token_ids_kernel[(num_reqs,)](
306
+ _combine_sampled_and_draft_tokens_kernel[(num_reqs,)](
307
input_ids,
308
idx_mapping,
259
- last_token_ids,
309
+ last_sampled_tokens,
310
query_start_loc,
311
seq_lens,
312
prefill_len,
313
)
314
return input_ids
315
+
316
+
317
+@triton.jit
318
+def _update_num_computed_tokens_kernel(
319
+ idx_mapping_ptr,
320
+ num_computed_tokens_ptr,
321
+ query_start_loc_ptr,
322
+):
323
+ req_id = tl.program_id(0)
324
+ req_state_idx = tl.load(idx_mapping_ptr + req_id)
325
+
326
+ start = tl.load(query_start_loc_ptr + req_id)
327
+ end = tl.load(query_start_loc_ptr + req_id + 1)
328
+ query_len = end - start
329
+
330
+ n = tl.load(num_computed_tokens_ptr + req_state_idx)
331
+ tl.store(num_computed_tokens_ptr + req_state_idx, n + query_len)
332
+
333
+
334
+def update_num_computed_tokens(
335
+ idx_mapping: torch.Tensor,
336
+ num_computed_tokens: torch.Tensor,
337
+ query_start_loc: torch.Tensor,
338
+) -> None:
339
+ num_reqs = idx_mapping.shape[0]
340
+ _update_num_computed_tokens_kernel[(num_reqs,)](
341
+ idx_mapping,
342
+ num_computed_tokens,
343
+ query_start_loc,
344
+ )
vllm/v1/worker/gpu/model_runner.py
+115
-76
@@ -39,8 +39,10 @@ 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,
42
+ combine_sampled_and_draft_tokens,
43
+ prepare_pos_seq_lens,
44
+ prepare_prefill_inputs,
45
+ update_num_computed_tokens,
46
)
47
from vllm.v1.worker.gpu.sampler import Sampler, compute_prompt_logprobs
48
from vllm.v1.worker.gpu.states import RequestState, SamplingMetadata
@@ -179,6 +181,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
181
self.vllm_config,
182
self.device,
183
)
184
+ # TODO(woosuk): Support other backends.
185
+ if not all(b.get_name() == "FLASH_ATTN" for b in self.attn_backends.values()):
186
+ raise NotImplementedError("Only FLASH_ATTN backend is supported currently.")
187
188
self.kv_caches: list[torch.Tensor] = []
189
init_kv_cache(
@@ -196,8 +201,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
201
slot_mappings = self.block_tables.get_dummy_slot_mappings(
202
input_batch.num_tokens
203
)
199
- num_computed_tokens_cpu = torch.zeros(
200
- input_batch.num_reqs, dtype=torch.int32, device="cpu"
204
+ num_computed_tokens = torch.zeros(
205
+ input_batch.num_reqs, dtype=torch.int32, device=self.device
206
)
207
attn_metadata = build_attn_metadata(
208
attn_metadata_builders=self.attn_metadata_builders,
@@ -205,7 +210,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
210
num_tokens=input_batch.num_tokens,
211
query_start_loc=self.input_buffers.query_start_loc,
212
seq_lens=self.input_buffers.seq_lens,
208
- num_computed_tokens_cpu=num_computed_tokens_cpu,
213
+ seq_lens_np=input_batch.seq_lens_np,
214
+ num_computed_tokens_cpu=num_computed_tokens,
215
block_tables=block_tables,
216
slot_mappings=slot_mappings,
217
kv_cache_config=self.kv_cache_config,
@@ -368,6 +374,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
374
cu_num_new_blocks[i].append(x + len(block_ids))
375
new_block_ids[i].extend(block_ids)
376
overwrite.append(True)
377
+ # Update the GPU tensors for request states.
378
+ if scheduler_output.scheduled_new_reqs:
379
+ self.req_states.prefill_len.copy_to_gpu()
380
381
# Add new blocks for the existing requests.
382
cached_reqs = scheduler_output.scheduled_cached_reqs
@@ -421,46 +430,60 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
430
# Block tables: num_kv_cache_groups x [num_reqs, max_num_blocks]
431
block_tables = self.block_tables.gather_block_tables(idx_mapping)
432
424
- prepare_inputs(
433
+ # Copy prefill tokens from CPU to GPU and get query_start_loc.
434
+ prepare_prefill_inputs(
435
idx_mapping_np,
426
- self.req_states.prefill_token_ids,
427
- self.req_states.num_computed_tokens,
436
num_scheduled_tokens,
437
+ num_tokens,
438
+ self.req_states.prefill_token_ids,
439
+ self.req_states.num_computed_prefill_tokens,
440
+ self.req_states.prefill_len.np,
441
self.input_buffers.input_ids,
430
- self.input_buffers.positions,
442
self.input_buffers.query_start_loc,
432
- self.input_buffers.seq_lens,
433
- num_tokens,
443
)
435
-
444
query_start_loc = self.input_buffers.query_start_loc
445
query_start_loc_gpu = query_start_loc.gpu[: num_reqs + 1]
446
query_start_loc_np = query_start_loc.np[: num_reqs + 1]
439
- seq_lens_gpu = self.input_buffers.seq_lens.gpu[:num_reqs]
440
- seq_lens_np = self.input_buffers.seq_lens.np[:num_reqs]
447
442
- # Some input token ids are directly read from the last sampled tokens.
443
- combine_last_token_ids(
448
+ # Prepare positions and seq_lens.
449
+ prepare_pos_seq_lens(
450
+ idx_mapping,
451
+ query_start_loc_gpu,
452
+ self.req_states.num_computed_tokens,
453
+ self.input_buffers.positions,
454
+ self.input_buffers.seq_lens,
455
+ )
456
+ seq_lens = self.input_buffers.seq_lens[:num_reqs]
457
+
458
+ # Some input token ids are directly read from the last sampled tokens
459
+ # and draft tokens.
460
+ combine_sampled_and_draft_tokens(
461
self.input_buffers.input_ids.gpu,
462
idx_mapping,
463
self.req_states.last_sampled_tokens,
464
query_start_loc_gpu,
448
- seq_lens_gpu,
449
- self.req_states.prefill_len.copy_to_gpu(),
465
+ seq_lens,
466
+ self.req_states.prefill_len.gpu,
467
)
468
469
# Compute slot mappings: [num_kv_cache_groups, num_tokens]
470
slot_mappings = self.block_tables.compute_slot_mappings(
454
- query_start_loc_gpu, self.input_buffers.positions.gpu[:num_tokens]
455
- )
456
-
457
- num_computed_tokens_cpu = torch.from_numpy(
458
- self.req_states.num_computed_tokens[idx_mapping_np]
471
+ query_start_loc_gpu, self.input_buffers.positions[:num_tokens]
472
)
473
474
# Logits indices to sample next token from.
475
logits_indices = query_start_loc_gpu[1:] - 1
476
477
+ # Get num_computed_tokens.
478
+ # HACK(woosuk): Here, we use num_computed_tokens on GPU instead of
479
+ # num_computed_tokens_cpu. This works for most cases.
480
+ num_computed_tokens = self.req_states.num_computed_tokens[idx_mapping]
481
+ # HACK(woosuk): Only GPU has the exact seq_lens because at this point
482
+ # CPU does not know how many draft tokens are accepted/rejected in the
483
+ # previous step. Therefore, we use max_model_len to be safe.
484
+ # NOTE(woosuk): This only works for FA3 backend.
485
+ seq_lens_np = np.full(num_reqs, self.max_model_len, dtype=np.int32)
486
+
487
# Layer name -> attention metadata.
488
attn_metadata = build_attn_metadata(
489
attn_metadata_builders=self.attn_metadata_builders,
@@ -468,14 +491,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
491
num_tokens=num_tokens,
492
query_start_loc=self.input_buffers.query_start_loc,
493
seq_lens=self.input_buffers.seq_lens,
471
- num_computed_tokens_cpu=num_computed_tokens_cpu,
494
+ seq_lens_np=seq_lens_np,
495
+ num_computed_tokens_cpu=num_computed_tokens,
496
block_tables=block_tables,
497
slot_mappings=slot_mappings,
498
kv_cache_config=self.kv_cache_config,
499
)
500
501
input_ids = self.input_buffers.input_ids.gpu[:num_tokens_after_padding]
478
- positions = self.input_buffers.positions.gpu[:num_tokens_after_padding]
502
+ positions = self.input_buffers.positions[:num_tokens_after_padding]
503
return InputBatch(
504
req_ids=req_ids,
505
num_reqs=num_reqs,
@@ -486,7 +510,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
510
num_tokens_after_padding=num_tokens_after_padding,
511
query_start_loc=query_start_loc_gpu,
512
query_start_loc_np=query_start_loc_np,
489
- seq_lens=seq_lens_gpu,
513
+ seq_lens=seq_lens,
514
seq_lens_np=seq_lens_np,
515
input_ids=input_ids,
516
positions=positions,
@@ -500,11 +524,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
524
input_batch: InputBatch,
525
sampling_metadata: SamplingMetadata,
526
grammar_output: GrammarOutput | None,
503
- ) -> SamplerOutput:
527
+ ) -> tuple[SamplerOutput, torch.Tensor]:
528
sample_hidden_states = hidden_states[input_batch.logits_indices]
529
logits = self.model.compute_logits(sample_hidden_states)
530
if grammar_output is not None:
531
# Apply grammar bitmask to the logits in-place.
532
+ # TODO(woosuk): Make compatible with spec decoding.
533
with async_barrier(self.structured_outputs_event):
534
apply_grammar_bitmask(
535
logits,
@@ -513,8 +538,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
538
grammar_output.grammar_bitmask,
539
self.input_buffers,
540
)
541
+
542
sampler_output = self.sampler(logits, sampling_metadata)
517
- return sampler_output
543
+ # Get the number of sampled tokens.
544
+ # 0 if chunked-prefilling, 1 if not.
545
+ prefill_len = self.req_states.prefill_len.gpu[input_batch.idx_mapping]
546
+ is_chunked_prefilling = input_batch.seq_lens < prefill_len
547
+ num_sampled = (~is_chunked_prefilling).int()
548
+ return sampler_output, num_sampled
549
550
def compute_prompt_logprobs(
551
self,
@@ -527,11 +558,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
558
# No request asks for prompt logprobs.
559
return {}
560
530
- num_computed_tokens = self.req_states.num_computed_tokens[idx_mapping_np]
561
prompt_lens = self.req_states.prompt_len[idx_mapping_np]
562
# NOTE(woosuk): -1 because the last prompt token's hidden state is not
563
# needed for prompt logprobs.
534
- includes_prompt = num_computed_tokens < prompt_lens - 1
564
+ computed_prefill = self.req_states.num_computed_prefill_tokens[idx_mapping_np]
565
+ includes_prompt = computed_prefill < prompt_lens - 1
566
# NOTE(woosuk): If the request was resumed after preemption, its prompt
567
# logprobs must have been computed before preemption. Skip.
568
resumed_after_prompt = (
@@ -550,8 +581,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
581
token_ids[n - 1] = 0
582
583
# Handle chunked prompts.
553
- seq_lens = self.input_buffers.seq_lens.np[: input_batch.num_reqs]
554
- is_prompt_chunked = seq_lens < prompt_lens
584
+ pos_after_step = computed_prefill + input_batch.num_scheduled_tokens
585
+ is_prompt_chunked = pos_after_step < prompt_lens
586
prefill_token_ids = self.req_states.prefill_token_ids
587
query_start_loc = self.input_buffers.query_start_loc.np
588
for i, req_id in enumerate(input_batch.req_ids):
@@ -561,7 +592,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
592
continue
593
# The prompt is chunked. Get the next prompt token.
594
req_idx = input_batch.idx_mapping_np[i]
564
- next_prompt_token = int(prefill_token_ids[req_idx, seq_lens[i]])
595
+ next_prompt_token = int(prefill_token_ids[req_idx, pos_after_step[i]])
596
idx = int(query_start_loc[i + 1] - 1)
597
# Set the next prompt token.
598
# NOTE(woosuk): This triggers a GPU operation.
@@ -617,48 +648,27 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
648
649
def postprocess(
650
self,
620
- sampler_output: SamplerOutput,
621
- prompt_logprobs_dict: dict[str, LogprobsTensors],
651
input_batch: InputBatch,
623
- ) -> AsyncOutput | ModelRunnerOutput:
624
- # Store the last sampled token ids.
625
- self.req_states.last_sampled_tokens[input_batch.idx_mapping] = (
626
- sampler_output.sampled_token_ids
652
+ sampled_tokens: torch.Tensor,
653
+ num_sampled: torch.Tensor,
654
+ ) -> None:
655
+ # Update the number of computed tokens.
656
+ update_num_computed_tokens(
657
+ input_batch.idx_mapping,
658
+ self.req_states.num_computed_tokens,
659
+ input_batch.query_start_loc,
660
)
628
- # Get the number of sampled tokens.
629
- # 0 if chunked-prefilling, 1 if not.
661
idx_mapping_np = input_batch.idx_mapping_np
631
- is_chunked_prefilling = (
632
- input_batch.seq_lens_np < self.req_states.num_tokens[idx_mapping_np]
633
- )
634
- num_sampled_tokens = (~is_chunked_prefilling).astype(np.int32)
635
- # Increment the number of tokens.
636
- self.req_states.num_tokens[idx_mapping_np] += num_sampled_tokens
637
- # Increment the number of computed tokens.
638
- self.req_states.num_computed_tokens[idx_mapping_np] += (
639
- input_batch.num_scheduled_tokens
662
+ computed_prefill = self.req_states.num_computed_prefill_tokens
663
+ # TODO(woosuk): Simplify this.
664
+ computed_prefill[idx_mapping_np] = np.minimum(
665
+ computed_prefill[idx_mapping_np] + input_batch.num_scheduled_tokens,
666
+ self.req_states.prefill_len.np[idx_mapping_np],
667
)
668
642
- model_runner_output = ModelRunnerOutput(
643
- req_ids=input_batch.req_ids,
644
- req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
645
- sampled_token_ids=None, # type: ignore
646
- logprobs=None,
647
- prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore
648
- pooler_output=[],
649
- kv_connector_output=None,
650
- num_nans_in_logits=None,
651
- )
652
- async_output = AsyncOutput(
653
- model_runner_output=model_runner_output,
654
- sampler_output=sampler_output,
655
- num_sampled_tokens=num_sampled_tokens,
656
- copy_stream=self.output_copy_stream,
657
- copy_event=self.output_copy_event,
658
- )
659
- if self.use_async_scheduling:
660
- return async_output
661
- return async_output.get_output()
669
+ # Store the last sampled token ids.
670
+ last_sampled = sampled_tokens
671
+ self.req_states.last_sampled_tokens[input_batch.idx_mapping] = last_sampled
672
673
def get_cudagraph_and_dp_padding(
674
self,
@@ -782,6 +792,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
792
)
793
else:
794
# Run PyTorch model in eager mode.
795
+ # TODO(woosuk): Support piecewise CUDA graph.
796
with set_forward_context(
797
input_batch.attn_metadata,
798
self.vllm_config,
@@ -807,13 +818,41 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
818
self.execute_model_state = None # type: ignore
819
assert sampling_metadata is not None
820
810
- sampler_output = self.sample(
821
+ sampler_output, num_sampled_tokens = self.sample(
822
hidden_states, input_batch, sampling_metadata, grammar_output
823
)
824
prompt_logprobs_dict = self.compute_prompt_logprobs(hidden_states, input_batch)
814
- output = self.postprocess(
815
- sampler_output,
816
- prompt_logprobs_dict,
817
- input_batch,
825
+
826
+ # Prepare the model runner output.
827
+ model_runner_output = ModelRunnerOutput(
828
+ req_ids=input_batch.req_ids,
829
+ # NOTE(woosuk): req_id_to_index is unused in this model runner.
830
+ # Only for compatibility with the existing model runner and scheduler.
831
+ req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
832
+ sampled_token_ids=None, # type: ignore
833
+ logprobs=None,
834
+ prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore
835
+ pooler_output=[],
836
+ kv_connector_output=None,
837
+ num_nans_in_logits=None,
838
+ )
839
+ async_output = AsyncOutput(
840
+ model_runner_output=model_runner_output,
841
+ sampler_output=sampler_output,
842
+ num_sampled_tokens=num_sampled_tokens,
843
+ copy_stream=self.output_copy_stream,
844
+ copy_event=self.output_copy_event,
845
+ )
846
+
847
+ # Postprocess results and update request states.
848
+ # NOTE: This is intentionally done after creating the AsyncOutput,
849
+ # ensuring that `copy_event` is recorded before calling postprocess.
850
+ # This sequencing may slightly reduce latency as async D2H copy does not
851
+ # need to wait for the postprocess to finish.
852
+ self.postprocess(
853
+ input_batch, sampler_output.sampled_token_ids, num_sampled_tokens
854
)
819
- return output
855
+
856
+ if self.use_async_scheduling:
857
+ return async_output
858
+ return async_output.get_output()
vllm/v1/worker/gpu/states.py
+10
-3
@@ -85,8 +85,12 @@ class RequestState:
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)
88
+
89
+ # Number of computed tokens.
90
+ self.num_computed_prefill_tokens = np.zeros(self.max_num_reqs, dtype=np.int32)
91
+ self.num_computed_tokens = torch.zeros(
92
+ self.max_num_reqs, dtype=torch.int32, device=device
93
+ )
94
95
# Last sampled tokens.
96
self.last_sampled_tokens = torch.zeros(
@@ -145,7 +149,10 @@ class RequestState:
149
)
150
self.prefill_len.np[req_idx] = prefill_len
151
self.prefill_token_ids[req_idx, :prefill_len] = prefill_token_ids
148
- self.num_tokens[req_idx] = prefill_len
152
+
153
+ self.num_computed_prefill_tokens[req_idx] = num_computed_tokens
154
+ # FIXME(woosuk): This triggers a GPU operation whenever adding a new request.
155
+ # Optimize this.
156
self.num_computed_tokens[req_idx] = num_computed_tokens
157
158
if lora_request is not None: