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, 08:30 UTC
[Model Runner V2] Refactor prefill token preparation (#29712)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 28, 2025 at 19:49 UTC
ca1b1e7296e44c75e023b3baaec55d9431a33d78
5 files changed
+83
-78
vllm/v1/worker/gpu/cudagraph_utils.py
+1
-1
@@ -78,7 +78,7 @@ class CudaGraphManager:
78
kv_cache_config: KVCacheConfig,
79
) -> None:
80
num_reqs = min(num_tokens, self.max_num_reqs)
81
- input_ids = input_buffers.input_ids.gpu[:num_tokens]
81
+ input_ids = input_buffers.input_ids[:num_tokens]
82
positions = input_buffers.positions[:num_tokens]
83
attn_metadata = prepare_inputs_to_capture(
84
num_reqs,
vllm/v1/worker/gpu/input_batch.py
+52
-34
@@ -3,7 +3,6 @@
3
from dataclasses import dataclass
4
from typing import Any
5
6
-import numba
6
import numpy as np
7
import torch
8
@@ -30,15 +29,12 @@ class InputBuffers:
29
self.pin_memory = pin_memory
30
31
self.idx_mapping = self._make_buffer(max_num_reqs, dtype=torch.int32)
33
- self.input_ids = self._make_buffer(max_num_tokens, dtype=torch.int32)
32
+ self.input_ids = torch.zeros(max_num_tokens, dtype=torch.int32, device=device)
33
self.positions = torch.zeros(max_num_tokens, dtype=torch.int64, device=device)
34
self.query_start_loc = self._make_buffer(max_num_reqs + 1, dtype=torch.int32)
35
self.seq_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
36
self.cu_num_logits = self._make_buffer(max_num_reqs + 1, dtype=torch.int32)
37
39
- # Spec decoding.
40
- self.next_prefill_tokens = self._make_buffer(max_num_reqs, dtype=torch.int32)
41
-
38
# Structured outputs.
39
self.bitmask_indices = self._make_buffer(max_num_reqs, dtype=torch.int32)
40
self.grammar_bitmask = self._make_buffer(
@@ -120,7 +116,7 @@ class InputBatch:
116
input_buffers.seq_lens[num_reqs:] = 0
117
seq_lens = input_buffers.seq_lens[:num_reqs]
118
123
- input_ids = input_buffers.input_ids.copy_to_gpu(num_tokens)
119
+ input_ids = input_buffers.input_ids[:num_tokens]
120
positions = input_buffers.positions[:num_tokens]
121
# attn_metadata = defaultdict(lambda: None)
122
logits_indices = query_start_loc[1:] - 1
@@ -146,41 +142,63 @@ class InputBatch:
142
)
143
144
149
-@numba.njit(cache=True)
150
-def _prepare_prefill_inputs(
151
- idx_mapping: np.ndarray, # [B]
152
- query_lens: np.ndarray, # [B]
153
- query_start_loc: np.ndarray, # [B + 1]
154
- prefill_token_ids: np.ndarray, # [N, max_model_len]
155
- num_computed_prefill_tokens: np.ndarray, # [N]
156
- input_ids: np.ndarray, # [num_input_tokens]
157
-) -> None:
158
- num_reqs = idx_mapping.shape[0]
159
- query_starts = query_start_loc[:num_reqs]
160
- query_ends = query_start_loc[1 : num_reqs + 1]
161
- starts = num_computed_prefill_tokens[idx_mapping]
162
- ends = starts + query_lens
163
- for i in range(num_reqs):
164
- input_ids[query_starts[i] : query_ends[i]] = prefill_token_ids[
165
- idx_mapping[i], starts[i] : ends[i]
166
- ]
145
+@triton.jit
146
+def _prepare_prefill_inputs_kernel(
147
+ input_ids_ptr,
148
+ next_prefill_tokens_ptr,
149
+ idx_mapping_ptr,
150
+ query_start_loc_ptr,
151
+ prefill_token_ids_ptr,
152
+ prefill_token_ids_stride,
153
+ prefill_lens_ptr,
154
+ num_computed_tokens_ptr,
155
+ BLOCK_SIZE: tl.constexpr,
156
+):
157
+ batch_idx = tl.program_id(0)
158
+ req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
159
+ prefill_len = tl.load(prefill_lens_ptr + req_state_idx)
160
+ num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
161
+ if num_computed >= prefill_len:
162
+ # Not prefill.
163
+ return
164
+
165
+ query_start = tl.load(query_start_loc_ptr + batch_idx)
166
+ query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
167
+ query_len = query_end - query_start
168
+
169
+ prefill_ptr = prefill_token_ids_ptr + req_state_idx * prefill_token_ids_stride
170
+ for i in range(0, query_len, BLOCK_SIZE):
171
+ block = i + tl.arange(0, BLOCK_SIZE)
172
+ mask = block < query_len
173
+ tokens = tl.load(prefill_ptr + num_computed + block, mask=mask)
174
+ tl.store(input_ids_ptr + query_start + block, tokens, mask=mask)
175
+
176
+ next_pos = num_computed + query_len
177
+ if next_pos < prefill_len:
178
+ next_token = tl.load(prefill_ptr + next_pos)
179
+ tl.store(next_prefill_tokens_ptr + req_state_idx, next_token)
180
181
182
def prepare_prefill_inputs(
170
- idx_mapping: np.ndarray,
171
- num_scheduled_tokens: np.ndarray,
172
- query_start_loc: np.ndarray,
173
- prefill_token_ids: np.ndarray,
174
- num_computed_prefill_tokens: np.ndarray,
175
- input_ids: np.ndarray,
183
+ input_ids: torch.Tensor,
184
+ next_prefill_tokens: torch.Tensor,
185
+ idx_mapping: torch.Tensor,
186
+ query_start_loc: torch.Tensor,
187
+ prefill_token_ids: torch.Tensor,
188
+ prefill_len: torch.Tensor,
189
+ num_computed_tokens: torch.Tensor,
190
) -> None:
177
- _prepare_prefill_inputs(
191
+ num_reqs = idx_mapping.shape[0]
192
+ _prepare_prefill_inputs_kernel[(num_reqs,)](
193
+ input_ids,
194
+ next_prefill_tokens,
195
idx_mapping,
179
- num_scheduled_tokens,
196
query_start_loc,
197
prefill_token_ids,
182
- num_computed_prefill_tokens,
183
- input_ids,
198
+ prefill_token_ids.stride(0),
199
+ prefill_len,
200
+ num_computed_tokens,
201
+ BLOCK_SIZE=1024,
202
)
203
204
vllm/v1/worker/gpu/model_runner.py
+17
-29
@@ -104,11 +104,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
104
if self.use_async_scheduling:
105
self.input_prep_event = torch.cuda.Event()
106
self.structured_outputs_event = torch.cuda.Event()
107
- self.spec_decode_event = torch.cuda.Event()
107
else:
108
self.input_prep_event = None
109
self.structured_outputs_event = None
111
- self.spec_decode_event = None
110
111
if self.speculative_config is not None:
112
self.do_spec_decode = True
@@ -412,9 +410,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
410
cu_num_new_blocks[i].append(x + len(block_ids))
411
new_block_ids[i].extend(block_ids)
412
overwrite.append(True)
415
- # Update the GPU tensors for request states.
416
- if scheduler_output.scheduled_new_reqs:
417
- self.req_states.prefill_len.copy_to_gpu()
413
414
# Add new blocks for the existing requests.
415
cached_reqs = scheduler_output.scheduled_cached_reqs
@@ -507,16 +502,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
502
query_start_loc_cpu = self.input_buffers.query_start_loc.cpu[: num_reqs + 1]
503
query_start_loc_np = self.input_buffers.query_start_loc.np[: num_reqs + 1]
504
510
- # Copy prefill tokens from CPU to GPU.
505
+ # Get prefill tokens.
506
prepare_prefill_inputs(
512
- idx_mapping_np,
513
- num_scheduled_tokens,
514
- query_start_loc_np,
515
- self.req_states.prefill_token_ids.np,
516
- self.req_states.num_computed_prefill_tokens,
517
- self.input_buffers.input_ids.np,
507
+ self.input_buffers.input_ids,
508
+ self.req_states.next_prefill_tokens,
509
+ idx_mapping,
510
+ query_start_loc_gpu,
511
+ self.req_states.prefill_token_ids.gpu,
512
+ self.req_states.prefill_len.gpu,
513
+ self.req_states.num_computed_tokens,
514
)
519
- self.input_buffers.input_ids.copy_to_gpu(num_tokens)
515
516
# Prepare positions and seq_lens.
517
prepare_pos_seq_lens(
@@ -531,7 +526,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
526
# Some input token ids are directly read from the last sampled tokens
527
# and draft tokens. Also, get the logits indices to sample tokens from.
528
logits_indices = combine_sampled_and_draft_tokens(
534
- self.input_buffers.input_ids.gpu,
529
+ self.input_buffers.input_ids,
530
idx_mapping,
531
self.req_states.last_sampled_tokens,
532
query_start_loc_gpu,
@@ -572,7 +567,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
567
kv_cache_config=self.kv_cache_config,
568
)
569
575
- input_ids = self.input_buffers.input_ids.gpu[:num_tokens_after_padding]
570
+ input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]
571
positions = self.input_buffers.positions[:num_tokens_after_padding]
572
return InputBatch(
573
req_ids=req_ids,
@@ -782,20 +777,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
777
num_sampled: torch.Tensor,
778
num_rejected: torch.Tensor,
779
) -> torch.Tensor:
785
- num_reqs = input_batch.num_reqs
786
- idx_mapping_np = input_batch.idx_mapping_np
787
- with async_barrier(self.spec_decode_event):
788
- self.input_buffers.next_prefill_tokens.np[:num_reqs] = (
789
- self.req_states.prefill_token_ids.np[
790
- idx_mapping_np,
791
- self.req_states.num_computed_prefill_tokens[idx_mapping_np],
792
- ]
793
- )
794
- next_prefill_tokens = self.input_buffers.next_prefill_tokens.copy_to_gpu(
795
- num_reqs
796
- )
797
-
780
assert self.speculator is not None
781
+ last_sampled_tokens = self.req_states.last_sampled_tokens[
782
+ input_batch.idx_mapping
783
+ ]
784
+ next_prefill_tokens = self.req_states.next_prefill_tokens[
785
+ input_batch.idx_mapping
786
+ ]
787
draft_tokens = self.speculator.propose(
788
input_batch,
789
sampling_metadata,
@@ -803,7 +791,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
791
aux_hidden_states,
792
num_sampled,
793
num_rejected,
806
- self.req_states.last_sampled_tokens,
794
+ last_sampled_tokens,
795
next_prefill_tokens,
796
)
797
return draft_tokens
vllm/v1/worker/gpu/spec_decode/eagle.py
+8
-11
@@ -121,7 +121,7 @@ class EagleSpeculator:
121
num_tokens_across_dp=num_tokens_across_dp,
122
):
123
ret_hidden_states = self.model(
124
- input_ids=self.input_buffers.input_ids.gpu[:num_tokens],
124
+ input_ids=self.input_buffers.input_ids[:num_tokens],
125
positions=self.input_buffers.positions[:num_tokens],
126
hidden_states=self.hidden_states[:num_tokens],
127
)
@@ -194,7 +194,7 @@ class EagleSpeculator:
194
num_sampled: torch.Tensor,
195
# [num_reqs]
196
num_rejected: torch.Tensor,
197
- # [max_num_reqs, 1]
197
+ # [num_reqs]
198
last_sampled: torch.Tensor,
199
# [num_reqs]
200
next_prefill_tokens: torch.Tensor,
@@ -316,7 +316,6 @@ def _prepare_eagle_inputs_kernel(
316
eagle_positions_ptr,
317
target_input_ids_ptr,
318
target_positions_ptr,
319
- idx_mapping_ptr,
319
last_sampled_ptr,
320
next_prefill_tokens_ptr,
321
num_sampled_ptr,
@@ -335,8 +334,7 @@ def _prepare_eagle_inputs_kernel(
334
335
num_sampled = tl.load(num_sampled_ptr + batch_idx)
336
if num_sampled > 0:
338
- req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
339
- next_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32)
337
+ next_token = tl.load(last_sampled_ptr + batch_idx).to(tl.int32)
338
else:
339
# Chunked prefilling.
340
# Get the next prefill token.
@@ -368,9 +366,9 @@ def prepare_eagle_inputs(
366
num_sampled: torch.Tensor,
367
# [num_reqs]
368
num_rejected: torch.Tensor,
371
- # [max_num_reqs, 1]
369
+ # [num_reqs]
370
last_sampled: torch.Tensor,
373
- # [max_num_reqs]
371
+ # [num_reqs]
372
next_prefill_tokens: torch.Tensor,
373
) -> torch.Tensor:
374
num_reqs = input_batch.num_reqs
@@ -381,11 +379,10 @@ def prepare_eagle_inputs(
379
)
380
_prepare_eagle_inputs_kernel[(num_reqs,)](
381
last_token_indices,
384
- input_buffers.input_ids.gpu,
382
+ input_buffers.input_ids,
383
input_buffers.positions,
384
input_batch.input_ids,
385
input_batch.positions,
388
- input_batch.idx_mapping,
386
last_sampled,
387
next_prefill_tokens,
388
num_sampled,
@@ -485,7 +482,7 @@ def prepare_eagle_decode(
482
last_token_indices,
483
target_seq_lens,
484
num_rejected,
488
- input_buffers.input_ids.gpu,
485
+ input_buffers.input_ids,
486
input_buffers.positions,
487
input_hidden_states,
488
input_hidden_states.stride(0),
@@ -553,7 +550,7 @@ def update_eagle_inputs(
550
):
551
num_reqs, hidden_size = output_hidden_states.shape
552
_update_eagle_inputs_kernel[(num_reqs,)](
556
- input_buffers.input_ids.gpu,
553
+ input_buffers.input_ids,
554
input_buffers.positions,
555
hidden_states,
556
hidden_states.stride(0),
vllm/v1/worker/gpu/states.py
+5
-3
@@ -117,8 +117,7 @@ class RequestState:
117
self.prefill_token_ids = UvaBuffer(
118
self.max_num_reqs, self.max_model_len, dtype=torch.int32
119
)
120
- self.prefill_len = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
121
-
120
+ self.prefill_len = UvaBuffer(self.max_num_reqs, dtype=torch.int32)
121
# Number of computed tokens.
122
self.num_computed_prefill_tokens = np.zeros(self.max_num_reqs, dtype=np.int32)
123
self.num_computed_tokens = torch.zeros(
@@ -140,6 +139,9 @@ class RequestState:
139
dtype=torch.int64,
140
device=device,
141
)
142
+ self.next_prefill_tokens = torch.zeros(
143
+ self.max_num_reqs, dtype=torch.int32, device=device
144
+ )
145
146
# LoRA.
147
self.lora_ids = np.zeros(self.max_num_reqs, dtype=np.int32)
@@ -380,13 +382,13 @@ def _expand_sampling_metadata_kernel(
382
expanded_top_p_ptr,
383
top_k_ptr,
384
expanded_top_k_ptr,
383
- seeds_ptr,
385
rep_penalty_ptr,
386
expanded_rep_penalty_ptr,
387
freq_penalty_ptr,
388
expanded_freq_penalty_ptr,
389
pres_penalty_ptr,
390
expanded_pres_penalty_ptr,
391
+ seeds_ptr,
392
expanded_seeds_ptr,
393
cu_num_logits_ptr,
394
BLOCK_SIZE: tl.constexpr,