| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for draft config overrides used by SpeculativeConfig. |
| 4 | |
| 5 | Callable ``hf_overrides`` on the target model config (e.g. the |
| 6 | ``dummy_hf_overrides`` shrink used by ``tests/models/test_initialization.py``) |
| 7 | must also be applied when building the draft ``ModelConfig``. Otherwise a |
| 8 | draft belonging to a large target model is instantiated at full size even |
| 9 | when the target itself is shrunk — which is what kept spec-decode archs like |
| 10 | ``EagleMistralLarge3ForCausalLM`` stuck at ``is_available_online=False`` |
| 11 | ("TODO: revert once figuring out OOM in CI"). |
| 12 | """ |
| 13 | |
| 14 | import functools |
| 15 | from unittest.mock import MagicMock, patch |
| 16 | |
| 17 | import pytest |
| 18 | from transformers import PreTrainedConfig |
| 19 | |
| 20 | from vllm.config.parallel import ParallelConfig |
| 21 | from vllm.config.speculative import SpeculativeConfig |
| 22 | |
| 23 | |
| 24 | def _make_hf_config(**kwargs) -> PreTrainedConfig: |
| 25 | defaults = dict( |
| 26 | architectures=["LlamaForCausalLM"], |
| 27 | model_type="llama", |
| 28 | num_hidden_layers=64, |
| 29 | ) |
| 30 | defaults.update(kwargs) |
| 31 | return PreTrainedConfig(**defaults) |
| 32 | |
| 33 | |
| 34 | @pytest.mark.cpu_test |
| 35 | def test_dict_overrides_are_not_forwarded_to_draft(): |
| 36 | """Dict overrides are target-specific key patches; the draft must get |
| 37 | only the architecture-mapping override.""" |
| 38 | composed = SpeculativeConfig.compose_draft_hf_overrides( |
| 39 | {"max_position_embeddings": 1234} |
| 40 | ) |
| 41 | assert composed is SpeculativeConfig.hf_config_override |
| 42 | |
| 43 | |
| 44 | @pytest.mark.cpu_test |
| 45 | def test_none_overrides_fall_back_to_arch_mapping(): |
| 46 | composed = SpeculativeConfig.compose_draft_hf_overrides(None) |
| 47 | assert composed is SpeculativeConfig.hf_config_override |
| 48 | |
| 49 | |
| 50 | def _make_speculative_config( |
| 51 | hf_kwargs: dict, architecture: str, **speculative_kwargs |
| 52 | ) -> SpeculativeConfig: |
| 53 | hf_config = SpeculativeConfig.hf_config_override(_make_hf_config(**hf_kwargs)) |
| 54 | draft = MagicMock( |
| 55 | hf_config=hf_config, architectures=hf_config.architectures, max_model_len=128 |
| 56 | ) |
| 57 | draft.registry.inspect_model_cls.return_value = (None, architecture) |
| 58 | target = MagicMock(max_model_len=128, quantization=None, hf_overrides={}) |
| 59 | target.hf_config.model_type = hf_kwargs["model_type"] |
| 60 | |
| 61 | with patch("vllm.config.speculative.ModelConfig", return_value=draft): |
| 62 | return SpeculativeConfig( |
| 63 | model="draft", |
| 64 | **speculative_kwargs, |
| 65 | target_model_config=target, |
| 66 | target_parallel_config=ParallelConfig(), |
| 67 | ) |
| 68 | |
| 69 | |
| 70 | @pytest.mark.cpu_test |
| 71 | @pytest.mark.parametrize( |
| 72 | ("hf_kwargs", "n_predict", "architecture"), |
| 73 | [ |
| 74 | ( |
| 75 | dict( |
| 76 | model_type="deepseek_v4", |
| 77 | num_nextn_predict_layers=3, |
| 78 | dspark_block_size=5, |
| 79 | ), |
| 80 | 3, |
| 81 | "DSparkDraftModel", |
| 82 | ), |
| 83 | ( |
| 84 | dict( |
| 85 | model_type="deepseek_v41", |
| 86 | num_nextn_predict_layers=3, |
| 87 | dspark_block_size=5, |
| 88 | ), |
| 89 | 3, |
| 90 | "DSparkV41DraftModel", |
| 91 | ), |
| 92 | ( |
| 93 | dict( |
| 94 | model_type="gemma4_text", |
| 95 | architectures=["Gemma4DSparkModel"], |
| 96 | block_size=5, |
| 97 | ), |
| 98 | None, |
| 99 | "Gemma4DSparkModel", |
| 100 | ), |
| 101 | ], |
| 102 | ) |
| 103 | @pytest.mark.parametrize("speculative_kwargs", [{"num_speculative_tokens": 5}, {}]) |
| 104 | def test_dspark_width_is_independent_of_mtp_stages( |
| 105 | hf_kwargs, n_predict, architecture, speculative_kwargs |
| 106 | ): |
| 107 | config = _make_speculative_config( |
| 108 | hf_kwargs, architecture, method="dspark", **speculative_kwargs |
| 109 | ) |
| 110 | |
| 111 | assert config.num_speculative_tokens == 5 |
| 112 | assert getattr(config.draft_model_config.hf_config, "n_predict", None) == n_predict |
| 113 | assert config.draft_model_config.hf_config.architectures == [architecture] |
| 114 | |
| 115 | |
| 116 | @pytest.mark.cpu_test |
| 117 | def test_mtp_stages_are_independent_of_dspark_width(): |
| 118 | hf_kwargs = dict( |
| 119 | model_type="deepseek_v4", num_nextn_predict_layers=3, dspark_block_size=5 |
| 120 | ) |
| 121 | config = _make_speculative_config( |
| 122 | hf_kwargs, "DeepSeekV4MTPModel", method="mtp", num_speculative_tokens=6 |
| 123 | ) |
| 124 | |
| 125 | assert config.num_speculative_tokens == 6 |
| 126 | assert config.draft_model_config.hf_config.n_predict == 3 |
| 127 | assert config.draft_model_config.hf_config.architectures == ["DeepSeekV4MTPModel"] |
| 128 | |
| 129 | with pytest.raises(ValueError, match="must be divisible by n_predict=3"): |
| 130 | _make_speculative_config( |
| 131 | hf_kwargs, "DeepSeekV4MTPModel", method="mtp", num_speculative_tokens=5 |
| 132 | ) |
| 133 | |
| 134 | |
| 135 | @pytest.mark.cpu_test |
| 136 | def test_callable_overrides_reach_the_draft_config(): |
| 137 | """A callable override (config-to-config transform) composes with the |
| 138 | architecture-mapping override and is applied to the draft config.""" |
| 139 | |
| 140 | def shrink(hf_config: PreTrainedConfig) -> PreTrainedConfig: |
| 141 | hf_config.num_hidden_layers = 1 |
| 142 | return hf_config |
| 143 | |
| 144 | composed = SpeculativeConfig.compose_draft_hf_overrides(shrink) |
| 145 | assert composed is not SpeculativeConfig.hf_config_override |
| 146 | |
| 147 | out = composed(_make_hf_config()) |
| 148 | # The shrink transform must have been applied to the draft config. |
| 149 | assert out.num_hidden_layers == 1 |
| 150 | |
| 151 | |
| 152 | @pytest.mark.cpu_test |
| 153 | def test_arch_mapping_applies_before_callable_override(): |
| 154 | """The static arch-mapping override runs first, so the user callable |
| 155 | observes (and may adjust) the post-mapping config.""" |
| 156 | seen_architectures: list[str] = [] |
| 157 | |
| 158 | def record(hf_config: PreTrainedConfig) -> PreTrainedConfig: |
| 159 | seen_architectures.append(hf_config.architectures[0]) |
| 160 | return hf_config |
| 161 | |
| 162 | composed = SpeculativeConfig.compose_draft_hf_overrides(record) |
| 163 | |
| 164 | # MiMo is one of the arch-mapped model types: hf_config_override |
| 165 | # rewrites architectures to ["MiMoMTPModel"]. |
| 166 | mimo = _make_hf_config( |
| 167 | architectures=["MiMoForCausalLM"], |
| 168 | model_type="mimo", |
| 169 | num_nextn_predict_layers=1, |
| 170 | ) |
| 171 | composed(mimo) |
| 172 | assert seen_architectures == ["MiMoMTPModel"] |
| 173 | |
| 174 | |
| 175 | @pytest.mark.cpu_test |
| 176 | def test_inkling_override_exposes_all_mtp_depths(): |
| 177 | text_config = _make_hf_config( |
| 178 | architectures=["InklingForCausalLM"], |
| 179 | model_type="inkling_model", |
| 180 | local_layer_ids=[1, 3], |
| 181 | ) |
| 182 | config = _make_hf_config( |
| 183 | architectures=["InklingForConditionalGeneration"], |
| 184 | model_type="inkling_mm_model", |
| 185 | text_config=text_config, |
| 186 | mtp_config={ |
| 187 | "num_nextn_predict_layers": 8, |
| 188 | "local_layer_ids": [0, 2, 4], |
| 189 | }, |
| 190 | ) |
| 191 | |
| 192 | out = SpeculativeConfig.hf_config_override(config) |
| 193 | |
| 194 | assert out is text_config |
| 195 | assert out.model_type == "inkling_mtp" |
| 196 | assert out.architectures == ["InklingMTPModel"] |
| 197 | # Multi-module MTP: every checkpoint depth is exposed (module i drafts |
| 198 | # speculative token i), no longer clamped to the first depth. |
| 199 | assert out.n_predict == 8 |
| 200 | assert out.num_nextn_predict_layers == 8 |
| 201 | assert out.chain_hidden_post_norm is False |
| 202 | assert out.local_layer_ids == [0, 2, 4] |
| 203 | |
| 204 | |
| 205 | def _module_level_shrink(hf_config: PreTrainedConfig) -> PreTrainedConfig: |
| 206 | hf_config.num_hidden_layers = 1 |
| 207 | return hf_config |
| 208 | |
| 209 | |
| 210 | @pytest.mark.cpu_test |
| 211 | def test_composed_override_is_picklable(): |
| 212 | """The draft ``ModelConfig`` is sent to spawned engine-core processes, so |
| 213 | the composed override must be picklable. A nested local closure is not |
| 214 | (it raised ``Can't get local object`` on DFlashDraftModel); a |
| 215 | ``functools.partial`` over a module-referenceable static method is. |
| 216 | Guard against regressing to a closure.""" |
| 217 | composed = SpeculativeConfig.compose_draft_hf_overrides(_module_level_shrink) |
| 218 | |
| 219 | assert isinstance(composed, functools.partial) |
| 220 | assert composed.func is SpeculativeConfig._apply_composed_hf_override |
| 221 | |
| 222 | out = composed(_make_hf_config()) |
| 223 | assert out.num_hidden_layers == 1 |
| 224 | |
| 225 | |
| 226 | def _make_mtp_speculative_config( |
| 227 | override: bool | None, |
| 228 | checkpoint_value: bool, |
| 229 | ) -> SpeculativeConfig: |
| 230 | draft_hf_config = _make_hf_config( |
| 231 | architectures=["Qwen4ExpMTP"], |
| 232 | model_type="qwen4_exp_mtp", |
| 233 | n_predict=1, |
| 234 | index_share_for_mtp_iteration=checkpoint_value, |
| 235 | ) |
| 236 | draft_model_config = MagicMock( |
| 237 | model="draft", |
| 238 | hf_config=draft_hf_config, |
| 239 | architectures=draft_hf_config.architectures, |
| 240 | max_model_len=128, |
| 241 | ) |
| 242 | target_model_config = MagicMock( |
| 243 | model="target", |
| 244 | max_model_len=128, |
| 245 | quantization=None, |
| 246 | hf_overrides={}, |
| 247 | ) |
| 248 | |
| 249 | with patch("vllm.config.speculative.ModelConfig", return_value=draft_model_config): |
| 250 | return SpeculativeConfig( |
| 251 | model="draft", |
| 252 | method="mtp", |
| 253 | num_speculative_tokens=1, |
| 254 | index_share_for_mtp_iteration=override, |
| 255 | target_model_config=target_model_config, |
| 256 | target_parallel_config=ParallelConfig(), |
| 257 | ) |
| 258 | |
| 259 | |
| 260 | @pytest.mark.cpu_test |
| 261 | @pytest.mark.parametrize( |
| 262 | ("override", "checkpoint_value", "expected"), |
| 263 | [(None, True, True), (False, True, False), (True, False, True)], |
| 264 | ) |
| 265 | def test_mtp_index_share_override( |
| 266 | override: bool | None, checkpoint_value: bool, expected: bool |
| 267 | ): |
| 268 | speculative_config = _make_mtp_speculative_config(override, checkpoint_value) |
| 269 | assert ( |
| 270 | speculative_config.draft_model_config.hf_config.index_share_for_mtp_iteration |
| 271 | is expected |
| 272 | ) |