[Perf] Extend Qwen Triton warmup to avoid first-request latency spikes (#54797)
Signed-off-by: Juqi Li <2223621784@qq.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
vhagor committed
Sep 7, 2026 at 11:46 UTC
3dc7a68ce45ce1b98c2879139465833611f117cb
7 files changed
+368
-19
tests/model_executor/test_mamba_triton_warmup.py
new
+14
@@ -0,0 +1,14 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.model_executor.warmup.mamba_triton_warmup import _warm_batch_memcpy_kernel
8
+from vllm.platforms import current_platform
9
+
10
+
11
+@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="CUDA is required")
12
+def test_mamba_batch_memcpy_kernel_compiles_on_gpu() -> None:
13
+ _warm_batch_memcpy_kernel(torch.device("cuda"))
14
+ torch.accelerator.synchronize("cuda")
tests/model_executor/test_qwen_triton_warmup.py
new
+55
@@ -0,0 +1,55 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.model_executor.warmup.qwen_triton_warmup import (
8
+ _FLA_POST_CONV_WARMUP_LENGTHS,
9
+ _QwenGDNWarmupConfig,
10
+ _warm_causal_conv1d_fwd_kernel,
11
+ _warm_fused_post_conv_kernel,
12
+ _warm_gated_rms_norm_kernel,
13
+)
14
+from vllm.platforms import current_platform
15
+
16
+
17
+def _cuda_gdn_config() -> _QwenGDNWarmupConfig:
18
+ h, hv, k, v = 2, 2, 16, 16
19
+ conv_kernel_size = 4
20
+ conv_dim = 2 * h * k + hv * v
21
+ device = torch.device("cuda")
22
+ conv_state = torch.empty(
23
+ (8, conv_dim, conv_kernel_size - 1),
24
+ dtype=torch.bfloat16,
25
+ device=device,
26
+ )
27
+ return _QwenGDNWarmupConfig(
28
+ h=h,
29
+ hv=hv,
30
+ k=k,
31
+ v=v,
32
+ conv_kernel_size=conv_kernel_size,
33
+ conv_state=conv_state,
34
+ conv_dtype=conv_state.dtype,
35
+ norm_weight_dtype=torch.bfloat16,
36
+ norm_before_gate=True,
37
+ norm_activation="silu",
38
+ a_log=torch.zeros(hv, dtype=torch.float32, device=device),
39
+ dt_bias=torch.zeros(hv, dtype=torch.float32, device=device),
40
+ state_stride_token=hv * v * k,
41
+ state_dtype=torch.float32,
42
+ )
43
+
44
+
45
+@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="CUDA is required")
46
+def test_qwen_gdn_prefill_warmup_kernels_compile_on_gpu() -> None:
47
+ config = _cuda_gdn_config()
48
+ device = torch.device("cuda")
49
+ _warm_gated_rms_norm_kernel(
50
+ device, config, max_num_tokens=16, x_dtype=config.conv_dtype
51
+ )
52
+ _warm_causal_conv1d_fwd_kernel(device, config)
53
+ _warm_fused_post_conv_kernel(device, config)
54
+ assert _FLA_POST_CONV_WARMUP_LENGTHS == (1, 2, 16)
55
+ torch.accelerator.synchronize(device)
tests/model_executor/test_qwen_vl_triton_warmup.py
new
+117
@@ -0,0 +1,117 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import torch
7
+
8
+from vllm.model_executor.warmup.qwen_vl_triton_warmup import (
9
+ _warm_mrope,
10
+ _warm_vision,
11
+)
12
+
13
+
14
+def test_vision_warmup_calls_only_position_and_rotary_paths() -> None:
15
+ from vllm.model_executor.models.qwen3_vl import Qwen3_VisionTransformer
16
+
17
+ calls: list[tuple[str, object]] = []
18
+
19
+ class FakeAttention:
20
+ num_attention_heads_per_partition = 4
21
+ hidden_size_per_attention_head = 8
22
+
23
+ def apply_rotary_emb(self, qk, cos, sin):
24
+ calls.append(("rotary", qk.shape))
25
+
26
+ class FakeVisual(Qwen3_VisionTransformer):
27
+ spatial_merge_size = 2
28
+
29
+ def __init__(self) -> None:
30
+ torch.nn.Module.__init__(self)
31
+ self.blocks = [SimpleNamespace(attn=FakeAttention())]
32
+ self.pos_embed = torch.nn.Embedding(4, 4, device="cpu")
33
+
34
+ def fast_pos_embed_interpolate(self, grid_thw):
35
+ calls.append(("position", grid_thw[0]))
36
+
37
+ def rot_pos_emb(self, grid_thw):
38
+ _, h, w = grid_thw[0]
39
+ shape = (h * w, 4)
40
+ return torch.empty(shape), torch.empty(shape)
41
+
42
+ def forward(self, *args, **kwargs):
43
+ raise AssertionError("vision warmup must not run the full tower")
44
+
45
+ model = torch.nn.Module()
46
+ model.visual = FakeVisual()
47
+ _warm_vision(model, torch.device("cpu"))
48
+
49
+ grids = [(1, 16, 16), (1, 16, 2), (1, 2, 16), (1, 2, 2)]
50
+ assert [value for name, value in calls if name == "position"] == [
51
+ list(grid) for grid in grids
52
+ ]
53
+ assert [value for name, value in calls if name == "rotary"] == [
54
+ torch.Size((2, h * w, 4, 8)) for _, h, w in grids
55
+ ]
56
+
57
+ # CPU offloading leaves tower weights off the accelerator, so the Triton
58
+ # paths must not be launched on host pointers.
59
+ calls.clear()
60
+ _warm_vision(model, torch.device("cuda"))
61
+ assert calls == []
62
+
63
+
64
+def test_mrope_warmup_launches_triton_shapes() -> None:
65
+ from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding
66
+
67
+ launched: list[tuple[torch.Size, torch.Size]] = []
68
+
69
+ class FakeRope(MRotaryEmbedding):
70
+ def __init__(self) -> None:
71
+ torch.nn.Module.__init__(self)
72
+ self.head_size = 8
73
+ self.rotary_dim = 8
74
+ self.mrope_section = [2, 3, 3]
75
+ self.mrope_interleaved = False
76
+ self.is_neox_style = True
77
+
78
+ def forward(self, positions, query, key):
79
+ launched.append((positions.shape, query.shape))
80
+ return query, key
81
+
82
+ class FakeMropeModel(torch.nn.Module):
83
+ supports_mrope = True
84
+
85
+ def __init__(self) -> None:
86
+ super().__init__()
87
+ self.rotary_emb = FakeRope()
88
+
89
+ def get_mrope_input_positions(self, input_tokens, mm_features):
90
+ raise AssertionError("warmup must not build live M-RoPE positions")
91
+
92
+ runner = SimpleNamespace(
93
+ model_config=SimpleNamespace(
94
+ get_num_attention_heads=lambda parallel_config: 4,
95
+ get_num_kv_heads=lambda parallel_config: 2,
96
+ ),
97
+ parallel_config=object(),
98
+ dtype=torch.bfloat16,
99
+ device=torch.device("cpu"),
100
+ )
101
+ _warm_mrope(runner, FakeMropeModel())
102
+ assert [shape for shape, _ in launched] == [
103
+ torch.Size((3, 1)),
104
+ torch.Size((3, 2)),
105
+ torch.Size((3, 16)),
106
+ ]
107
+
108
+
109
+def test_mrope_warmup_skips_models_without_supports_mrope() -> None:
110
+ def fail(*_args, **_kwargs):
111
+ raise AssertionError("M-RoPE warmup must not run without SupportsMRoPE")
112
+
113
+ runner = SimpleNamespace(
114
+ model_config=SimpleNamespace(get_num_attention_heads=fail),
115
+ get_model=fail,
116
+ )
117
+ _warm_mrope(runner, torch.nn.Module())
vllm/model_executor/warmup/kernel_warmup.py
+4
@@ -34,10 +34,12 @@ from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
34
from vllm.model_executor.warmup.kimi_k3_triton_warmup import (
35
kimi_k3_triton_warmup,
36
)
37
+from vllm.model_executor.warmup.mamba_triton_warmup import mamba_triton_warmup
38
from vllm.model_executor.warmup.qwen4_exp_qsa_warmup import (
39
qwen4_exp_qsa_triton_warmup,
40
)
41
from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
42
+from vllm.model_executor.warmup.qwen_vl_triton_warmup import qwen_vl_triton_warmup
43
from vllm.model_executor.warmup.replayssm_warmup import (
44
replayssm_autotune_warmup,
45
)
@@ -149,6 +151,8 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
151
)
152
153
qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config)
154
+ qwen_vl_triton_warmup(worker.model_runner)
155
+ mamba_triton_warmup(worker.model_runner)
156
157
compilation_config = worker.vllm_config.compilation_config
158
cudagraph_capture_sizes = list(compilation_config.cudagraph_capture_sizes or [])
vllm/model_executor/warmup/mamba_triton_warmup.py
new
+53
@@ -0,0 +1,53 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Warm Mamba-style Triton kernels shared across GDN / Mamba / KDA models."""
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
+
14
+logger = init_logger(__name__)
15
+
16
+
17
+def _has_mamba_style_cache(runner: "GPUModelRunner") -> bool:
18
+ from vllm.v1.kv_cache_interface import MambaSpec, UniformTypeKVCacheSpecs
19
+
20
+ for group in runner.kv_cache_config.kv_cache_groups:
21
+ spec = group.kv_cache_spec
22
+ if isinstance(spec, UniformTypeKVCacheSpecs):
23
+ spec = spec.first_spec
24
+ if isinstance(spec, MambaSpec):
25
+ return True
26
+ return False
27
+
28
+
29
+def _warm_batch_memcpy_kernel(device: torch.device) -> None:
30
+ """Warm the Mamba prefix-cache state copy specialization."""
31
+ from vllm.v1.worker.mamba_utils import batch_memcpy
32
+
33
+ src = torch.empty(1024, dtype=torch.uint8, device=device)
34
+ dst = torch.empty_like(src)
35
+ batch_memcpy(
36
+ torch.tensor([src.data_ptr()], dtype=torch.uint64, device=device),
37
+ torch.tensor([dst.data_ptr()], dtype=torch.uint64, device=device),
38
+ # Keep this int32: Triton's compile key includes pointer element dtypes,
39
+ # and the production prefix-cache path passes an int32 sizes tensor.
40
+ torch.tensor([src.numel()], dtype=torch.int32, device=device),
41
+ )
42
+ logger.info("Warmed Mamba batch_memcpy_kernel.")
43
+
44
+
45
+@torch.inference_mode()
46
+def mamba_triton_warmup(runner: "GPUModelRunner") -> None:
47
+ """Warm prefix-cache memcpy for every Mamba-style KV cache group."""
48
+ device = runner.device
49
+ if device.type != "cuda":
50
+ return
51
+ if not _has_mamba_style_cache(runner):
52
+ return
53
+ _warm_batch_memcpy_kernel(device)
vllm/model_executor/warmup/qwen_triton_warmup.py
+13
-19
@@ -283,33 +283,27 @@ def qwen_triton_warmup(
283
runner: "GPUModelRunner",
284
model_config: "ModelConfig",
285
) -> None:
286
- """Warm Qwen Triton kernels reported by the JIT monitor."""
287
- if runner.is_pooling_model:
288
- return
289
-
290
- hf_text_config = getattr(model_config, "hf_text_config", None)
291
- hf_config = getattr(model_config, "hf_config", None)
292
- model_type = None
293
- for config in (hf_text_config, hf_config):
294
- model_type = getattr(config, "model_type", None)
295
- if model_type is not None:
296
- model_type = str(model_type)
297
- break
286
+ """Warm Qwen GDN Triton kernels reported by the JIT monitor."""
287
+ model_type = getattr(model_config.hf_text_config, "model_type", "") or getattr(
288
+ model_config.hf_config, "model_type", ""
289
+ )
290
if model_type not in _QWEN_MODEL_TYPES:
291
return
292
301
- device = getattr(runner, "device", torch.device("cuda"))
302
- logger.info("Warming up Qwen Triton kernels for model_type=%s.", model_type)
293
+ device = runner.device
294
+ logger.info("Warming up Qwen GDN Triton kernels for model_type=%s.", model_type)
295
304
- compilation_config = getattr(runner, "compilation_config", None)
305
- static_forward_context = getattr(compilation_config, "static_forward_context", None)
306
- gdn_config = _qwen_gdn_warmup_config(static_forward_context)
296
+ gdn_config = _qwen_gdn_warmup_config(
297
+ runner.compilation_config.static_forward_context
298
+ )
299
if gdn_config is None:
300
return
301
310
- max_num_tokens = max(1, int(getattr(runner, "max_num_tokens", 1)))
302
+ max_num_tokens = max(1, int(runner.max_num_tokens))
303
_warm_gated_rms_norm_kernel(device, gdn_config, max_num_tokens, model_config.dtype)
304
_warm_causal_conv1d_fwd_kernel(device, gdn_config)
305
_warm_fused_post_conv_kernel(device, gdn_config)
314
- _warm_fused_sigmoid_gating_delta_rule_update_kernel(device, gdn_config)
306
+ # Pooling only runs full prefills; the decode update kernel is unused.
307
+ if not runner.is_pooling_model:
308
+ _warm_fused_sigmoid_gating_delta_rule_update_kernel(device, gdn_config)
309
_synchronize_device(device)
vllm/model_executor/warmup/qwen_vl_triton_warmup.py
new
+112
@@ -0,0 +1,112 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Warm Qwen3-VL ViT kernels and M-RoPE for SupportsMRoPE models."""
4
+
5
+from typing import TYPE_CHECKING
6
+
7
+import torch
8
+
9
+from vllm.logger import init_logger
10
+from vllm.utils.math_utils import round_up
11
+
12
+if TYPE_CHECKING:
13
+ from vllm.v1.worker.gpu_model_runner import GPUModelRunner
14
+
15
+logger = init_logger(__name__)
16
+
17
+# Covers Triton's ==1 / %16==0 / other integer buckets.
18
+_MROPE_TOKEN_COUNTS = (1, 2, 16)
19
+
20
+
21
+def _warm_vision(model: torch.nn.Module, device: torch.device) -> None:
22
+ from vllm.model_executor.models.qwen3_vl import Qwen3_VisionTransformer
23
+
24
+ for visual in model.modules():
25
+ if not isinstance(visual, Qwen3_VisionTransformer):
26
+ continue
27
+ # CPU offloading keeps tower weights off the accelerator until the
28
+ # module forward hook copies them back, so warming here would launch
29
+ # Triton on host pointers.
30
+ if visual.pos_embed.weight.device.type != device.type:
31
+ continue
32
+ attention = visual.blocks[0].attn
33
+ num_heads = int(attention.num_attention_heads_per_partition)
34
+ head_size = int(attention.hidden_size_per_attention_head)
35
+ merge_size = int(visual.spatial_merge_size)
36
+ divisible = round_up(16, merge_size)
37
+ while divisible % 16:
38
+ divisible += merge_size
39
+ spatial_sizes = [divisible]
40
+ if merge_size % 16:
41
+ spatial_sizes.append(merge_size)
42
+ if merge_size == 1:
43
+ spatial_sizes.append(2)
44
+ grids = [(1, h, w) for h in spatial_sizes for w in spatial_sizes]
45
+
46
+ for grid in grids:
47
+ grid_thw = [list(grid)]
48
+ visual.fast_pos_embed_interpolate(grid_thw)
49
+ cos, sin = visual.rot_pos_emb(grid_thw)
50
+ qk = torch.empty(
51
+ (2, grid[0] * grid[1] * grid[2], num_heads, head_size),
52
+ dtype=cos.dtype,
53
+ device=cos.device,
54
+ )
55
+ attention.apply_rotary_emb(qk, cos, sin)
56
+
57
+ logger.info(
58
+ "Warmed position embedding and vision rotary kernels on grids=%s.",
59
+ grids,
60
+ )
61
+
62
+
63
+def _warm_mrope(runner: "GPUModelRunner", model: torch.nn.Module) -> None:
64
+ from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding
65
+ from vllm.model_executor.models.interfaces import supports_mrope
66
+
67
+ is_mrope_model = supports_mrope(model)
68
+ if not is_mrope_model:
69
+ return
70
+
71
+ rope = next(
72
+ (module for module in model.modules() if isinstance(module, MRotaryEmbedding)),
73
+ None,
74
+ )
75
+ if rope is None:
76
+ return
77
+
78
+ model_config = runner.model_config
79
+ parallel_config = runner.parallel_config
80
+ num_query_heads = int(model_config.get_num_attention_heads(parallel_config))
81
+ num_kv_heads = int(model_config.get_num_kv_heads(parallel_config))
82
+ for num_tokens in _MROPE_TOKEN_COUNTS:
83
+ positions = torch.arange(
84
+ num_tokens, dtype=torch.long, device=runner.device
85
+ ).expand(3, -1)
86
+ query = torch.empty(
87
+ (num_tokens, num_query_heads * rope.head_size),
88
+ dtype=runner.dtype,
89
+ device=runner.device,
90
+ )
91
+ key_tensor = torch.empty(
92
+ (num_tokens, num_kv_heads * rope.head_size),
93
+ dtype=runner.dtype,
94
+ device=runner.device,
95
+ )
96
+ rope(positions, query, key_tensor)
97
+
98
+ logger.info("Warmed M-RoPE Triton kernels.")
99
+
100
+
101
+def _synchronize_device(device: torch.device) -> None:
102
+ if device.type == "cuda":
103
+ torch.accelerator.synchronize(device)
104
+
105
+
106
+@torch.inference_mode()
107
+def qwen_vl_triton_warmup(runner: "GPUModelRunner") -> None:
108
+ """Warm Qwen3-VL ViT kernels and M-RoPE for ``SupportsMRoPE`` models."""
109
+ model = runner.get_model()
110
+ _warm_vision(model, runner.device)
111
+ _warm_mrope(runner, model)
112
+ _synchronize_device(runner.device)