From 3b1911c0cac39420f1b2acd2e5d76d403254bb88 Mon Sep 17 00:00:00 2001 From: lesj0610 Date: Mon, 14 Sep 2026 07:25:01 +0900 Subject: [PATCH] [Bugfix][Model] Fix Qwen4Exp MTP buffer placement, config parsing, and Triton warmup Three initialization and config defects found serving Qwen4Exp, each with a regression test. **The MTP hidden buffer lands on the ambient device.** torch.empty for the spec-decode hidden buffer takes no device, so it follows whatever device is current rather than vllm_config.device_config.device. The allocation was duplicated verbatim in the NVIDIA and AMD backends; it moves to common/mtp.py so the two cannot drift, and the caller passes is_last_rank rather than reaching for the process group from shared code. **M-RoPE keys are reported as unrecognized.** Qwen4Exp drives M-RoPE through mrope_section and mrope_interleaved, which transformers' rope validator does not know for rope_type="default", so every load warns about a configuration the model genuinely uses. Declare them at rope validation on the text config and on the outer config that mirrors it; every other key still warns. **transformers' layer_types spelling is unread.** Its Qwen4Exp config canonicalizes "full_attention" entries to "qwen_sparse_attention" in memory, so checkpoints re-exported through it serialize that name. vLLM keys QSA off "full_attention" plus indexer_n_heads, so map the spelling back. **Qwen4Exp is missing from the Triton warmup set.** Its GDN prefill and QSA run on other backends, but the kernels this warmup covers -- the gated RMS norm, the causal conv, the post-conv preparation, the decode state update -- are Triton whichever prefill kernel is chosen, so they compile on the first prompt instead. qwen4_exp_mtp stays out: kernel_warmup passes the worker's own model_config, which is the target model's. tests/models/qwen4_exp/test_config.py: 15 passed. tests/model_executor/test_qwen_triton_warmup.py: 6 passed. ruff check, ruff format --check and mypy (Python 3.10) report what they already report on the base commit. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: lesj0610 --- .../model_executor/test_qwen_triton_warmup.py | 57 ++++++++ tests/models/qwen4_exp/test_config.py | 122 ++++++++++++++++++ .../warmup/qwen_triton_warmup.py | 7 + vllm/models/qwen4_exp/amd/model.py | 23 +--- vllm/models/qwen4_exp/common/mtp.py | 37 ++++++ vllm/models/qwen4_exp/config.py | 21 +++ vllm/models/qwen4_exp/nvidia/model.py | 23 +--- 7 files changed, 254 insertions(+), 36 deletions(-) create mode 100644 vllm/models/qwen4_exp/common/mtp.py diff --git a/tests/model_executor/test_qwen_triton_warmup.py b/tests/model_executor/test_qwen_triton_warmup.py index 81bef605d5..7699be25c2 100644 --- a/tests/model_executor/test_qwen_triton_warmup.py +++ b/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 diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index 9d30365b81..4b9d42a4aa 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/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, @@ -246,3 +250,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 diff --git a/vllm/model_executor/warmup/qwen_triton_warmup.py b/vllm/model_executor/warmup/qwen_triton_warmup.py index c6b6abea2e..38f9051975 100644 --- a/vllm/model_executor/warmup/qwen_triton_warmup.py +++ b/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", } ) diff --git a/vllm/models/qwen4_exp/amd/model.py b/vllm/models/qwen4_exp/amd/model.py index aeb3a14dad..ef06a6fe41 100644 --- a/vllm/models/qwen4_exp/amd/model.py +++ b/vllm/models/qwen4_exp/amd/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 Qwen4ExpConfig from .hyperconnection import GatedResidual, HyperConnectionConfig from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm @@ -446,25 +447,11 @@ class Qwen4ExpModel(nn.Module): else: self.hyper_connection_mixer = None - spec_config = vllm_config.speculative_config - # MTP HC multi-stream outputs: when speculative method=="mtp" and the - # model uses HC with hc_count>1, retain the pre-final-mixer multi-stream - # hidden state [T, hc_count*H] so the MTP drafter can feed a real - # 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 + self._mtp_hidden_buffer = make_mtp_hidden_buffer( + vllm_config, + config, + is_last_rank=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) diff --git a/vllm/models/qwen4_exp/common/mtp.py b/vllm/models/qwen4_exp/common/mtp.py new file mode 100644 index 0000000000..457e8bc064 --- /dev/null +++ b/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, + ) diff --git a/vllm/models/qwen4_exp/config.py b/vllm/models/qwen4_exp/config.py index a8eedab83f..b0e1a4a39e 100644 --- a/vllm/models/qwen4_exp/config.py +++ b/vllm/models/qwen4_exp/config.py @@ -29,6 +29,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, @@ -49,6 +54,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 @@ -197,6 +215,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, diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 28c3ea69d1..83a3c74dfa 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -75,6 +75,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 Qwen4ExpConfig from .hyperconnection import GatedResidual, HyperConnectionConfig from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm @@ -438,25 +439,11 @@ class Qwen4ExpModel(nn.Module): else: self.hyper_connection_mixer = None - spec_config = vllm_config.speculative_config - # MTP HC multi-stream outputs: when speculative method=="mtp" and the - # model uses HC with hc_count>1, retain the pre-final-mixer multi-stream - # hidden state [T, hc_count*H] so the MTP drafter can feed a real - # 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 + self._mtp_hidden_buffer = make_mtp_hidden_buffer( + vllm_config, + config, + is_last_rank=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)