[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>

tests/model_executor/test_mamba_triton_warmup.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.model_executor.warmup.mamba_triton_warmup import _warm_batch_memcpy_kernel
+from vllm.platforms import current_platform
+
+
+@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="CUDA is required")
+def test_mamba_batch_memcpy_kernel_compiles_on_gpu() -> None:
+    _warm_batch_memcpy_kernel(torch.device("cuda"))
+    torch.accelerator.synchronize("cuda")

tests/model_executor/test_qwen_triton_warmup.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.model_executor.warmup.qwen_triton_warmup import (
+    _FLA_POST_CONV_WARMUP_LENGTHS,
+    _QwenGDNWarmupConfig,
+    _warm_causal_conv1d_fwd_kernel,
+    _warm_fused_post_conv_kernel,
+    _warm_gated_rms_norm_kernel,
+)
+from vllm.platforms import current_platform
+
+
+def _cuda_gdn_config() -> _QwenGDNWarmupConfig:
+    h, hv, k, v = 2, 2, 16, 16
+    conv_kernel_size = 4
+    conv_dim = 2 * h * k + hv * v
+    device = torch.device("cuda")
+    conv_state = torch.empty(
+        (8, conv_dim, conv_kernel_size - 1),
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    return _QwenGDNWarmupConfig(
+        h=h,
+        hv=hv,
+        k=k,
+        v=v,
+        conv_kernel_size=conv_kernel_size,
+        conv_state=conv_state,
+        conv_dtype=conv_state.dtype,
+        norm_weight_dtype=torch.bfloat16,
+        norm_before_gate=True,
+        norm_activation="silu",
+        a_log=torch.zeros(hv, dtype=torch.float32, device=device),
+        dt_bias=torch.zeros(hv, dtype=torch.float32, device=device),
+        state_stride_token=hv * v * k,
+        state_dtype=torch.float32,
+    )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="CUDA is required")
+def test_qwen_gdn_prefill_warmup_kernels_compile_on_gpu() -> None:
+    config = _cuda_gdn_config()
+    device = torch.device("cuda")
+    _warm_gated_rms_norm_kernel(
+        device, config, max_num_tokens=16, x_dtype=config.conv_dtype
+    )
+    _warm_causal_conv1d_fwd_kernel(device, config)
+    _warm_fused_post_conv_kernel(device, config)
+    assert _FLA_POST_CONV_WARMUP_LENGTHS == (1, 2, 16)
+    torch.accelerator.synchronize(device)

tests/model_executor/test_qwen_vl_triton_warmup.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import torch
+
+from vllm.model_executor.warmup.qwen_vl_triton_warmup import (
+    _warm_mrope,
+    _warm_vision,
+)
+
+
+def test_vision_warmup_calls_only_position_and_rotary_paths() -> None:
+    from vllm.model_executor.models.qwen3_vl import Qwen3_VisionTransformer
+
+    calls: list[tuple[str, object]] = []
+
+    class FakeAttention:
+        num_attention_heads_per_partition = 4
+        hidden_size_per_attention_head = 8
+
+        def apply_rotary_emb(self, qk, cos, sin):
+            calls.append(("rotary", qk.shape))
+
+    class FakeVisual(Qwen3_VisionTransformer):
+        spatial_merge_size = 2
+
+        def __init__(self) -> None:
+            torch.nn.Module.__init__(self)
+            self.blocks = [SimpleNamespace(attn=FakeAttention())]
+            self.pos_embed = torch.nn.Embedding(4, 4, device="cpu")
+
+        def fast_pos_embed_interpolate(self, grid_thw):
+            calls.append(("position", grid_thw[0]))
+
+        def rot_pos_emb(self, grid_thw):
+            _, h, w = grid_thw[0]
+            shape = (h * w, 4)
+            return torch.empty(shape), torch.empty(shape)
+
+        def forward(self, *args, **kwargs):
+            raise AssertionError("vision warmup must not run the full tower")
+
+    model = torch.nn.Module()
+    model.visual = FakeVisual()
+    _warm_vision(model, torch.device("cpu"))
+
+    grids = [(1, 16, 16), (1, 16, 2), (1, 2, 16), (1, 2, 2)]
+    assert [value for name, value in calls if name == "position"] == [
+        list(grid) for grid in grids
+    ]
+    assert [value for name, value in calls if name == "rotary"] == [
+        torch.Size((2, h * w, 4, 8)) for _, h, w in grids
+    ]
+
+    # CPU offloading leaves tower weights off the accelerator, so the Triton
+    # paths must not be launched on host pointers.
+    calls.clear()
+    _warm_vision(model, torch.device("cuda"))
+    assert calls == []
+
+
+def test_mrope_warmup_launches_triton_shapes() -> None:
+    from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding
+
+    launched: list[tuple[torch.Size, torch.Size]] = []
+
+    class FakeRope(MRotaryEmbedding):
+        def __init__(self) -> None:
+            torch.nn.Module.__init__(self)
+            self.head_size = 8
+            self.rotary_dim = 8
+            self.mrope_section = [2, 3, 3]
+            self.mrope_interleaved = False
+            self.is_neox_style = True
+
+        def forward(self, positions, query, key):
+            launched.append((positions.shape, query.shape))
+            return query, key
+
+    class FakeMropeModel(torch.nn.Module):
+        supports_mrope = True
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.rotary_emb = FakeRope()
+
+        def get_mrope_input_positions(self, input_tokens, mm_features):
+            raise AssertionError("warmup must not build live M-RoPE positions")
+
+    runner = SimpleNamespace(
+        model_config=SimpleNamespace(
+            get_num_attention_heads=lambda parallel_config: 4,
+            get_num_kv_heads=lambda parallel_config: 2,
+        ),
+        parallel_config=object(),
+        dtype=torch.bfloat16,
+        device=torch.device("cpu"),
+    )
+    _warm_mrope(runner, FakeMropeModel())
+    assert [shape for shape, _ in launched] == [
+        torch.Size((3, 1)),
+        torch.Size((3, 2)),
+        torch.Size((3, 16)),
+    ]
+
+
+def test_mrope_warmup_skips_models_without_supports_mrope() -> None:
+    def fail(*_args, **_kwargs):
+        raise AssertionError("M-RoPE warmup must not run without SupportsMRoPE")
+
+    runner = SimpleNamespace(
+        model_config=SimpleNamespace(get_num_attention_heads=fail),
+        get_model=fail,
+    )
+    _warm_mrope(runner, torch.nn.Module())

vllm/model_executor/warmup/kernel_warmup.py

@@ -34,10 +34,12 @@ from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
 from vllm.model_executor.warmup.kimi_k3_triton_warmup import (
     kimi_k3_triton_warmup,
 )
+from vllm.model_executor.warmup.mamba_triton_warmup import mamba_triton_warmup
 from vllm.model_executor.warmup.qwen4_exp_qsa_warmup import (
     qwen4_exp_qsa_triton_warmup,
 )
 from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
+from vllm.model_executor.warmup.qwen_vl_triton_warmup import qwen_vl_triton_warmup
 from vllm.model_executor.warmup.replayssm_warmup import (
     replayssm_autotune_warmup,
 )
@@ -149,6 +151,8 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
         )
 
     qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config)
+    qwen_vl_triton_warmup(worker.model_runner)
+    mamba_triton_warmup(worker.model_runner)
 
     compilation_config = worker.vllm_config.compilation_config
     cudagraph_capture_sizes = list(compilation_config.cudagraph_capture_sizes or [])

vllm/model_executor/warmup/mamba_triton_warmup.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Warm Mamba-style Triton kernels shared across GDN / Mamba / KDA models."""
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+logger = init_logger(__name__)
+
+
+def _has_mamba_style_cache(runner: "GPUModelRunner") -> bool:
+    from vllm.v1.kv_cache_interface import MambaSpec, UniformTypeKVCacheSpecs
+
+    for group in runner.kv_cache_config.kv_cache_groups:
+        spec = group.kv_cache_spec
+        if isinstance(spec, UniformTypeKVCacheSpecs):
+            spec = spec.first_spec
+        if isinstance(spec, MambaSpec):
+            return True
+    return False
+
+
+def _warm_batch_memcpy_kernel(device: torch.device) -> None:
+    """Warm the Mamba prefix-cache state copy specialization."""
+    from vllm.v1.worker.mamba_utils import batch_memcpy
+
+    src = torch.empty(1024, dtype=torch.uint8, device=device)
+    dst = torch.empty_like(src)
+    batch_memcpy(
+        torch.tensor([src.data_ptr()], dtype=torch.uint64, device=device),
+        torch.tensor([dst.data_ptr()], dtype=torch.uint64, device=device),
+        # Keep this int32: Triton's compile key includes pointer element dtypes,
+        # and the production prefix-cache path passes an int32 sizes tensor.
+        torch.tensor([src.numel()], dtype=torch.int32, device=device),
+    )
+    logger.info("Warmed Mamba batch_memcpy_kernel.")
+
+
+@torch.inference_mode()
+def mamba_triton_warmup(runner: "GPUModelRunner") -> None:
+    """Warm prefix-cache memcpy for every Mamba-style KV cache group."""
+    device = runner.device
+    if device.type != "cuda":
+        return
+    if not _has_mamba_style_cache(runner):
+        return
+    _warm_batch_memcpy_kernel(device)

vllm/model_executor/warmup/qwen_triton_warmup.py

@@ -283,33 +283,27 @@ def qwen_triton_warmup(
     runner: "GPUModelRunner",
     model_config: "ModelConfig",
 ) -> None:
-    """Warm Qwen Triton kernels reported by the JIT monitor."""
-    if runner.is_pooling_model:
-        return
-
-    hf_text_config = getattr(model_config, "hf_text_config", None)
-    hf_config = getattr(model_config, "hf_config", None)
-    model_type = None
-    for config in (hf_text_config, hf_config):
-        model_type = getattr(config, "model_type", None)
-        if model_type is not None:
-            model_type = str(model_type)
-            break
+    """Warm Qwen GDN Triton kernels reported by the JIT monitor."""
+    model_type = getattr(model_config.hf_text_config, "model_type", "") or getattr(
+        model_config.hf_config, "model_type", ""
+    )
     if model_type not in _QWEN_MODEL_TYPES:
         return
 
-    device = getattr(runner, "device", torch.device("cuda"))
-    logger.info("Warming up Qwen Triton kernels for model_type=%s.", model_type)
+    device = runner.device
+    logger.info("Warming up Qwen GDN Triton kernels for model_type=%s.", model_type)
 
-    compilation_config = getattr(runner, "compilation_config", None)
-    static_forward_context = getattr(compilation_config, "static_forward_context", None)
-    gdn_config = _qwen_gdn_warmup_config(static_forward_context)
+    gdn_config = _qwen_gdn_warmup_config(
+        runner.compilation_config.static_forward_context
+    )
     if gdn_config is None:
         return
 
-    max_num_tokens = max(1, int(getattr(runner, "max_num_tokens", 1)))
+    max_num_tokens = max(1, int(runner.max_num_tokens))
     _warm_gated_rms_norm_kernel(device, gdn_config, max_num_tokens, model_config.dtype)
     _warm_causal_conv1d_fwd_kernel(device, gdn_config)
     _warm_fused_post_conv_kernel(device, gdn_config)
-    _warm_fused_sigmoid_gating_delta_rule_update_kernel(device, gdn_config)
+    # Pooling only runs full prefills; the decode update kernel is unused.
+    if not runner.is_pooling_model:
+        _warm_fused_sigmoid_gating_delta_rule_update_kernel(device, gdn_config)
     _synchronize_device(device)

vllm/model_executor/warmup/qwen_vl_triton_warmup.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Warm Qwen3-VL ViT kernels and M-RoPE for SupportsMRoPE models."""
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.utils.math_utils import round_up
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+logger = init_logger(__name__)
+
+# Covers Triton's ==1 / %16==0 / other integer buckets.
+_MROPE_TOKEN_COUNTS = (1, 2, 16)
+
+
+def _warm_vision(model: torch.nn.Module, device: torch.device) -> None:
+    from vllm.model_executor.models.qwen3_vl import Qwen3_VisionTransformer
+
+    for visual in model.modules():
+        if not isinstance(visual, Qwen3_VisionTransformer):
+            continue
+        # CPU offloading keeps tower weights off the accelerator until the
+        # module forward hook copies them back, so warming here would launch
+        # Triton on host pointers.
+        if visual.pos_embed.weight.device.type != device.type:
+            continue
+        attention = visual.blocks[0].attn
+        num_heads = int(attention.num_attention_heads_per_partition)
+        head_size = int(attention.hidden_size_per_attention_head)
+        merge_size = int(visual.spatial_merge_size)
+        divisible = round_up(16, merge_size)
+        while divisible % 16:
+            divisible += merge_size
+        spatial_sizes = [divisible]
+        if merge_size % 16:
+            spatial_sizes.append(merge_size)
+            if merge_size == 1:
+                spatial_sizes.append(2)
+        grids = [(1, h, w) for h in spatial_sizes for w in spatial_sizes]
+
+        for grid in grids:
+            grid_thw = [list(grid)]
+            visual.fast_pos_embed_interpolate(grid_thw)
+            cos, sin = visual.rot_pos_emb(grid_thw)
+            qk = torch.empty(
+                (2, grid[0] * grid[1] * grid[2], num_heads, head_size),
+                dtype=cos.dtype,
+                device=cos.device,
+            )
+            attention.apply_rotary_emb(qk, cos, sin)
+
+        logger.info(
+            "Warmed position embedding and vision rotary kernels on grids=%s.",
+            grids,
+        )
+
+
+def _warm_mrope(runner: "GPUModelRunner", model: torch.nn.Module) -> None:
+    from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding
+    from vllm.model_executor.models.interfaces import supports_mrope
+
+    is_mrope_model = supports_mrope(model)
+    if not is_mrope_model:
+        return
+
+    rope = next(
+        (module for module in model.modules() if isinstance(module, MRotaryEmbedding)),
+        None,
+    )
+    if rope is None:
+        return
+
+    model_config = runner.model_config
+    parallel_config = runner.parallel_config
+    num_query_heads = int(model_config.get_num_attention_heads(parallel_config))
+    num_kv_heads = int(model_config.get_num_kv_heads(parallel_config))
+    for num_tokens in _MROPE_TOKEN_COUNTS:
+        positions = torch.arange(
+            num_tokens, dtype=torch.long, device=runner.device
+        ).expand(3, -1)
+        query = torch.empty(
+            (num_tokens, num_query_heads * rope.head_size),
+            dtype=runner.dtype,
+            device=runner.device,
+        )
+        key_tensor = torch.empty(
+            (num_tokens, num_kv_heads * rope.head_size),
+            dtype=runner.dtype,
+            device=runner.device,
+        )
+        rope(positions, query, key_tensor)
+
+    logger.info("Warmed M-RoPE Triton kernels.")
+
+
+def _synchronize_device(device: torch.device) -> None:
+    if device.type == "cuda":
+        torch.accelerator.synchronize(device)
+
+
+@torch.inference_mode()
+def qwen_vl_triton_warmup(runner: "GPUModelRunner") -> None:
+    """Warm Qwen3-VL ViT kernels and M-RoPE for ``SupportsMRoPE`` models."""
+    model = runner.get_model()
+    _warm_vision(model, runner.device)
+    _warm_mrope(runner, model)
+    _synchronize_device(runner.device)