[Perf][2/N] Expand Triton kernel warmup coverage, Qwen (#46750)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Roberto L. Castro committed
Jun 29, 2026 at 12:10 UTC
eddfd4cf219359296758272ca736d38cb2c327b1
2 files changed
+389
vllm/model_executor/warmup/kernel_warmup.py
+3
@@ -24,6 +24,7 @@ from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
24
deepseek_v4_sparse_mla_attention_warmup,
25
flashinfer_sparse_mla_decode_autotune_warmup,
26
)
27
+from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
28
from vllm.platforms import current_platform
29
from vllm.utils.deep_gemm import is_deep_gemm_supported
30
from vllm.utils.flashinfer import has_flashinfer
@@ -40,6 +41,8 @@ def kernel_warmup(worker: "Worker"):
41
minimax_m3_msa_warmup,
42
)
43
44
+ qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config)
45
+
46
# DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder
47
# layer per token; warm them across token sizes first so the first real
48
# request doesn't pay JIT cost. No-op for non-DSv4 models (gated inside).
vllm/model_executor/warmup/qwen_triton_warmup.py
new
+386
@@ -0,0 +1,386 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Warm up Qwen Triton kernels from the loaded model's compile keys."""
4
+
5
+from dataclasses import dataclass
6
+from typing import TYPE_CHECKING
7
+
8
+import torch
9
+
10
+from vllm.logger import init_logger
11
+from vllm.model_executor.layers.fla.ops.fused_gdn_prefill_post_conv import (
12
+ fused_post_conv_prep,
13
+)
14
+from vllm.model_executor.layers.fla.ops.fused_sigmoid_gating import (
15
+ fused_sigmoid_gating_delta_rule_update,
16
+)
17
+from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
18
+from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
19
+ causal_conv1d_fn,
20
+)
21
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID
22
+from vllm.v1.worker.block_table import BlockTable
23
+from vllm.v1.worker.utils import _zero_kv_blocks_kernel
24
+
25
+if TYPE_CHECKING:
26
+ from vllm.v1.worker.gpu_model_runner import GPUModelRunner
27
+
28
+logger = init_logger(__name__)
29
+
30
+_QWEN_MODEL_TYPES = frozenset(
31
+ {
32
+ "qwen3_next",
33
+ "qwen3_5",
34
+ "qwen3_5_text",
35
+ "qwen3_5_moe",
36
+ "qwen3_5_moe_text",
37
+ }
38
+)
39
+
40
+_ZERO_KV_N_BLOCKS = (1, 2)
41
+
42
+_SLOT_MAPPING_KV_BLOCK_SIZE = 16
43
+_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE = 1
44
+_SLOT_MAPPING_BLOCK_TABLE_STRIDES = (1, 3)
45
+
46
+# Covers L=1 constexpr, non-divisible runtime L, and divisible runtime L.
47
+_FLA_POST_CONV_WARMUP_LENGTHS = (1, 2, 16)
48
+
49
+
50
+@dataclass(frozen=True)
51
+class _ZeroKvWarmupConfig:
52
+ page_size_el: int
53
+ block_size: int
54
+ n_segs: int
55
+
56
+
57
+@dataclass(frozen=True)
58
+class _QwenGDNWarmupConfig:
59
+ h: int
60
+ hv: int
61
+ k: int
62
+ v: int
63
+ conv_kernel_size: int
64
+ conv_state: torch.Tensor
65
+ conv_dtype: torch.dtype
66
+ a_log: torch.Tensor
67
+ dt_bias: torch.Tensor
68
+ state_stride_token: int
69
+ state_dtype: torch.dtype
70
+
71
+ @property
72
+ def conv_dim(self) -> int:
73
+ return 2 * self.h * self.k + self.hv * self.v
74
+
75
+
76
+def _is_non_empty_tensor(value: object) -> bool:
77
+ return isinstance(value, torch.Tensor) and value.numel() > 0
78
+
79
+
80
+def _is_qwen_gdn_layer(module: object) -> bool:
81
+ return all(
82
+ hasattr(module, attr)
83
+ for attr in (
84
+ "num_k_heads",
85
+ "num_v_heads",
86
+ "head_k_dim",
87
+ "head_v_dim",
88
+ "conv_kernel_size",
89
+ "tp_size",
90
+ "kv_cache",
91
+ "A_log",
92
+ "dt_bias",
93
+ )
94
+ )
95
+
96
+
97
+def _iter_qwen_gdn_layers(static_forward_context: object):
98
+ if not isinstance(static_forward_context, dict):
99
+ return
100
+
101
+ for module in static_forward_context.values():
102
+ if _is_qwen_gdn_layer(module):
103
+ yield module
104
+
105
+
106
+def _split_qwen_gdn_cache(kv_cache: object) -> tuple[torch.Tensor, torch.Tensor] | None:
107
+ if isinstance(kv_cache, (list, tuple)) and len(kv_cache) >= 2:
108
+ conv_cache, ssm_state = kv_cache[:2]
109
+ if _is_non_empty_tensor(conv_cache) and _is_non_empty_tensor(ssm_state):
110
+ return conv_cache, ssm_state
111
+
112
+ if isinstance(kv_cache, torch.Tensor) and kv_cache.size(0) >= 2:
113
+ conv_cache = kv_cache[0]
114
+ ssm_state = kv_cache[1]
115
+ if _is_non_empty_tensor(conv_cache) and _is_non_empty_tensor(ssm_state):
116
+ return conv_cache, ssm_state
117
+ return None
118
+
119
+
120
+def _qwen_gdn_warmup_config(
121
+ static_forward_context: object,
122
+) -> _QwenGDNWarmupConfig | None:
123
+ found_layer = False
124
+ for layer in _iter_qwen_gdn_layers(static_forward_context):
125
+ found_layer = True
126
+ cache_tensors = _split_qwen_gdn_cache(getattr(layer, "kv_cache", None))
127
+ if cache_tensors is None:
128
+ continue
129
+
130
+ conv_cache, ssm_state = cache_tensors
131
+ conv_state = (
132
+ conv_cache if is_conv_state_dim_first() else conv_cache.transpose(-1, -2)
133
+ )
134
+ tp_size = int(layer.tp_size)
135
+ h = int(layer.num_k_heads) // tp_size
136
+ hv = int(layer.num_v_heads) // tp_size
137
+
138
+ return _QwenGDNWarmupConfig(
139
+ h=h,
140
+ hv=hv,
141
+ k=int(layer.head_k_dim),
142
+ v=int(layer.head_v_dim),
143
+ conv_kernel_size=int(layer.conv_kernel_size),
144
+ conv_state=conv_state,
145
+ conv_dtype=conv_state.dtype,
146
+ a_log=layer.A_log,
147
+ dt_bias=layer.dt_bias,
148
+ state_stride_token=int(ssm_state.stride(0)),
149
+ state_dtype=ssm_state.dtype,
150
+ )
151
+
152
+ if found_layer:
153
+ logger.info("Skipping Qwen GDN Triton warmup: no bound Qwen GDN cache found.")
154
+ else:
155
+ logger.info("Skipping Qwen GDN Triton warmup: no Qwen GDN layer found.")
156
+ return None
157
+
158
+
159
+def _get_kv_block_zeroer(runner: object) -> object | None:
160
+ zeroer = getattr(runner, "kv_block_zeroer", None)
161
+ if zeroer is None:
162
+ zeroer = getattr(runner, "_kv_block_zeroer", None)
163
+ return zeroer
164
+
165
+
166
+def _zero_kv_warmup_config(runner: object) -> _ZeroKvWarmupConfig | None:
167
+ zeroer = _get_kv_block_zeroer(runner)
168
+ meta = getattr(zeroer, "_meta", None)
169
+ if meta is None:
170
+ return None
171
+
172
+ _, page_size_el, block_size, n_segs = meta
173
+ return _ZeroKvWarmupConfig(
174
+ page_size_el=int(page_size_el),
175
+ block_size=int(block_size),
176
+ n_segs=int(n_segs),
177
+ )
178
+
179
+
180
+def _warm_zero_kv_blocks_with_runner_zeroer(runner: object) -> bool:
181
+ zeroer = _get_kv_block_zeroer(runner)
182
+ zero_block_ids = getattr(zeroer, "zero_block_ids", None)
183
+ if not callable(zero_block_ids):
184
+ return False
185
+
186
+ for n_blocks in _ZERO_KV_N_BLOCKS:
187
+ zero_block_ids(list(range(n_blocks)))
188
+ return True
189
+
190
+
191
+def _warm_zero_kv_blocks_kernel(
192
+ device: torch.device, config: _ZeroKvWarmupConfig
193
+) -> None:
194
+ max_n_blocks = max(_ZERO_KV_N_BLOCKS)
195
+ scratch = torch.empty(
196
+ max_n_blocks * config.page_size_el,
197
+ dtype=torch.int32,
198
+ device=device,
199
+ )
200
+ seg_addrs = torch.tensor(
201
+ [scratch.data_ptr()] * config.n_segs,
202
+ dtype=torch.uint64,
203
+ device=device,
204
+ )
205
+
206
+ for n_blocks in _ZERO_KV_N_BLOCKS:
207
+ block_ids = torch.arange(n_blocks, dtype=torch.int64, device=device)
208
+ grid = (n_blocks * config.n_segs * (config.page_size_el // config.block_size),)
209
+ _zero_kv_blocks_kernel[grid](
210
+ seg_addrs,
211
+ block_ids,
212
+ n_blocks,
213
+ N_SEGS=config.n_segs,
214
+ PAGE_SIZE_EL=config.page_size_el,
215
+ BLOCK_SIZE=config.block_size,
216
+ num_warps=4,
217
+ num_stages=3,
218
+ )
219
+
220
+
221
+def _warm_compute_slot_mapping_kernel(device: torch.device) -> None:
222
+ # num_tokens/max_num_tokens are do_not_specialize; keep the launch tiny.
223
+ num_tokens = 1
224
+ query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
225
+ positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
226
+
227
+ for block_table_stride in _SLOT_MAPPING_BLOCK_TABLE_STRIDES:
228
+ # Use BlockTable so the JIT key matches the production slot-mapping call.
229
+ block_table = BlockTable(
230
+ block_size=_SLOT_MAPPING_KV_BLOCK_SIZE,
231
+ max_num_reqs=1,
232
+ max_num_blocks_per_req=block_table_stride,
233
+ max_num_batched_tokens=num_tokens,
234
+ pin_memory=False,
235
+ device=device,
236
+ kernel_block_size=_SLOT_MAPPING_KV_BLOCK_SIZE,
237
+ cp_kv_cache_interleave_size=_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE,
238
+ )
239
+ block_table.add_row(list(range(block_table_stride)), 0)
240
+ block_table.commit_block_table(num_reqs=1)
241
+ block_table.compute_slot_mapping(1, query_start_loc, positions)
242
+
243
+
244
+def _warm_causal_conv1d_fwd_kernel(
245
+ device: torch.device, config: _QwenGDNWarmupConfig
246
+) -> None:
247
+ x_storage = torch.empty(
248
+ (1, config.conv_dim), dtype=config.conv_dtype, device=device
249
+ )
250
+ x = x_storage.t()
251
+ weight = torch.empty(
252
+ (config.conv_dim, config.conv_kernel_size),
253
+ dtype=config.conv_dtype,
254
+ device=device,
255
+ )
256
+ cache_indices = torch.full((1,), NULL_BLOCK_ID, dtype=torch.int32, device=device)
257
+ has_initial_state = torch.empty(1, dtype=torch.bool, device=device)
258
+ query_start_loc = torch.tensor([0, 1], dtype=torch.int32, device=device)
259
+
260
+ causal_conv1d_fn(
261
+ x,
262
+ weight,
263
+ None,
264
+ config.conv_state,
265
+ query_start_loc,
266
+ cache_indices=cache_indices,
267
+ has_initial_state=has_initial_state,
268
+ activation="silu",
269
+ pad_slot_id=PAD_SLOT_ID,
270
+ null_block_id=NULL_BLOCK_ID,
271
+ metadata=None,
272
+ validate_data=False,
273
+ )
274
+
275
+
276
+def _warm_fused_post_conv_kernel(
277
+ device: torch.device, config: _QwenGDNWarmupConfig
278
+) -> None:
279
+ qkv_dim = 2 * config.h * config.k + config.hv * config.v
280
+ for length in _FLA_POST_CONV_WARMUP_LENGTHS:
281
+ conv_output = torch.empty(
282
+ (length, qkv_dim), dtype=config.conv_dtype, device=device
283
+ )
284
+ a = torch.empty((length, config.hv), dtype=config.conv_dtype, device=device)
285
+ b = torch.empty_like(a)
286
+
287
+ fused_post_conv_prep(
288
+ conv_output,
289
+ a,
290
+ b,
291
+ config.a_log,
292
+ config.dt_bias,
293
+ config.h,
294
+ config.k,
295
+ config.v,
296
+ apply_l2norm=True,
297
+ output_g_exp=False,
298
+ )
299
+
300
+
301
+def _warm_fused_sigmoid_gating_delta_rule_update_kernel(
302
+ device: torch.device,
303
+ config: _QwenGDNWarmupConfig,
304
+) -> None:
305
+ q = torch.empty((1, 1, config.h, config.k), dtype=config.conv_dtype, device=device)
306
+ k = torch.empty_like(q)
307
+ v = torch.empty((1, 1, config.hv, config.v), dtype=config.conv_dtype, device=device)
308
+ a = torch.empty((1, 1, config.hv), dtype=config.conv_dtype, device=device)
309
+ b = torch.empty_like(a)
310
+ state = torch.empty(
311
+ (1, config.state_stride_token),
312
+ dtype=config.state_dtype,
313
+ device=device,
314
+ )
315
+ cu_seqlens = torch.tensor([0, 1], dtype=torch.int32, device=device)
316
+ ssm_state_indices = torch.empty((1, 1), dtype=torch.int32, device=device)
317
+ ssm_state_indices.zero_()
318
+
319
+ fused_sigmoid_gating_delta_rule_update(
320
+ A_log=config.a_log,
321
+ a=a,
322
+ b=b,
323
+ dt_bias=config.dt_bias,
324
+ q=q,
325
+ k=k,
326
+ v=v,
327
+ beta=1.0,
328
+ threshold=20.0,
329
+ initial_state=state,
330
+ inplace_final_state=True,
331
+ cu_seqlens=cu_seqlens,
332
+ ssm_state_indices=ssm_state_indices,
333
+ use_qk_l2norm_in_kernel=True,
334
+ is_kda=False,
335
+ )
336
+
337
+
338
+def _synchronize_device(device: torch.device) -> None:
339
+ if device.type == "cuda":
340
+ torch.accelerator.synchronize(device)
341
+
342
+
343
+@torch.inference_mode()
344
+def qwen_triton_warmup(
345
+ runner: "GPUModelRunner",
346
+ model_config: object,
347
+) -> None:
348
+ """Warm Qwen Triton kernels reported by the JIT monitor."""
349
+ if runner.is_pooling_model:
350
+ return
351
+
352
+ hf_text_config = getattr(model_config, "hf_text_config", None)
353
+ hf_config = getattr(model_config, "hf_config", None)
354
+ model_type = None
355
+ for config in (hf_text_config, hf_config):
356
+ model_type = getattr(config, "model_type", None)
357
+ if model_type is not None:
358
+ model_type = str(model_type)
359
+ break
360
+ if model_type not in _QWEN_MODEL_TYPES:
361
+ return
362
+
363
+ device = getattr(runner, "device", torch.device("cuda"))
364
+ logger.info("Warming up Qwen Triton kernels for model_type=%s.", model_type)
365
+
366
+ zero_config = _zero_kv_warmup_config(runner)
367
+ if _warm_zero_kv_blocks_with_runner_zeroer(runner):
368
+ pass
369
+ elif zero_config is not None:
370
+ _warm_zero_kv_blocks_kernel(device, zero_config)
371
+ else:
372
+ logger.info("Skipping Qwen zero-kv warmup: no KVBlockZeroer metadata.")
373
+
374
+ _warm_compute_slot_mapping_kernel(device)
375
+ _synchronize_device(device)
376
+
377
+ compilation_config = getattr(runner, "compilation_config", None)
378
+ static_forward_context = getattr(compilation_config, "static_forward_context", None)
379
+ gdn_config = _qwen_gdn_warmup_config(static_forward_context)
380
+ if gdn_config is None:
381
+ return
382
+
383
+ _warm_causal_conv1d_fwd_kernel(device, gdn_config)
384
+ _warm_fused_post_conv_kernel(device, gdn_config)
385
+ _warm_fused_sigmoid_gating_delta_rule_update_kernel(device, gdn_config)
386
+ _synchronize_device(device)