[Bugfix][Spec Decode] Resolve n_predict from text_config for Qwen3.5 multimodal MTP (#55369)

Signed-off-by: SOUMYAJIT GHOSH <23051387@kiit.ac.in>
Signed-off-by: Soumyajit Ghosh <jobsoumyajit6124@gmail.com>
Co-authored-by: SOUMYAJIT GHOSH <23051387@kiit.ac.in>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

.buildkite/test_areas/models_basic.yaml

@@ -134,7 +134,8 @@ steps:
   - tests/models/test_utils.py
   - tests/models/test_vision.py
   - tests/models/test_adapters.py
+  - tests/models/test_qwen3_5_mtp_config.py
   - tests/models/transformers/fusers/
   device: cpu-small
   commands:
-    - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/transformers/fusers/
+    - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/

tests/models/test_qwen3_5_mtp_config.py

@@ -1,19 +1,45 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""CPU-only tests for Qwen3.5 text-only MTP speculative decoding."""
+"""CPU-only tests for Qwen3.5 MTP speculative decoding config overrides."""
+
+from typing import Any
 
 import pytest
-from transformers import PretrainedConfig
+from transformers import AutoConfig, PretrainedConfig
 
 from vllm.config.speculative import SpeculativeConfig
 
+_CHECKPOINTS = {
+    "qwen3_5": "Qwen/Qwen3.8-27B",
+    "qwen3_5_moe": "Qwen/Qwen3.6-35B-A3B",
+}
+
 
 def _mtp_config(model_type: str) -> PretrainedConfig:
-    return PretrainedConfig(
-        model_type=model_type,
-        architectures=["SomeArch"],
-        mtp_num_hidden_layers=1,
-    )
+    """Create a top-level MTP configuration with mtp_num_hidden_layers."""
+    kwargs: dict[str, Any] = {
+        "model_type": model_type,
+        "architectures": ["SomeArch"],
+        "mtp_num_hidden_layers": 1,
+    }
+    return PretrainedConfig(**kwargs)
+
+
+def _multimodal_wrapper_mtp_config(
+    model_type: str, mtp_layers: int = 1
+) -> PretrainedConfig:
+    """Download a multimodal wrapper config via AutoConfig.from_pretrained
+    and configure mtp_num_hidden_layers in text_config.
+
+    Uses real-world Hugging Face Hub checkpoints:
+    - Dense (qwen3_5):     Qwen/Qwen3.8-27B
+    - MoE   (qwen3_5_moe): Qwen/Qwen3.6-35B-A3B
+    """
+    repo = _CHECKPOINTS[model_type]
+    config: PretrainedConfig = AutoConfig.from_pretrained(repo)
+    text_config = config.get_text_config()
+    text_config.mtp_num_hidden_layers = mtp_layers
+    return config
 
 
 @pytest.mark.parametrize(
@@ -26,8 +52,70 @@ def _mtp_config(model_type: str) -> PretrainedConfig:
         ("qwen3_5_moe_text", "Qwen3_5MoeMTP"),
     ],
 )
-def test_mtp_override_recognizes_text_only_types(model_type, expected_arch):
+def test_mtp_override_recognizes_text_only_types(
+    model_type: str, expected_arch: str
+) -> None:
+    """Verify that text-only config variants map to the expected MTP architectures."""
     cfg = SpeculativeConfig.hf_config_override(_mtp_config(model_type))
     assert cfg.model_type == "qwen3_5_mtp"
     assert cfg.architectures == [expected_arch]
     assert cfg.n_predict == 1
+
+
+@pytest.mark.parametrize(
+    "model_type,expected_arch",
+    [
+        ("qwen3_5", "Qwen3_5MTP"),
+        ("qwen3_5_moe", "Qwen3_5MoeMTP"),
+    ],
+)
+def test_mtp_override_extracts_n_predict_from_multimodal_wrapper(
+    model_type: str, expected_arch: str
+) -> None:
+    """Verify that multimodal wrapper checkpoints with mtp_num_hidden_layers
+    in text_config resolve n_predict and architecture correctly."""
+    cfg = SpeculativeConfig.hf_config_override(
+        _multimodal_wrapper_mtp_config(model_type, mtp_layers=2)
+    )
+    assert cfg.model_type == "qwen3_5_mtp"
+    assert cfg.architectures == [expected_arch]
+    assert cfg.n_predict == 2
+
+
+@pytest.mark.parametrize(
+    "model_type,expected_arch",
+    [
+        ("qwen3_5", "Qwen3_5MTP"),
+        ("qwen3_5_moe", "Qwen3_5MoeMTP"),
+    ],
+)
+def test_mtp_override_top_level_precedence_over_nested_text_config(
+    model_type: str, expected_arch: str
+) -> None:
+    """Verify that an explicit top-level mtp_num_hidden_layers takes precedence
+    over a nested text_config value."""
+    cfg = _multimodal_wrapper_mtp_config(model_type, mtp_layers=2)
+    cfg.mtp_num_hidden_layers = 3
+    overridden = SpeculativeConfig.hf_config_override(cfg)
+    assert overridden.model_type == "qwen3_5_mtp"
+    assert overridden.architectures == [expected_arch]
+    assert overridden.n_predict == 3
+
+
+@pytest.mark.parametrize(
+    "model_id,expected_arch",
+    [
+        ("Qwen/Qwen3.8-27B", "Qwen3_5MTP"),
+        ("Qwen/Qwen3.6-35B-A3B", "Qwen3_5MoeMTP"),
+    ],
+)
+def test_mtp_override_downloads_real_hf_hub_configs(
+    model_id: str, expected_arch: str
+) -> None:
+    """Verify that unmodified real-world checkpoints downloaded via
+    AutoConfig.from_pretrained resolve n_predict=1 from text_config."""
+    hf_config = AutoConfig.from_pretrained(model_id)
+    cfg = SpeculativeConfig.hf_config_override(hf_config)
+    assert cfg.model_type == "qwen3_5_mtp"
+    assert cfg.architectures == [expected_arch]
+    assert cfg.n_predict == 1

vllm/config/speculative.py

@@ -897,6 +897,9 @@ class SpeculativeConfig:
             is_moe = hf_config.model_type in ("qwen3_5_moe", "qwen3_5_moe_text")
             hf_config.model_type = "qwen3_5_mtp"
             n_predict = getattr(hf_config, "mtp_num_hidden_layers", None)
+            if n_predict is None:
+                text_config = get_hf_text_config(hf_config)
+                n_predict = getattr(text_config, "mtp_num_hidden_layers", None)
             hf_config.update(
                 {
                     "n_predict": n_predict,