[Core] Warm up runner-owned Triton kernels before the first request (#49903)
Nick Hill committed
Jul 28, 2026 at 07:58 UTC
0d0504b54c73119ed643c80b4ed56ef3cf80e209
7 files changed
+183
-189
tests/v1/worker/test_kv_block_zeroer.py
+65
-1
@@ -4,7 +4,7 @@
4
import pytest
5
import torch
6
7
-from vllm.v1.worker.utils import KVBlockZeroer
7
+from vllm.v1.worker.utils import KVBlockZeroer, _zero_kv_blocks_kernel
8
9
10
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@@ -86,3 +86,67 @@ def test_non_uniform_page_sizes():
86
assert torch.all(storage[1] == 0)
87
assert torch.all(storage[2] == 0)
88
assert torch.all(storage[3] == 1)
89
+
90
+
91
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
92
+def test_warmup_compiles_every_n_blocks_specialization():
93
+ """After warmup, no launch should trigger a first-request JIT compile.
94
+
95
+ ``n_blocks`` is ``do_not_specialize``, so a single warmup launch must
96
+ cover every block count.
97
+ """
98
+ device = torch.device("cuda")
99
+ num_blocks = 64
100
+ page_size_el = 4
101
+ storage = torch.ones((num_blocks, page_size_el), dtype=torch.int32, device=device)
102
+
103
+ zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
104
+ zeroer.device = device
105
+ zeroer._meta = (
106
+ torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
107
+ torch.tensor([page_size_el], dtype=torch.int64, device=device),
108
+ 1, # max_chunks
109
+ page_size_el, # blk_size
110
+ 1, # n_segs
111
+ )
112
+
113
+ def compiled_variants() -> set:
114
+ return {
115
+ key
116
+ for caches in _zero_kv_blocks_kernel.device_caches.values()
117
+ for key in caches[0]
118
+ }
119
+
120
+ zeroer.warmup(num_blocks)
121
+ torch.accelerator.synchronize()
122
+ warmed = compiled_variants()
123
+ assert warmed
124
+
125
+ for n_blocks in (1, 2, 3, 16, 32):
126
+ zeroer.zero_block_ids(list(range(n_blocks)))
127
+ torch.accelerator.synchronize()
128
+
129
+ assert compiled_variants() == warmed
130
+
131
+
132
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
133
+def test_warmup_respects_available_block_count():
134
+ """An empty KV cache must not be warmed with out-of-range block IDs."""
135
+ device = torch.device("cuda")
136
+ page_size_el = 4
137
+ storage = torch.ones((1, page_size_el), dtype=torch.int32, device=device)
138
+
139
+ zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
140
+ zeroer.device = device
141
+ zeroer._meta = (
142
+ torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
143
+ torch.tensor([page_size_el], dtype=torch.int64, device=device),
144
+ 1,
145
+ page_size_el,
146
+ 1,
147
+ )
148
+
149
+ zeroer.warmup(0)
150
+ torch.accelerator.synchronize()
151
+
152
+ assert torch.all(storage == 1)
vllm/model_executor/warmup/kernel_warmup.py
+10
-6
@@ -98,12 +98,16 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
98
minimax_m3_msa_warmup,
99
)
100
101
- # Pooling models do not use the generation slot-mapping path.
102
- if not worker.use_v2_model_runner and not worker.model_runner.is_pooling_model:
103
- warm_v1_block_table_kernels(
104
- getattr(worker.model_runner, "device", torch.device("cuda")),
105
- worker.scheduler_config.max_num_batched_tokens,
106
- )
101
+ if not worker.use_v2_model_runner:
102
+ # Pooling models do not use the generation slot-mapping path.
103
+ if not worker.model_runner.is_pooling_model:
104
+ warm_v1_block_table_kernels(worker.model_runner)
105
+ # The KV-block zeroing kernel is driven by the scheduler's
106
+ # `new_block_ids_to_zero`, so no dummy run ever reaches it.
107
+ zeroer = getattr(worker.model_runner, "_kv_block_zeroer", None)
108
+ if zeroer is not None:
109
+ zeroer.warmup(worker.model_runner.kv_cache_config.num_blocks)
110
+
111
qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config)
112
113
# DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder
vllm/model_executor/warmup/qwen_triton_warmup.py
-114
@@ -24,24 +24,10 @@ _QWEN_MODEL_TYPES = frozenset(
24
}
25
)
26
27
-_ZERO_KV_N_BLOCKS = (1, 2)
28
-
29
-_SLOT_MAPPING_KV_BLOCK_SIZE = 16
30
-_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE = 1
31
-_SLOT_MAPPING_BLOCK_TABLE_STRIDES = (1, 3)
32
-
27
# Covers L=1 constexpr, non-divisible runtime L, and divisible runtime L.
28
_FLA_POST_CONV_WARMUP_LENGTHS = (1, 2, 16)
29
30
37
-@dataclass(frozen=True)
38
-class _ZeroKvWarmupConfig:
39
- seg_page_sizes: torch.Tensor
40
- max_chunks: int
41
- block_size: int
42
- n_segs: int
43
-
44
-
31
@dataclass(frozen=True)
32
class _QwenGDNWarmupConfig:
33
h: int
@@ -148,96 +134,6 @@ def _qwen_gdn_warmup_config(
134
return None
135
136
151
-def _get_kv_block_zeroer(runner: object) -> object | None:
152
- zeroer = getattr(runner, "kv_block_zeroer", None)
153
- if zeroer is None:
154
- zeroer = getattr(runner, "_kv_block_zeroer", None)
155
- return zeroer
156
-
157
-
158
-def _zero_kv_warmup_config(runner: object) -> _ZeroKvWarmupConfig | None:
159
- zeroer = _get_kv_block_zeroer(runner)
160
- meta = getattr(zeroer, "_meta", None)
161
- if meta is None:
162
- return None
163
-
164
- _, seg_page_sizes, max_chunks, block_size, n_segs = meta
165
- return _ZeroKvWarmupConfig(
166
- seg_page_sizes=seg_page_sizes,
167
- max_chunks=int(max_chunks),
168
- block_size=int(block_size),
169
- n_segs=int(n_segs),
170
- )
171
-
172
-
173
-def _warm_zero_kv_blocks_with_runner_zeroer(runner: object) -> bool:
174
- zeroer = _get_kv_block_zeroer(runner)
175
- zero_block_ids = getattr(zeroer, "zero_block_ids", None)
176
- if not callable(zero_block_ids):
177
- return False
178
-
179
- for n_blocks in _ZERO_KV_N_BLOCKS:
180
- zero_block_ids(list(range(n_blocks)))
181
- return True
182
-
183
-
184
-def _warm_zero_kv_blocks_kernel(
185
- device: torch.device, config: _ZeroKvWarmupConfig
186
-) -> None:
187
- from vllm.v1.worker.utils import _zero_kv_blocks_kernel
188
-
189
- max_n_blocks = max(_ZERO_KV_N_BLOCKS)
190
- max_page_size = int(config.seg_page_sizes.max().item())
191
- scratch = torch.empty(
192
- max_n_blocks * max_page_size,
193
- dtype=torch.int32,
194
- device=device,
195
- )
196
- seg_addrs = torch.tensor(
197
- [scratch.data_ptr()] * config.n_segs,
198
- dtype=torch.uint64,
199
- device=device,
200
- )
201
-
202
- for n_blocks in _ZERO_KV_N_BLOCKS:
203
- block_ids = torch.arange(n_blocks, dtype=torch.int64, device=device)
204
- grid = (n_blocks * config.n_segs * config.max_chunks,)
205
- _zero_kv_blocks_kernel[grid](
206
- seg_addrs,
207
- config.seg_page_sizes,
208
- block_ids,
209
- n_blocks,
210
- N_SEGS=config.n_segs,
211
- MAX_CHUNKS=config.max_chunks,
212
- BLOCK_SIZE=config.block_size,
213
- )
214
-
215
-
216
-def _warm_compute_slot_mapping_kernel(device: torch.device) -> None:
217
- from vllm.v1.worker.block_table import BlockTable
218
-
219
- # num_tokens/max_num_tokens are do_not_specialize; keep the launch tiny.
220
- num_tokens = 1
221
- query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
222
- positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
223
-
224
- for block_table_stride in _SLOT_MAPPING_BLOCK_TABLE_STRIDES:
225
- # Use BlockTable so the JIT key matches the production slot-mapping call.
226
- block_table = BlockTable(
227
- block_size=_SLOT_MAPPING_KV_BLOCK_SIZE,
228
- max_num_reqs=1,
229
- max_num_blocks_per_req=block_table_stride,
230
- max_num_batched_tokens=num_tokens,
231
- pin_memory=False,
232
- device=device,
233
- kernel_block_size=_SLOT_MAPPING_KV_BLOCK_SIZE,
234
- cp_kv_cache_interleave_size=_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE,
235
- )
236
- block_table.add_row(list(range(block_table_stride)), 0)
237
- block_table.commit_block_table(num_reqs=1)
238
- block_table.compute_slot_mapping(1, query_start_loc, positions)
239
-
240
-
137
def _warm_causal_conv1d_fwd_kernel(
138
device: torch.device, config: _QwenGDNWarmupConfig
139
) -> None:
@@ -373,16 +269,6 @@ def qwen_triton_warmup(
269
device = getattr(runner, "device", torch.device("cuda"))
270
logger.info("Warming up Qwen Triton kernels for model_type=%s.", model_type)
271
376
- zero_config = _zero_kv_warmup_config(runner)
377
- warmed_zeroer = _warm_zero_kv_blocks_with_runner_zeroer(runner)
378
- if zero_config is not None:
379
- _warm_zero_kv_blocks_kernel(device, zero_config)
380
- elif not warmed_zeroer:
381
- logger.info("Skipping Qwen zero-kv warmup: no KVBlockZeroer metadata.")
382
-
383
- _warm_compute_slot_mapping_kernel(device)
384
- _synchronize_device(device)
385
-
272
compilation_config = getattr(runner, "compilation_config", None)
273
static_forward_context = getattr(compilation_config, "static_forward_context", None)
274
gdn_config = _qwen_gdn_warmup_config(static_forward_context)
vllm/model_executor/warmup/v1_block_table_warmup.py
+14
-28
@@ -2,42 +2,28 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Warm up v1 block-table Triton kernels."""
4
5
+from typing import TYPE_CHECKING
6
+
7
import torch
8
9
+if TYPE_CHECKING:
10
+ from vllm.v1.worker.gpu_model_runner import GPUModelRunner
11
+
12
_SLOT_MAPPING_WARMUP_TOKENS = 8
8
-_SLOT_MAPPING_WARMUP_BLOCK_SIZES = (3, 16)
9
-_SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE = 1
13
14
12
-def warm_v1_block_table_kernels(
13
- device: torch.device,
14
- max_tokens: int,
15
-) -> None:
16
- from vllm.v1.worker.block_table import BlockTable
15
+def warm_v1_block_table_kernels(runner: "GPUModelRunner") -> None:
16
+ """JIT-compile ``_compute_slot_mapping_kernel`` for the real block tables."""
17
18
- num_tokens = max(0, min(_SLOT_MAPPING_WARMUP_TOKENS, max_tokens))
18
+ device = runner.device
19
+ block_table = runner.input_batch.block_table
20
+ num_tokens = min(
21
+ _SLOT_MAPPING_WARMUP_TOKENS,
22
+ runner.scheduler_config.max_num_batched_tokens,
23
+ )
24
if num_tokens <= 0:
25
return
26
27
query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
28
positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
24
- for block_size in _SLOT_MAPPING_WARMUP_BLOCK_SIZES:
25
- max_num_blocks_per_req = max(
26
- 1, (max(num_tokens, max_tokens) + block_size - 1) // block_size
27
- )
28
- max_num_blocks_per_req = ((max_num_blocks_per_req + 15) // 16) * 16
29
- block_table = BlockTable(
30
- block_size=block_size,
31
- max_num_reqs=1,
32
- max_num_blocks_per_req=max_num_blocks_per_req,
33
- max_num_batched_tokens=max(num_tokens, max_tokens),
34
- pin_memory=False,
35
- device=device,
36
- kernel_block_size=block_size,
37
- cp_kv_cache_interleave_size=(
38
- _SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE
39
- ),
40
- )
41
- block_table.add_row(list(range(max_num_blocks_per_req)), 0)
42
- block_table.commit_block_table(1)
43
- block_table.compute_slot_mapping(1, query_start_loc, positions)
29
+ block_table.compute_slot_mapping(1, query_start_loc, positions)
vllm/v1/worker/gpu/warmup.py
+85
-36
@@ -157,13 +157,10 @@ def warmup_kernels(
157
worker_execute_model: Callable[[SchedulerOutput], Any],
158
worker_sample_tokens: Callable[[GrammarOutput | None], Any],
159
) -> None:
160
- """Run two execute_model + sample_tokens iterations to JIT compile
161
- triton kernels. We must call the provided worker's execute_model for
162
- pipeline parallel coordination.
160
+ """Run scheduler-realistic prefill and decode steps to JIT compile kernels.
161
164
- The first iteration simulates a prefill with requests of
165
- decode_query_len + 1 prompt tokens each. The second iteration simulates
166
- a decode step with all requests generating decode_query_len tokens.
162
+ We must call the provided worker's execute_model for pipeline parallel
163
+ coordination.
164
"""
165
num_spec_steps = model_runner.num_speculative_steps
166
decode_query_len = model_runner.decode_query_len
@@ -172,8 +169,13 @@ def warmup_kernels(
169
# a uniform decode batch.
170
prompt_len = decode_query_len + 1
171
prompt_token_ids = list(range(prompt_len))
175
- # After prefill, decode generates decode_query_len tokens.
176
- decode_len = prompt_len + decode_query_len
172
+ # Upper bound on the decode steps built in `decode_steps` below.
173
+ num_decode_steps = 1
174
+ if not model_runner.is_pooling_model:
175
+ num_decode_steps = 5 if num_spec_steps > 0 else 3
176
+ # Size the block allocation for the worst case: every request advancing
177
+ # decode_query_len tokens on every decode step.
178
+ decode_len = prompt_len + num_decode_steps * decode_query_len
179
180
kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
181
num_kv_cache_groups = len(kv_cache_groups)
@@ -208,9 +210,6 @@ def warmup_kernels(
210
kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups]
211
prefill_block_counts = [_warmup_block_count(prompt_len, s) for s in kv_cache_specs]
212
decode_block_counts = [_warmup_block_count(decode_len, s) for s in kv_cache_specs]
211
- decode_block_deltas = [
212
- d - p for d, p in zip(decode_block_counts, prefill_block_counts)
213
- ]
213
max_blocks_per_req = sum(decode_block_counts)
214
215
num_reqs = min(
@@ -243,6 +242,11 @@ def warmup_kernels(
242
nonlocal next_block_id
243
return list(range(next_block_id, next_block_id := next_block_id + num_blocks))
244
245
+ # The KV-block zeroing kernel is driven by the scheduler's
246
+ # new_block_ids_to_zero, so none of the steps below reach it.
247
+ if model_runner.kv_block_zeroer is not None:
248
+ model_runner.kv_block_zeroer.warmup(model_runner.kv_cache_config.num_blocks)
249
+
250
# Step 1: Prefill all requests with 1 + decode_query_len prompt tokens each.
251
new_reqs = [
252
NewRequestData.from_request(
@@ -287,33 +291,78 @@ def warmup_kernels(
291
292
worker_sample_tokens(grammar_output)
293
290
- # Step 2: Decode all requests with decode_query_len tokens each.
291
- cached_req_data = CachedRequestData.make_empty()
292
- cached_req_data.req_ids = list(req_ids)
293
- cached_req_data.num_computed_tokens = [prompt_len] * num_reqs
294
- cached_req_data.num_output_tokens = [1] * num_reqs
295
- new_block = any(decode_block_deltas)
296
- cached_req_data.new_block_ids = [
297
- tuple(_alloc_blocks(n) for n in decode_block_deltas) if new_block else None
298
- for _ in range(num_reqs)
299
- ]
294
+ # Per-request state carried across the decode steps.
295
+ req_computed = [prompt_len] * num_reqs
296
+ req_blocks = [list(prefill_block_counts) for _ in range(num_reqs)]
297
+
298
+ def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None:
299
+ """Decode `indices`, spec-decoding the ones flagged in `spec_flags`."""
300
+ cached_req_data = CachedRequestData.make_empty()
301
+ cached_req_data.req_ids = [req_ids[i] for i in indices]
302
+ cached_req_data.num_computed_tokens = [req_computed[i] for i in indices]
303
+ cached_req_data.num_output_tokens = [1] * len(indices)
304
+ cached_req_data.new_block_ids = []
305
+
306
+ step_num_scheduled_tokens: dict[str, int] = {}
307
+ step_spec_tokens: dict[str, list[int]] = {}
308
+ for i, use_spec in zip(indices, spec_flags):
309
+ num_tokens = decode_query_len if use_spec else 1
310
+ after = req_computed[i] + num_tokens
311
+ deltas = [
312
+ _warmup_block_count(after, spec) - held
313
+ for spec, held in zip(kv_cache_specs, req_blocks[i])
314
+ ]
315
+ cached_req_data.new_block_ids.append(
316
+ tuple(_alloc_blocks(n) for n in deltas) if any(deltas) else None
317
+ )
318
+ req_blocks[i] = [
319
+ held + delta for held, delta in zip(req_blocks[i], deltas)
320
+ ]
321
+ step_num_scheduled_tokens[req_ids[i]] = num_tokens
322
+ if use_spec:
323
+ step_spec_tokens[req_ids[i]] = [0] * num_spec_steps
324
+
325
+ decode_output = SchedulerOutput.make_empty()
326
+ decode_output.scheduled_cached_reqs = cached_req_data
327
+ decode_output.num_scheduled_tokens = step_num_scheduled_tokens
328
+ decode_output.scheduled_spec_decode_tokens = step_spec_tokens
329
+ decode_output.total_num_scheduled_tokens = sum(
330
+ step_num_scheduled_tokens.values()
331
+ )
332
+ decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
333
301
- decode_output = SchedulerOutput.make_empty()
302
- decode_output.scheduled_cached_reqs = cached_req_data
303
- decode_output.num_scheduled_tokens = {
304
- req_id: decode_query_len for req_id in req_ids
305
- }
306
- if num_spec_steps > 0:
307
- decode_output.scheduled_spec_decode_tokens = {
308
- req_id: [0] * num_spec_steps for req_id in req_ids
309
- }
310
- decode_output.total_num_scheduled_tokens = sum(
311
- decode_output.num_scheduled_tokens.values()
312
- )
313
- decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
334
+ worker_execute_model(decode_output)
335
+ worker_sample_tokens(None)
336
315
- worker_execute_model(decode_output)
316
- worker_sample_tokens(None)
337
+ for i, use_spec in zip(indices, spec_flags):
338
+ req_computed[i] += decode_query_len if use_spec else 1
339
+
340
+ all_indices = list(range(num_reqs))
341
+ use_spec_decode = num_spec_steps > 0
342
+
343
+ # Decode steps to warm, as (request indices, per-request spec flag).
344
+ # Under spec decoding the scheduler drops requests the drafter proposed
345
+ # nothing for, so warm each batch shape with and without draft tokens.
346
+ decode_steps: list[tuple[list[int], list[bool]]] = [
347
+ (all_indices, [use_spec_decode] * num_reqs),
348
+ ]
349
+ if num_reqs >= 2:
350
+ # Mixed spec / non-spec: GDN and KDA reclassify the non-spec decode
351
+ # as a prefill and split the batch into spec/non-spec token indices.
352
+ decode_steps.append(([0, 1], [use_spec_decode, False]))
353
+ if use_spec_decode:
354
+ # Exercise the model paths that split a batch by whether each
355
+ # request received draft tokens.
356
+ decode_steps.append(([0, 1], [False, False]))
357
+ if num_reqs > 1:
358
+ decode_steps.append(([0], [use_spec_decode]))
359
+ if use_spec_decode:
360
+ decode_steps.append(([0], [False]))
361
+ elif use_spec_decode:
362
+ decode_steps.append(([0], [False]))
363
+
364
+ for step_indices, step_spec_flags in decode_steps:
365
+ _run_decode_step(step_indices, step_spec_flags)
366
367
# Clean up - process finish_req_ids.
368
cleanup_output = SchedulerOutput.make_empty()
vllm/v1/worker/mamba_utils.py
+3
-3
@@ -150,7 +150,7 @@ def _copy_mamba_state_block(
150
tl.store(tail_dst + tail_off, tail_data, mask=tail_mask)
151
152
153
-@triton.jit
153
+@triton.jit(do_not_specialize=["num_reqs"])
154
def postprocess_mamba_fused_kernel(
155
# Decision inputs (per-request)
156
num_accepted_tokens_ptr,
@@ -280,7 +280,7 @@ def postprocess_mamba_fused_kernel(
280
)
281
282
283
-@triton.jit
283
+@triton.jit(do_not_specialize=["num_reqs"])
284
def preprocess_mamba_align_fused_kernel(
285
idx_mapping_ptr,
286
state_idx_ptr,
@@ -326,7 +326,7 @@ def preprocess_mamba_align_fused_kernel(
326
tl.store(num_accepted_tokens_ptr + req_indices, 1, mask=mask & should_reset)
327
328
329
-@triton.jit
329
+@triton.jit(do_not_specialize=["num_reqs"])
330
def precopy_mamba_align_fused_kernel(
331
# Per-request-slot inputs (indexed by req_idx via idx_mapping), produced by
332
# the V2 fused align preprocess kernel for the current step:
vllm/v1/worker/utils.py
+6
-1
@@ -40,7 +40,7 @@ from vllm.v1.kv_cache_interface import (
40
logger = init_logger(__name__)
41
42
43
-@triton.jit
43
+@triton.jit(do_not_specialize=["n_blocks"])
44
def _zero_kv_blocks_kernel(
45
seg_addrs_ptr,
46
seg_page_sizes_ptr,
@@ -206,6 +206,11 @@ class KVBlockZeroer:
206
BLOCK_SIZE=blk_size,
207
)
208
209
+ def warmup(self, num_kv_blocks: int) -> None:
210
+ """JIT-compile the zeroing kernel before the first real request."""
211
+ if num_kv_blocks > 0:
212
+ self.zero_block_ids([0])
213
+
214
215
@dataclass
216
class AttentionGroup: