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, 14:45 UTC
[Model Runner V2] Change Numba AoT to JIT (#29328)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 24, 2025 at 09:34 UTC
cec418b5df3bf032a83b6a6795e8026d39e199bd
2 files changed
+32
-63
vllm/v1/worker/gpu/input_batch.py
+16
-55
@@ -4,7 +4,6 @@ from dataclasses import dataclass
4
from typing import Any
5
6
import numba
7
-import numba.types as types
7
import numpy as np
8
import torch
9
@@ -147,80 +146,42 @@ class InputBatch:
146
)
147
148
150
-# NOTE: With the type annotations, this function is pre-compiled
151
-# before the first call.
152
-@numba.jit(
153
- [
154
- types.none(
155
- types.int32[:], # idx_mapping
156
- types.int32[:], # num_scheduled_tokens
157
- types.int32[:, :], # prefill_token_ids
158
- types.int32[:], # num_computed_prefill_tokens
159
- types.int32[:], # prefill_len
160
- types.int32[:], # input_ids
161
- types.int32[:], # query_start_loc
162
- )
163
- ],
164
- nopython=True,
165
- cache=True,
166
-)
149
+@numba.njit(cache=True)
150
def _prepare_prefill_inputs(
168
- idx_mapping: np.ndarray, # batch_idx -> req_idx
169
- num_scheduled_tokens: np.ndarray, # [B]
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]
172
- prefill_len: np.ndarray, # [N]
156
input_ids: np.ndarray, # [num_input_tokens]
174
- query_start_loc: np.ndarray, # [B + 1]
157
) -> None:
176
- num_reqs = num_scheduled_tokens.shape[0]
177
- query_start_loc[0] = 0
178
-
179
- cu_num_tokens = 0
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):
181
- req_idx = idx_mapping[i]
182
- query_len = num_scheduled_tokens[i]
183
-
184
- start = num_computed_prefill_tokens[req_idx]
185
- end = min(start + query_len, prefill_len[req_idx])
186
- n = end - start
187
-
188
- start_idx = cu_num_tokens
189
- input_ids[start_idx : start_idx + n] = prefill_token_ids[req_idx, start:end]
190
-
191
- cu_num_tokens = start_idx + query_len
192
- query_start_loc[i + 1] = cu_num_tokens
193
-
194
- # Pad the inputs for CUDA graphs.
195
- # Note: pad query_start_loc to be non-decreasing, as kernels
196
- # like FlashAttention requires that
197
- query_start_loc[num_reqs + 1 :].fill(cu_num_tokens)
164
+ input_ids[query_starts[i] : query_ends[i]] = prefill_token_ids[
165
+ idx_mapping[i], starts[i] : ends[i]
166
+ ]
167
168
169
def prepare_prefill_inputs(
170
idx_mapping: np.ndarray,
171
num_scheduled_tokens: np.ndarray,
203
- total_num_tokens: int,
172
+ query_start_loc: np.ndarray,
173
prefill_token_ids: np.ndarray,
174
num_computed_prefill_tokens: np.ndarray,
206
- prefill_len: np.ndarray,
207
- input_ids: CpuGpuBuffer,
208
- query_start_loc: CpuGpuBuffer,
175
+ input_ids: np.ndarray,
176
) -> None:
177
_prepare_prefill_inputs(
178
idx_mapping,
179
num_scheduled_tokens,
180
+ query_start_loc,
181
prefill_token_ids,
182
num_computed_prefill_tokens,
215
- prefill_len,
216
- input_ids.np,
217
- query_start_loc.np,
183
+ input_ids,
184
)
219
- input_ids.copy_to_gpu(total_num_tokens)
220
- # NOTE(woosuk): We should copy the whole query_start_loc and seq_lens
221
- # tensors from CPU to GPU, because they may include paddings needed
222
- # for full CUDA graph mode.
223
- query_start_loc.copy_to_gpu()
185
186
187
@triton.jit
vllm/v1/worker/gpu/model_runner.py
+16
-8
@@ -502,20 +502,28 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
502
# Block tables: num_kv_cache_groups x [num_reqs, max_num_blocks]
503
block_tables = self.block_tables.gather_block_tables(idx_mapping)
504
505
- # Copy prefill tokens from CPU to GPU and get query_start_loc.
505
+ # Get query_start_loc.
506
+ np.cumsum(
507
+ num_scheduled_tokens,
508
+ out=self.input_buffers.query_start_loc.np[1 : num_reqs + 1],
509
+ )
510
+ # Pad for full CUDA graph mode.
511
+ # Some attention backends like FA3 require query_start_loc to be non-decreasing.
512
+ self.input_buffers.query_start_loc.np[num_reqs + 1 :] = num_tokens
513
+ self.input_buffers.query_start_loc.copy_to_gpu()
514
+ query_start_loc_gpu = self.input_buffers.query_start_loc.gpu[: num_reqs + 1]
515
+ query_start_loc_np = self.input_buffers.query_start_loc.np[: num_reqs + 1]
516
+
517
+ # Copy prefill tokens from CPU to GPU.
518
prepare_prefill_inputs(
519
idx_mapping_np,
520
num_scheduled_tokens,
509
- num_tokens,
521
+ query_start_loc_np,
522
self.req_states.prefill_token_ids,
523
self.req_states.num_computed_prefill_tokens,
512
- self.req_states.prefill_len.np,
513
- self.input_buffers.input_ids,
514
- self.input_buffers.query_start_loc,
524
+ self.input_buffers.input_ids.np,
525
)
516
- query_start_loc = self.input_buffers.query_start_loc
517
- query_start_loc_gpu = query_start_loc.gpu[: num_reqs + 1]
518
- query_start_loc_np = query_start_loc.np[: num_reqs + 1]
526
+ self.input_buffers.input_ids.copy_to_gpu(num_tokens)
527
528
# Prepare positions and seq_lens.
529
prepare_pos_seq_lens(