diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index bc5300e53fce..b81f1f4e292f 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/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) diff --git a/tests/models/qwen4_exp/test_b12x_qsa.py b/tests/models/qwen4_exp/test_b12x_qsa.py index 48540e675c36..36452ced19ea 100644 --- a/tests/models/qwen4_exp/test_b12x_qsa.py +++ b/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) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 42df7dbd911d..1384d7897c45 100644 --- a/vllm/config/speculative.py +++ b/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",