[Perf][1/N] Expand Triton kernel warmup coverage, DSv4 (#46634)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Roberto L. Castro committed
Jun 29, 2026 at 18:40 UTC
030c9523bdb6a6292545768c863fd747c195b06b
3 files changed
+386
vllm/model_executor/warmup/kernel_warmup.py
+13
@@ -25,6 +25,12 @@ from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
25
flashinfer_sparse_mla_decode_autotune_warmup,
26
)
27
from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
28
+from vllm.model_executor.warmup.sparse_mla_triton_warmup import (
29
+ sparse_mla_triton_warmup_if_needed,
30
+)
31
+from vllm.model_executor.warmup.v1_block_table_warmup import (
32
+ warm_v1_block_table_kernels,
33
+)
34
from vllm.platforms import current_platform
35
from vllm.utils.deep_gemm import is_deep_gemm_supported
36
from vllm.utils.flashinfer import has_flashinfer
@@ -41,6 +47,12 @@ def kernel_warmup(worker: "Worker"):
47
minimax_m3_msa_warmup,
48
)
49
50
+ # Pooling models do not use the generation slot-mapping path.
51
+ if not worker.use_v2_model_runner and not worker.model_runner.is_pooling_model:
52
+ warm_v1_block_table_kernels(
53
+ getattr(worker.model_runner, "device", torch.device("cuda")),
54
+ worker.scheduler_config.max_num_batched_tokens,
55
+ )
56
qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config)
57
58
# DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder
@@ -55,6 +67,7 @@ def kernel_warmup(worker: "Worker"):
67
)
68
69
# Run next so input-prep kernels JIT against pristine runner state.
70
+ sparse_mla_triton_warmup_if_needed(worker)
71
flashinfer_sparse_mla_decode_autotune_warmup(worker)
72
deepseek_v4_sparse_mla_attention_warmup(worker)
73
vllm/model_executor/warmup/sparse_mla_triton_warmup.py
new
+330
@@ -0,0 +1,330 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Warm up sparse-MLA Triton metadata kernels."""
4
+
5
+from typing import TYPE_CHECKING
6
+
7
+import torch
8
+
9
+from vllm.logger import init_logger
10
+
11
+if TYPE_CHECKING:
12
+ from vllm.v1.worker.gpu_model_runner import GPUModelRunner
13
+ from vllm.v1.worker.gpu_worker import Worker
14
+
15
+logger = init_logger(__name__)
16
+
17
+_DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset(
18
+ {
19
+ "FLASHMLA_SPARSE_DSV4",
20
+ "FLASHINFER_MLA_SPARSE_DSV4",
21
+ "ROCM_FLASHMLA_SPARSE_DSV4",
22
+ "DEEPSEEK_SPARSE_SWA",
23
+ }
24
+)
25
+_GENERIC_SPARSE_MLA_BACKENDS = frozenset(
26
+ {
27
+ "FLASHMLA_SPARSE",
28
+ "FLASHINFER_MLA_SPARSE",
29
+ "FLASHINFER_MLA_SPARSE_SM120",
30
+ }
31
+)
32
+
33
+_SPARSE_PREFILL_METADATA_NUM_PREFILLS = (1, 2, 4, 8)
34
+_SPARSE_PREFILL_METADATA_NUM_DECODES = (0, 1, 2)
35
+_DSV4_PREFILL_CHUNK_METADATA_COMPRESS_RATIOS = (4, 128)
36
+_PREFILL_CHUNK_METADATA_SEQ_LEN_MULTIPLIERS = (2, 3)
37
+_PREFILL_CHUNK_METADATA_QUERY_SLICE_OFFSETS = (
38
+ # query_slice_start offset, query_slice_stop offset
39
+ (0, 0),
40
+ (0, -1),
41
+ (1, 0),
42
+ (1, -1),
43
+)
44
+_COMBINE_TOPK_SWA_INPUT_VARIANTS = (
45
+ # offset_topk, offset_query_and_seq, offset_gather
46
+ (False, False, False),
47
+ (False, True, False),
48
+ (True, True, True),
49
+)
50
+_DSV4_COMBINE_TOPK_SWA_WARMUP_CASES = (
51
+ # compress_ratio, topk, topk_width, N
52
+ (1, 0, 512, 512),
53
+ (4, 512, 512, 512 * 4),
54
+ # DSv4-Pro C4A traffic uses top-k 1024 with N=1024.
55
+ (4, 1024, 1024, 1024),
56
+ (128, 8192, 8192, 8192 * 128),
57
+ # Real C128A traffic also specializes N=1 in one call path.
58
+ (128, 8192, 8192, 1),
59
+)
60
+
61
+
62
+def _clamp_warmup_tokens(num_tokens: int, max_tokens: int) -> int:
63
+ return max(0, min(num_tokens, max_tokens))
64
+
65
+
66
+def _next_power_of_2(x: int) -> int:
67
+ return 1 << (x - 1).bit_length()
68
+
69
+
70
+def _hf_config_int(runner: "GPUModelRunner", name: str, default: int) -> int:
71
+ model_config = getattr(runner.vllm_config, "model_config", None)
72
+ hf_config = getattr(model_config, "hf_config", None)
73
+ return int(getattr(hf_config, name, default) or default)
74
+
75
+
76
+def _attention_backend_name(backend: object) -> str | None:
77
+ get_name = getattr(backend, "get_name", None)
78
+ if get_name is None:
79
+ return None
80
+ try:
81
+ return get_name()
82
+ except NotImplementedError:
83
+ return None
84
+
85
+
86
+def _has_attention_backend(
87
+ runner: "GPUModelRunner",
88
+ backend_names: frozenset[str],
89
+) -> bool:
90
+ for groups in getattr(runner, "attn_groups", []) or ():
91
+ for group in groups:
92
+ name = _attention_backend_name(getattr(group, "backend", None))
93
+ if name in backend_names:
94
+ return True
95
+ return False
96
+
97
+
98
+def _warm_sparse_swa_prefill_metadata_kernel(
99
+ device: torch.device,
100
+ window_size: int,
101
+ prefill_tokens: int,
102
+) -> None:
103
+ from vllm.v1.attention.backends.mla.sparse_swa import (
104
+ _compute_prefill_metadata_kernel,
105
+ )
106
+
107
+ for num_prefills in _SPARSE_PREFILL_METADATA_NUM_PREFILLS:
108
+ for num_decodes in _SPARSE_PREFILL_METADATA_NUM_DECODES:
109
+ query_lens = [1] * num_decodes
110
+ query_lens += [prefill_tokens] * num_prefills
111
+ query_start_locs = [0]
112
+ for query_len in query_lens:
113
+ query_start_locs.append(query_start_locs[-1] + query_len)
114
+ query_start_loc = torch.tensor(
115
+ query_start_locs,
116
+ dtype=torch.int32,
117
+ device=device,
118
+ )
119
+ seq_lens = torch.tensor(
120
+ [1] * num_decodes + [window_size + q for q in query_lens[num_decodes:]],
121
+ dtype=torch.int32,
122
+ device=device,
123
+ )
124
+ prefill_gather_lens = torch.empty(
125
+ num_prefills, dtype=torch.int32, device=device
126
+ )
127
+ _compute_prefill_metadata_kernel[(1,)](
128
+ prefill_gather_lens,
129
+ seq_lens,
130
+ query_start_loc,
131
+ num_prefills,
132
+ num_decodes,
133
+ window_size,
134
+ BLOCK_SIZE=_next_power_of_2(num_prefills),
135
+ )
136
+
137
+
138
+def _warm_prefill_chunk_metadata_kernel(
139
+ device: torch.device,
140
+ compress_ratio: int,
141
+ query_len: int,
142
+) -> None:
143
+ from vllm.v1.attention.backends.mla.indexer import build_prefill_chunk_metadata
144
+
145
+ num_reqs = 2
146
+ query_start_loc_cpu = torch.arange(
147
+ 0, (num_reqs + 1) * query_len, query_len, dtype=torch.int32
148
+ )
149
+ query_start_loc = query_start_loc_cpu.to(device=device)
150
+
151
+ uncompressed_seq_lens_cpu = torch.tensor(
152
+ [
153
+ compress_ratio * multiplier + query_len
154
+ for multiplier in _PREFILL_CHUNK_METADATA_SEQ_LEN_MULTIPLIERS
155
+ ],
156
+ dtype=torch.int32,
157
+ )
158
+ compressed_seq_lens_cpu = uncompressed_seq_lens_cpu // compress_ratio
159
+ uncompressed_seq_lens = uncompressed_seq_lens_cpu.to(device=device)
160
+ compressed_seq_lens = compressed_seq_lens_cpu.to(device=device)
161
+ block_table = torch.zeros(
162
+ (num_reqs, int(compressed_seq_lens_cpu.max().item())),
163
+ dtype=torch.int32,
164
+ device=device,
165
+ )
166
+
167
+ offset_uncompressed_seq_lens = torch.empty(
168
+ num_reqs + 1, dtype=torch.int32, device=device
169
+ )[1:]
170
+ offset_uncompressed_seq_lens.copy_(uncompressed_seq_lens)
171
+ query_slices = tuple(
172
+ slice(start, num_reqs * query_len + stop)
173
+ for start, stop in _PREFILL_CHUNK_METADATA_QUERY_SLICE_OFFSETS
174
+ )
175
+ for warmup_uncompressed_seq_lens in (
176
+ uncompressed_seq_lens,
177
+ offset_uncompressed_seq_lens,
178
+ ):
179
+ for query_slice in query_slices:
180
+ build_prefill_chunk_metadata(
181
+ 0,
182
+ num_reqs,
183
+ query_start_loc,
184
+ query_start_loc_cpu,
185
+ warmup_uncompressed_seq_lens,
186
+ compressed_seq_lens,
187
+ compressed_seq_lens_cpu,
188
+ block_table,
189
+ compress_ratio,
190
+ query_slice=query_slice,
191
+ )
192
+
193
+
194
+def _warm_combine_topk_swa_indices_kernel(
195
+ device: torch.device,
196
+ num_tokens: int,
197
+ window_size: int,
198
+ compress_ratio: int,
199
+ topk: int,
200
+ topk_width: int,
201
+ n: int,
202
+) -> None:
203
+ from vllm.models.deepseek_v4.common.ops.cache_utils import combine_topk_swa_indices
204
+
205
+ if num_tokens <= 0:
206
+ return
207
+
208
+ def _make_topk_indices(*, offset: bool) -> torch.Tensor:
209
+ if offset:
210
+ topk_storage = torch.full(
211
+ (num_tokens * topk_width + 1,),
212
+ -1,
213
+ dtype=torch.int32,
214
+ device=device,
215
+ )
216
+ topk_indices = topk_storage[1:].reshape(num_tokens, topk_width)
217
+ else:
218
+ topk_indices = torch.full(
219
+ (num_tokens, topk_width), -1, dtype=torch.int32, device=device
220
+ )
221
+ if topk > 0:
222
+ topk_indices.copy_(
223
+ torch.arange(num_tokens * topk_width, dtype=torch.int32, device=device)
224
+ .reshape(num_tokens, topk_width)
225
+ .remainder(topk_width)
226
+ )
227
+ return topk_indices
228
+
229
+ query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
230
+ seq_lens = torch.tensor(
231
+ [window_size + num_tokens], dtype=torch.int32, device=device
232
+ )
233
+ gather_lens = torch.tensor(
234
+ [min(window_size + num_tokens, window_size + num_tokens - 1)],
235
+ dtype=torch.int32,
236
+ device=device,
237
+ )
238
+ offset_query_start_loc = torch.empty(3, dtype=torch.int32, device=device)[1:]
239
+ offset_query_start_loc.copy_(query_start_loc)
240
+ offset_seq_lens = torch.empty(2, dtype=torch.int32, device=device)[1:]
241
+ offset_seq_lens.copy_(seq_lens)
242
+ offset_gather_lens = torch.empty(2, dtype=torch.int32, device=device)[1:]
243
+ offset_gather_lens.copy_(gather_lens)
244
+
245
+ for (
246
+ offset_topk,
247
+ offset_query_and_seq,
248
+ offset_gather,
249
+ ) in _COMBINE_TOPK_SWA_INPUT_VARIANTS:
250
+ warmup_topk_indices = _make_topk_indices(offset=offset_topk)
251
+ warmup_query_start_loc = (
252
+ offset_query_start_loc if offset_query_and_seq else query_start_loc
253
+ )
254
+ warmup_seq_lens = offset_seq_lens if offset_query_and_seq else seq_lens
255
+ warmup_gather_lens = offset_gather_lens if offset_gather else gather_lens
256
+ n_values = (n,) if n == 1 else (n, n + 1)
257
+ for m in (window_size + num_tokens, topk_width):
258
+ for n_value in n_values:
259
+ combine_topk_swa_indices(
260
+ warmup_topk_indices,
261
+ warmup_query_start_loc,
262
+ warmup_seq_lens,
263
+ warmup_gather_lens,
264
+ window_size,
265
+ compress_ratio,
266
+ topk,
267
+ M=m,
268
+ N=n_value,
269
+ )
270
+
271
+
272
+@torch.inference_mode()
273
+def sparse_mla_triton_warmup(
274
+ runner: "GPUModelRunner",
275
+ num_tokens: int,
276
+ *,
277
+ compress_ratios: tuple[int, ...],
278
+ combine_topk_swa_cases: tuple[tuple[int, int, int, int], ...] = (),
279
+) -> None:
280
+ device = getattr(runner, "device", torch.device("cuda"))
281
+ window_size = _hf_config_int(runner, "sliding_window", 128)
282
+
283
+ _warm_sparse_swa_prefill_metadata_kernel(device, window_size, num_tokens)
284
+ for compress_ratio in compress_ratios:
285
+ _warm_prefill_chunk_metadata_kernel(device, compress_ratio, num_tokens)
286
+ for compress_ratio, topk, topk_width, n in combine_topk_swa_cases:
287
+ _warm_combine_topk_swa_indices_kernel(
288
+ device,
289
+ num_tokens,
290
+ window_size,
291
+ compress_ratio,
292
+ topk,
293
+ topk_width,
294
+ n,
295
+ )
296
+
297
+
298
+def deepseek_v4_sparse_triton_warmup(
299
+ runner: "GPUModelRunner",
300
+ num_tokens: int,
301
+) -> None:
302
+ sparse_mla_triton_warmup(
303
+ runner,
304
+ num_tokens,
305
+ compress_ratios=_DSV4_PREFILL_CHUNK_METADATA_COMPRESS_RATIOS,
306
+ combine_topk_swa_cases=_DSV4_COMBINE_TOPK_SWA_WARMUP_CASES,
307
+ )
308
+
309
+
310
+def sparse_mla_triton_warmup_if_needed(worker: "Worker") -> None:
311
+ runner = worker.model_runner
312
+ if runner.is_pooling_model:
313
+ return
314
+
315
+ max_tokens = worker.scheduler_config.max_num_batched_tokens
316
+ num_tokens = _clamp_warmup_tokens(8, max_tokens)
317
+ if num_tokens <= 0:
318
+ return
319
+
320
+ try:
321
+ if _has_attention_backend(runner, _DEEPSEEK_V4_SPARSE_MLA_BACKENDS):
322
+ deepseek_v4_sparse_triton_warmup(runner, num_tokens)
323
+ elif _has_attention_backend(runner, _GENERIC_SPARSE_MLA_BACKENDS):
324
+ sparse_mla_triton_warmup(
325
+ runner,
326
+ num_tokens,
327
+ compress_ratios=(1,),
328
+ )
329
+ except Exception:
330
+ logger.warning("Skipping sparse MLA Triton warmup.", exc_info=True)
vllm/model_executor/warmup/v1_block_table_warmup.py
new
+43
@@ -0,0 +1,43 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Warm up v1 block-table Triton kernels."""
4
+
5
+import torch
6
+
7
+_SLOT_MAPPING_WARMUP_TOKENS = 8
8
+_SLOT_MAPPING_WARMUP_BLOCK_SIZES = (3, 16)
9
+_SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE = 1
10
+
11
+
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
17
+
18
+ num_tokens = max(0, min(_SLOT_MAPPING_WARMUP_TOKENS, max_tokens))
19
+ if num_tokens <= 0:
20
+ return
21
+
22
+ query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
23
+ 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)