fix(qwen): propagate MTP positional overrides

tests/config/test_speculative_draft_hf_overrides.py

@@ -22,6 +22,8 @@
 from vllm.config.parallel import ParallelConfig
 from vllm.config.speculative import SpeculativeConfig
 
+pytestmark = pytest.mark.skip_global_cleanup
+
 
 def _make_hf_config(**kwargs) -> PretrainedConfig:
     defaults = dict(
@@ -43,6 +45,41 @@ def test_dict_overrides_are_not_forwarded_to_draft():
     assert composed is SpeculativeConfig.hf_config_override
 
 
+@pytest.mark.cpu_test
+@pytest.mark.parametrize("method", ["mtp", "draft_model"])
+def test_mtp_draft_receives_target_dict_overrides(method: str):
+    """An in-model MTP draft shares the target's positional geometry."""
+    target_hf_overrides = {
+        "architectures": ["Qwen3_8FlashNextForCausalLM"],
+        "model_type": "qwen3_8_flash_next",
+        "text_config": {
+            "max_position_embeddings": 1048576,
+            "rope_parameters": {"rope_type": "yarn", "factor": 4.0},
+        },
+    }
+    override = SpeculativeConfig.get_draft_hf_overrides(method, target_hf_overrides)
+    assert callable(override)
+    text_config = _make_hf_config(
+        max_position_embeddings=262144,
+        rope_parameters={"rope_type": "default"},
+        hc_count=4,
+        mtp_num_hidden_layers=2,
+        num_attention_heads=4,
+    )
+    source = _make_hf_config(
+        architectures=["Qwen3_8FlashNextForCausalLM"],
+        model_type="qwen3_8_flash_next",
+        text_config=text_config,
+    )
+
+    out = override(source)
+
+    assert out.model_type == "qwen4_exp_mtp"
+    assert out.architectures == ["Qwen4ExpMTP"]
+    assert out.text_config.max_position_embeddings == 1048576
+    assert out.text_config.rope_parameters == {"rope_type": "yarn", "factor": 4.0}
+
+
 @pytest.mark.cpu_test
 def test_none_overrides_fall_back_to_arch_mapping():
     composed = SpeculativeConfig.compose_draft_hf_overrides(None)

tests/models/qwen4_exp/test_b12x_qsa.py

@@ -490,6 +490,7 @@ def test_qsa_selector_fork_requires_one_full_graph(
 
 
 def test_qsa_prefill_context_capacities_cover_the_configured_limit() -> None:
+    """Prefill plan capacities cover each requested context ceiling."""
     assert qsa_module._qsa_prefill_context_capacities(262144, 4096) == (
         4096,
         8192,
@@ -513,6 +514,27 @@ def test_qsa_prefill_context_capacities_cover_the_configured_limit() -> None:
     )
 
 
+def test_qsa_prefill_binding_accepts_pass_one_at_one_million_tokens() -> None:
+    """Pass one uses a prefill plan at the target context ceiling."""
+    owner = Qwen4ExpQSAAttention.__new__(Qwen4ExpQSAAttention)
+    owner.max_decode_rows = 16
+    owner.max_seq_len = 1048576
+    prefill_plan = object()
+    owner._qsa_prefill_bindings = (
+        qsa_module._QSAContextPlan(
+            max_seq_len=1048576,
+            caps=None,
+            plan=prefill_plan,
+        ),
+    )
+    owner._qsa_decode_context = None
+
+    assert (
+        owner._qsa_binding_for_workload(rows=17, max_seq_len=1048576).plan
+        is prefill_plan
+    )
+
+
 def test_qsa_run_consumes_projection_views_and_writes_live_output(monkeypatch) -> None:
     rows = 2
     query = torch.randn(rows, 6, 256, dtype=torch.bfloat16)

vllm/config/speculative.py

@@ -1088,9 +1088,36 @@ def _apply_composed_hf_override(
         target_hf_overrides: Callable[[PretrainedConfig], PretrainedConfig],
         hf_config: PretrainedConfig,
     ) -> PretrainedConfig:
+        """Apply the draft normalization before a target config transform."""
         hf_config = SpeculativeConfig.hf_config_override(hf_config)
         return target_hf_overrides(hf_config)
 
+    @staticmethod
+    def _apply_mtp_dict_hf_override(
+        target_hf_overrides: dict[str, Any],
+        hf_config: PretrainedConfig,
+        *,
+        require_mtp: bool = False,
+    ) -> PretrainedConfig:
+        """Apply target patches without replacing the MTP discriminator."""
+        hf_config = SpeculativeConfig.hf_config_override(hf_config)
+        if require_mtp and hf_config.model_type not in get_args(MTPModelTypes):
+            return hf_config
+        for key, value in target_hf_overrides.items():
+            if key in {"architectures", "model_type"}:
+                continue
+            target = getattr(hf_config, key, None)
+            if isinstance(value, dict) and target is not None:
+                if isinstance(target, dict):
+                    target.update(value)
+                    continue
+                if hasattr(target, "__dict__"):
+                    for nested_key, nested_value in value.items():
+                        setattr(target, nested_key, nested_value)
+                    continue
+            setattr(hf_config, key, value)
+        return hf_config
+
     @staticmethod
     def compose_draft_hf_overrides(
         target_hf_overrides: HfOverrides | None,
@@ -1117,6 +1144,27 @@ def compose_draft_hf_overrides(
             SpeculativeConfig._apply_composed_hf_override, target_hf_overrides
         )
 
+    @staticmethod
+    def get_draft_hf_overrides(
+        method: str,
+        target_hf_overrides: HfOverrides | None,
+    ) -> HfOverrides:
+        """Return overrides appropriate for the selected draft method."""
+        if method == "medusa":
+            return {"model_type": "medusa"}
+        if method == "mtp" and isinstance(target_hf_overrides, dict):
+            return functools.partial(
+                SpeculativeConfig._apply_mtp_dict_hf_override,
+                target_hf_overrides,
+            )
+        if method == "draft_model" and isinstance(target_hf_overrides, dict):
+            return functools.partial(
+                SpeculativeConfig._apply_mtp_dict_hf_override,
+                target_hf_overrides,
+                require_mtp=True,
+            )
+        return SpeculativeConfig.compose_draft_hf_overrides(target_hf_overrides)
+
     @staticmethod
     def _is_custom_proposer_path(model: str | None) -> bool:
         """True if ``model`` is a dotted import path (e.g. ``pkg.MyProposer``)."""
@@ -1130,6 +1178,7 @@ def _is_custom_proposer_path(model: str | None) -> bool:
         return len(parts) >= 2 and all(part.isidentifier() for part in parts)
 
     def __post_init__(self):
+        """Validate speculative settings and construct the draft configuration."""
         # Note: "method" is a new parameter that helps to extend the
         # configuration of non-model-based proposers, and the "model" parameter
         # will be used to set the draft model, eagle head, or additional weight
@@ -1318,15 +1367,9 @@ def __post_init__(self):
                 # detect them. When the method is explicitly "medusa", inject
                 # model_type so MedusaConfig.from_pretrained is used instead.
                 draft_hf_overrides: HfOverrides
-                if self.method == "medusa":
-                    draft_hf_overrides = {"model_type": "medusa"}
-                else:
-                    # Compose any callable hf_overrides set on the target so the
-                    # draft config receives the same transform (e.g. the test
-                    # shrink). Dict overrides stay target-only.
-                    draft_hf_overrides = SpeculativeConfig.compose_draft_hf_overrides(
-                        self.target_model_config.hf_overrides
-                    )
+                draft_hf_overrides = SpeculativeConfig.get_draft_hf_overrides(
+                    self.method, self.target_model_config.hf_overrides
+                )
                 self.draft_model_config = ModelConfig(
                     model=self.model,
                     runner="draft",