[Qwen4Exp] Fix MTP buffer placement, config parsing, Triton warmup (PR #56742)

Intake of open upstream PR vllm-project/vllm#56742. Every rig here runs
MTP on both Qwen3.8 models, so these three bugs sit on our hot path.

The MTP hidden buffer was allocated with torch.empty and no device, so
it followed the ambient default instead of device_config.device; on a
CPU-default process the draft buffer lands on the wrong device. Both
NVIDIA and AMD backends now call make_mtp_hidden_buffer() in
vllm/models/qwen4_exp/common/mtp.py, which pins it to the configured
device.

Config parsing now accepts configs carrying M-RoPE keys and the
Transformers spelling of layer_types; both previously failed at load.

Triton warmup: Qwen4Exp was missing from the Qwen warmup dispatch, so
its Triton kernels compiled during the first served requests instead of
at startup.

Co-Authored-By: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>

tests/model_executor/test_qwen_triton_warmup.py

@@ -1,9 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
+import vllm.model_executor.warmup.qwen_triton_warmup as qwen_triton_warmup_module
 from vllm.model_executor.warmup.qwen_triton_warmup import (
     _FLA_POST_CONV_WARMUP_LENGTHS,
     _QwenGDNWarmupConfig,
@@ -53,3 +56,57 @@ def test_qwen_gdn_prefill_warmup_kernels_compile_on_gpu() -> None:
     _warm_fused_post_conv_kernel(device, config)
     assert _FLA_POST_CONV_WARMUP_LENGTHS == (1, 2, 16)
     torch.accelerator.synchronize(device)
+
+
+def _warmup_runner() -> SimpleNamespace:
+    return SimpleNamespace(
+        device=torch.device("cpu"),
+        compilation_config=SimpleNamespace(static_forward_context={}),
+        max_num_tokens=1,
+        is_pooling_model=False,
+    )
+
+
+def _warmup_model_config(model_type: str) -> SimpleNamespace:
+    return SimpleNamespace(
+        hf_text_config=SimpleNamespace(model_type=model_type),
+        hf_config=SimpleNamespace(model_type=model_type),
+        dtype=torch.bfloat16,
+    )
+
+
+@pytest.mark.parametrize(
+    ("model_type", "dispatched"),
+    [
+        ("qwen4_exp", True),
+        ("qwen4_exp_text", True),
+        ("qwen3_next", True),
+        # The drafter never reaches this warmup: the worker passes its own
+        # model_config, which is the target's.
+        ("qwen4_exp_mtp", False),
+        ("llama", False),
+    ],
+)
+def test_qwen_triton_warmup_dispatch_gate(
+    monkeypatch: pytest.MonkeyPatch,
+    model_type: str,
+    dispatched: bool,
+) -> None:
+    """Which model types reach the Qwen Triton warmup at all.
+
+    Stops at the gate: the kernels themselves need a GPU and are covered by the
+    CUDA test above.
+    """
+    reached: list[object] = []
+
+    def _record(static_forward_context: object) -> None:
+        reached.append(static_forward_context)
+        return None
+
+    monkeypatch.setattr(qwen_triton_warmup_module, "_qwen_gdn_warmup_config", _record)
+
+    qwen_triton_warmup_module.qwen_triton_warmup(
+        _warmup_runner(), _warmup_model_config(model_type)
+    )
+
+    assert bool(reached) is dispatched

tests/models/qwen4_exp/test_config.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import logging
+from collections.abc import Iterator
+from contextlib import contextmanager
 from importlib import import_module
 from types import SimpleNamespace
 from unittest.mock import patch
@@ -13,6 +16,7 @@ from vllm.model_executor.models.config import (
     Qwen3_5ForConditionalGenerationConfig,
     Qwen4ExpForConditionalGenerationConfig,
 )
+from vllm.models.qwen4_exp.common.mtp import make_mtp_hidden_buffer
 from vllm.models.qwen4_exp.config import (
     Qwen4ExpConfig,
     Qwen4ExpTextConfig,
@@ -281,3 +285,121 @@ def test_qwen4_exp_model_state_prepares_stable_dummy_ngram_inputs() -> None:
     )
     assert second["query_start_loc"].data_ptr() == query_start_loc_ptr
     assert second["ngram_context"].data_ptr() == ngram_context_ptr
+
+
+def test_text_config_normalizes_transformers_sparse_attention_spelling():
+    """Checkpoints re-exported through transformers' Qwen4Exp config class
+    serialize "qwen_sparse_attention"; vLLM keys QSA off "full_attention"."""
+    from vllm.models.qwen4_exp.config import Qwen4ExpTextConfig
+
+    config = Qwen4ExpTextConfig(
+        hc_count=4,
+        layer_types=["linear_attention", "qwen_sparse_attention"],
+        num_hidden_layers=2,
+    )
+
+    assert config.layer_types == ["linear_attention", "full_attention"]
+
+
+@contextmanager
+def _captured_transformers_warnings() -> Iterator[list[str]]:
+    """Records what transformers logs, without relying on propagation.
+
+    transformers keeps its own logger hierarchy, and the rope validator reaches
+    it through ``warning_once``, so a root-level capture can miss the record
+    entirely.
+    """
+    records: list[str] = []
+
+    class _Collect(logging.Handler):
+        def emit(self, record: logging.LogRecord) -> None:
+            records.append(record.getMessage())
+
+    handler = _Collect()
+    logger = logging.getLogger("transformers")
+    previous_level = logger.level
+    logger.addHandler(handler)
+    logger.setLevel(logging.WARNING)
+    try:
+        yield records
+    finally:
+        logger.removeHandler(handler)
+        logger.setLevel(previous_level)
+
+
+@pytest.mark.parametrize("wrapped_config", [False, True])
+def test_rope_validation_accepts_mrope_keys_and_still_warns_for_unknown_keys(
+    wrapped_config: bool,
+) -> None:
+    """M-RoPE is driven through two rope_parameters keys transformers does not
+    know for rope_type="default". Declaring them must not silence the
+    unrecognized-key warning for every other key."""
+
+    def build(rope_parameters: dict) -> None:
+        text_config = _text_config(
+            rope_parameters=dict(rope_parameters), rope_theta=10_000.0
+        )
+        if wrapped_config:
+            Qwen4ExpConfig(
+                architectures=["Qwen4ExpForConditionalGeneration"],
+                text_config=text_config.to_dict(),
+            )
+
+    with _captured_transformers_warnings() as records:
+        build(
+            {
+                "rope_type": "default",
+                "mrope_section": [16, 24, 24],
+                "mrope_interleaved": True,
+            }
+        )
+    assert not [line for line in records if "Unrecognized keys" in line], records
+
+    # warning_once dedupes on the message, so each parametrization needs a key
+    # no earlier run has reported.
+    unknown_key = f"not_a_rope_key_{int(wrapped_config)}"
+    with _captured_transformers_warnings() as records:
+        build({"rope_type": "default", unknown_key: 1})
+    assert [line for line in records if unknown_key in line], records
+
+
+@pytest.mark.parametrize("is_last_rank", [True, False])
+def test_mtp_hidden_buffer_is_allocated_on_the_configured_device(
+    is_last_rank: bool,
+) -> None:
+    """The drafter's hidden buffer belongs on the device the config names, not
+    on whichever device happens to be current."""
+    vllm_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(method="mtp"),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
+        model_config=SimpleNamespace(dtype=torch.bfloat16),
+        # "meta" is never the ambient default, so an implicit allocation lands
+        # somewhere else and this test sees it.
+        device_config=SimpleNamespace(device=torch.device("meta")),
+    )
+    config = SimpleNamespace(hc_count=2, hidden_size=4)
+
+    buffer = make_mtp_hidden_buffer(vllm_config, config, is_last_rank=is_last_rank)
+
+    if not is_last_rank:
+        assert buffer is None
+        return
+    assert buffer is not None
+    assert buffer.device.type == "meta"
+    assert buffer.shape == (8, 8)
+    assert buffer.dtype is torch.bfloat16
+
+
+def test_mtp_hidden_buffer_is_absent_without_mtp_speculation() -> None:
+    vllm_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(method="eagle"),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
+        model_config=SimpleNamespace(dtype=torch.bfloat16),
+        device_config=SimpleNamespace(device=torch.device("meta")),
+    )
+    config = SimpleNamespace(hc_count=2, hidden_size=4)
+
+    assert make_mtp_hidden_buffer(vllm_config, config, is_last_rank=True) is None
+
+    vllm_config.speculative_config = None
+    assert make_mtp_hidden_buffer(vllm_config, config, is_last_rank=True) is None

vllm/model_executor/warmup/qwen_triton_warmup.py

@@ -22,6 +22,13 @@ _QWEN_MODEL_TYPES = frozenset(
         "qwen3_5_text",
         "qwen3_5_moe",
         "qwen3_5_moe_text",
+        # Qwen4Exp runs its GDN prefill and QSA on other backends, but the
+        # kernels this file warms are the Triton ones around them -- the gated
+        # RMS norm, the causal conv, the post-conv preparation, the decode
+        # state update -- and those stay whichever prefill kernel is chosen.
+        # Left out of this set they compile on the first prompt instead.
+        "qwen4_exp",
+        "qwen4_exp_text",
     }
 )
 

vllm/models/qwen4_exp/amd/model.py

@@ -77,6 +77,7 @@ from vllm.transformers_utils.configs.qwen4_exp import (
 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.kv_cache_interface import MambaSpec
 
+from ..common.mtp import make_mtp_hidden_buffer
 from ..config import ATTENTION_LAYER_TYPES, QSA_LAYER_TYPE, Qwen4ExpConfig
 from .hyperconnection import GatedResidual, HyperConnectionConfig
 from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
@@ -451,25 +452,11 @@ class Qwen4ExpModel(nn.Module):
         else:
             self.hyper_connection_mixer = None
 
-        spec_config = vllm_config.speculative_config
+        self._mtp_hidden_buffer = make_mtp_hidden_buffer(
-        # MTP HC multi-stream outputs: when speculative method=="mtp" and the
+            vllm_config,
-        # model uses HC with hc_count>1, retain the pre-final-mixer multi-stream
+            config,
-        # hidden state [T, hc_count*H] so the MTP drafter can feed a real
+            is_last_rank=get_pp_group().is_last_rank,
-        # multi-stream backbone hidden on its first step (scheme A). Derived
-        # purely from config (NOT node identity) so P/D nodes stay consistent.
-        needs_mtp_hidden = (
-            spec_config is not None
-            and getattr(spec_config, "method", None) == "mtp"
-            and get_pp_group().is_last_rank
         )
-        if needs_mtp_hidden:
-            self._mtp_hidden_buffer = torch.empty(
-                vllm_config.scheduler_config.max_num_batched_tokens,
-                config.hc_count * config.hidden_size,
-                dtype=vllm_config.model_config.dtype,
-            )
-        else:
-            self._mtp_hidden_buffer = None
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids)

vllm/models/qwen4_exp/common/mtp.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared MTP drafter state for Qwen4Exp."""
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.transformers_utils.configs.qwen4_exp import Qwen4ExpTextConfig
+
+
+def make_mtp_hidden_buffer(
+    vllm_config: VllmConfig,
+    config: Qwen4ExpTextConfig,
+    *,
+    is_last_rank: bool,
+) -> torch.Tensor | None:
+    """Buffer holding the pre-final-mixer multi-stream hidden state, or None.
+
+    When speculative ``method == "mtp"`` the drafter feeds a real multi-stream
+    backbone hidden on its first step, so the last stage retains
+    ``[T, hc_count * H]``. Whether it is needed is derived from the config
+    alone, not from node identity, so prefill and decode nodes agree.
+    """
+    spec_config = vllm_config.speculative_config
+    needs_mtp_hidden = (
+        spec_config is not None
+        and getattr(spec_config, "method", None) == "mtp"
+        and is_last_rank
+    )
+    if not needs_mtp_hidden:
+        return None
+    return torch.empty(
+        vllm_config.scheduler_config.max_num_batched_tokens,
+        config.hc_count * config.hidden_size,
+        dtype=vllm_config.model_config.dtype,
+        device=vllm_config.device_config.device,
+    )

vllm/models/qwen4_exp/config.py

@@ -35,6 +35,11 @@ class Qwen4ExpTextConfig(Qwen3NextConfig):
     model_type = "qwen4_exp_text"
     base_config_key = "text_config"
     keys_to_ignore_at_inference = ["past_key_values"]
+    # Qwen4Exp drives M-RoPE through these two rope_parameters entries, which
+    # transformers' rope validator does not know about for rope_type="default".
+    # Declare them here so validation stays strict for every other key instead
+    # of warning on a configuration the model genuinely uses.
+    ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"}
 
     def __init__(
         self,
@@ -55,6 +60,19 @@ class Qwen4ExpTextConfig(Qwen3NextConfig):
         if hc_count <= 1:
             raise ValueError(f"Qwen4Exp requires hc_count > 1, got {hc_count}.")
 
+        if layer_types is not None:
+            # Transformers' Qwen4Exp config canonicalizes "full_attention"
+            # entries to "qwen_sparse_attention" in memory, so checkpoints
+            # re-exported through it (e.g. quantized exports) serialize that
+            # name. vLLM keys QSA off "full_attention" plus indexer_n_heads,
+            # so map the transformers spelling back.
+            layer_types = [
+                "full_attention"
+                if layer_type == "qwen_sparse_attention"
+                else layer_type
+                for layer_type in layer_types
+            ]
+
         if rope_parameters is not None:
             if kwargs.get("rope_scaling") is None:
                 kwargs["rope_scaling"] = rope_parameters
@@ -208,6 +226,9 @@ class Qwen4ExpConfig(PreTrainedConfig):
         "text_config": Qwen4ExpTextConfig,
     }
     keys_to_ignore_at_inference = ["past_key_values"]
+    # The outer config mirrors the text config's rope_parameters (see below),
+    # so it is validated a second time and needs the same allowance.
+    ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"}
 
     def __init__(
         self,

vllm/models/qwen4_exp/nvidia/model.py

@@ -76,6 +76,7 @@ from vllm.transformers_utils.configs.qwen4_exp import (
 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.kv_cache_interface import MambaSpec
 
+from ..common.mtp import make_mtp_hidden_buffer
 from ..config import ATTENTION_LAYER_TYPES, QSA_LAYER_TYPE, Qwen4ExpConfig
 from .hyperconnection import GatedResidual, HyperConnectionConfig
 from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
@@ -444,25 +445,11 @@ class Qwen4ExpModel(nn.Module):
         else:
             self.hyper_connection_mixer = None
 
-        spec_config = vllm_config.speculative_config
+        self._mtp_hidden_buffer = make_mtp_hidden_buffer(
-        # MTP HC multi-stream outputs: when speculative method=="mtp" and the
+            vllm_config,
-        # model uses HC with hc_count>1, retain the pre-final-mixer multi-stream
+            config,
-        # hidden state [T, hc_count*H] so the MTP drafter can feed a real
+            is_last_rank=get_pp_group().is_last_rank,
-        # multi-stream backbone hidden on its first step (scheme A). Derived
-        # purely from config (NOT node identity) so P/D nodes stay consistent.
-        needs_mtp_hidden = (
-            spec_config is not None
-            and getattr(spec_config, "method", None) == "mtp"
-            and get_pp_group().is_last_rank
         )
-        if needs_mtp_hidden:
-            self._mtp_hidden_buffer = torch.empty(
-                vllm_config.scheduler_config.max_num_batched_tokens,
-                config.hc_count * config.hidden_size,
-                dtype=vllm_config.model_config.dtype,
-            )
-        else:
-            self._mtp_hidden_buffer = None
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids)