@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2f55ef254c

[Model] Add Qwen3-Omni DSpark support (#52560)

Signed-off-by: Zhou248 <630563665@qq.com>

Zhou248 committed Aug 23, 2026 at 06:35 UTC 2f55ef254c70e110d637beeedf48238977ebb683
11 files changed +731 -22
tests/model_executor/test_qwen3_omni.py
+171 -1
@@ -1,9 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from unittest.mock import Mock
4 +from contextlib import nullcontext
5 +from types import SimpleNamespace
6 +from unittest.mock import Mock, patch
7
8 import pytest
9 +import torch
10 +import torch.nn as nn
11 from transformers import PretrainedConfig
12
13 from vllm.multimodal.processing import InputProcessingContext
@@ -218,5 +222,171 @@ def test_qwen3_omni_get_updates_use_audio_in_video(
222 )
223
224
225 +@pytest.mark.skip_global_cleanup
226 +def test_qwen3_omni_exposes_eagle3_to_its_text_backbone():
227 + from vllm.model_executor.models.interfaces import EagleModelMixin, supports_eagle3
228 + from vllm.model_executor.models.qwen3_omni_moe_thinker import (
229 + Qwen3OmniMoeThinkerForConditionalGeneration,
230 + )
231 +
232 + class DummyBackbone(nn.Module, EagleModelMixin):
233 + def __init__(self):
234 + super().__init__()
235 + self.layers = nn.ModuleList([nn.Identity(), nn.Identity()])
236 +
237 + class DummyLanguageModel(nn.Module):
238 + def __init__(self):
239 + super().__init__()
240 + self.model = DummyBackbone()
241 +
242 + def embed_input_ids(self, input_ids):
243 + return input_ids
244 +
245 + model = Qwen3OmniMoeThinkerForConditionalGeneration.__new__(
246 + Qwen3OmniMoeThinkerForConditionalGeneration
247 + )
248 + nn.Module.__init__(model)
249 + model.language_model = DummyLanguageModel()
250 +
251 + assert supports_eagle3(model)
252 + model.set_aux_hidden_state_layers((1, 2))
253 + assert model.language_model.model.aux_hidden_state_layers == (1, 2)
254 +
255 +
256 +@pytest.mark.skip_global_cleanup
257 +def test_qwen3_omni_text_model_collects_post_deepstack_aux_hidden_states():
258 + from vllm.model_executor.models.qwen3_omni_moe_thinker import Qwen3MoeLLMModel
259 +
260 + class DummyLayer(nn.Module):
261 + def forward(self, positions, hidden_states, residual):
262 + return hidden_states + 1, torch.full_like(hidden_states, 10)
263 +
264 + class DummyNorm(nn.Module):
265 + def forward(self, hidden_states, residual):
266 + return hidden_states + residual, None
267 +
268 + model = Qwen3MoeLLMModel.__new__(Qwen3MoeLLMModel)
269 + nn.Module.__init__(model)
270 + model.start_layer = 0
271 + model.end_layer = 1
272 + model.layers = nn.ModuleList([DummyLayer()])
273 + model.norm = DummyNorm()
274 + model.aux_hidden_state_layers = (1,)
275 +
276 + pp_group = Mock(is_first_rank=True, is_last_rank=True)
277 + inputs_embeds = torch.tensor([[1.0]])
278 + deepstack_inputs = {"deepstack_input_embeds_0": torch.tensor([[3.0]])}
279 + with patch(
280 + "vllm.model_executor.models.qwen3_omni_moe_thinker.get_pp_group",
281 + return_value=pp_group,
282 + ):
283 + output, aux_hidden_states = model.forward(
284 + input_ids=None,
285 + positions=torch.tensor([0]),
286 + inputs_embeds=inputs_embeds,
287 + deepstack_input_embeds=deepstack_inputs,
288 + )
289 +
290 + torch.testing.assert_close(output, torch.tensor([[15.0]]))
291 + assert len(aux_hidden_states) == 1
292 + torch.testing.assert_close(aux_hidden_states[0], torch.tensor([[15.0]]))
293 +
294 +
295 +@pytest.mark.skip_global_cleanup
296 +@pytest.mark.parametrize(
297 + ("input_vocab_size", "draft_vocab_size", "weights", "error"),
298 + [
299 + (101, 100, [], "must include embed_tokens weights"),
300 + (99, 99, [], "must include lm_head weights"),
301 + (100, 40, [], "must include lm_head weights"),
302 + (
303 + 100,
304 + 40,
305 + [("lm_head.weight", torch.empty(40, 8))],
306 + "must include a d2t mapping",
307 + ),
308 + ],
309 +)
310 +def test_qwen3_dspark_rejects_incomplete_vocab_weights(
311 + input_vocab_size, draft_vocab_size, weights, error
312 +):
313 + from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM
314 +
315 + model = Qwen3DSparkForCausalLM.__new__(Qwen3DSparkForCausalLM)
316 + nn.Module.__init__(model)
317 + object.__setattr__(
318 + model,
319 + "config",
320 + SimpleNamespace(
321 + vocab_size=input_vocab_size,
322 + draft_vocab_size=draft_vocab_size,
323 + ),
324 + )
325 + object.__setattr__(model, "target_vocab_size", 100)
326 +
327 + with pytest.raises(ValueError, match=error):
328 + model.load_weights(weights)
329 +
330 +
331 +@pytest.mark.skip_global_cleanup
332 +def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary():
333 + from vllm.v1.worker.gpu.spec_decode.dspark import utils as dspark_utils
334 +
335 + target_embedding = nn.Embedding(100, 8)
336 + draft_embedding = nn.Embedding(99, 8)
337 + target_model = SimpleNamespace(model=SimpleNamespace(embed_tokens=target_embedding))
338 + draft_model = SimpleNamespace(
339 + model=SimpleNamespace(embed_tokens=draft_embedding),
340 + has_own_embed_tokens=False,
341 + )
342 + draft_model_config = SimpleNamespace(
343 + hf_config=SimpleNamespace(model_type="qwen3"),
344 + get_vocab_size=Mock(return_value=99),
345 + )
346 + vllm_config = SimpleNamespace(
347 + speculative_config=SimpleNamespace(
348 + draft_model_config=draft_model_config,
349 + attention_backend=None,
350 + kv_cache_dtype=None,
351 + ),
352 + attention_config=SimpleNamespace(backend=None),
353 + cache_config=SimpleNamespace(),
354 + model_config=SimpleNamespace(get_vocab_size=Mock(return_value=100)),
355 + )
356 +
357 + def fake_replace(config, **changes):
358 + values = vars(config).copy()
359 + values.update(changes)
360 + return SimpleNamespace(**values)
361 +
362 + with (
363 + patch.object(dspark_utils, "replace", side_effect=fake_replace),
364 + patch.object(
365 + dspark_utils,
366 + "get_pp_group",
367 + return_value=SimpleNamespace(world_size=1),
368 + ),
369 + patch(
370 + "vllm.compilation.backends.set_model_tag",
371 + return_value=nullcontext(),
372 + ),
373 + patch(
374 + "vllm.model_executor.model_loader.get_model",
375 + return_value=draft_model,
376 + ),
377 + patch(
378 + "vllm.model_executor.models.qwen3_dflash.dflash_has_any_non_causal",
379 + return_value=False,
380 + ),
381 + patch(
382 + "vllm.model_executor.models.utils.get_draft_quant_config",
383 + return_value=None,
384 + ),
385 + ):
386 + loaded_model = dspark_utils.load_dspark_model(target_model, vllm_config)
387 +
388 + assert loaded_model.model.embed_tokens is target_embedding
389 +
390 +
391 if __name__ == "__main__":
392 pytest.main([__file__, "-v"])
tests/models/registry.py
+7
@@ -1557,6 +1557,13 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1557 is_available_online=False,
1558 use_original_num_layers=True, # DSpark backbone requires all layers
1559 ),
1560 + "Qwen3OmniDSparkModel": _HfExamplesInfo(
1561 + "Qwen/Qwen3-Omni-30B-A3B-Instruct",
1562 + speculative_model="Qwen/Qwen3-Omni-30B-A3B-DSpark",
1563 + is_available_online=False,
1564 + use_original_num_layers=True,
1565 + min_transformers_version="4.57",
1566 + ),
1567 "Gemma4DSparkModel": _HfExamplesInfo(
1568 "google/gemma-4-12B-it",
1569 speculative_model="deepseek-ai/dspark_gemma4_12b_block7",
tests/test_config.py
+121
@@ -32,6 +32,7 @@ from vllm.config.compilation import CompilationMode, CUDAGraphMode
32 from vllm.config.kernel import IrOpPriorityConfig
33 from vllm.config.load import LoadConfig
34 from vllm.config.mamba import MambaBackendEnum
35 +from vllm.config.speculative import _validate_qwen3_omni_dspark
36 from vllm.config.utils import get_field
37 from vllm.config.vllm import OPTIMIZATION_LEVEL_TO_CONFIG, OptimizationLevel
38 from vllm.platforms import current_platform
@@ -2100,6 +2101,126 @@ def test_mtp_draft_uses_model_weights_not_local_cache(mock_model_config_cls):
2101 assert actual_model == s3_url
2102
2103
2104 +def _make_qwen3_omni_dspark_configs():
2105 + text_config = SimpleNamespace(
2106 + hidden_size=2048,
2107 + num_hidden_layers=48,
2108 + num_attention_heads=32,
2109 + num_key_value_heads=4,
2110 + head_dim=128,
2111 + vocab_size=152064,
2112 + )
2113 + target_model_config = SimpleNamespace(
2114 + hf_config=SimpleNamespace(
2115 + model_type="qwen3_omni_moe",
2116 + architectures=["Qwen3OmniMoeForConditionalGeneration"],
2117 + thinker_config=SimpleNamespace(text_config=text_config),
2118 + ),
2119 + hf_text_config=text_config,
2120 + architectures=["Qwen3OmniMoeForConditionalGeneration"],
2121 + get_hidden_size=lambda: 2048,
2122 + get_total_num_hidden_layers=lambda: 48,
2123 + get_vocab_size=lambda: 152064,
2124 + )
2125 + draft_hf_config = SimpleNamespace(
2126 + model_type="qwen3",
2127 + architectures=["Qwen3OmniDSparkModel"],
2128 + hidden_size=2048,
2129 + num_attention_heads=32,
2130 + num_key_value_heads=4,
2131 + head_dim=128,
2132 + block_size=7,
2133 + target_hidden_size=2048,
2134 + target_layer_ids=[1, 9, 17, 25, 33],
2135 + use_aux_hidden_state=True,
2136 + markov_rank=64,
2137 + markov_head_type="vanilla",
2138 + sample_from_anchor=True,
2139 + dspark_bonus_anchor=False,
2140 + vocab_size=152064,
2141 + draft_vocab_size=32000,
2142 + mask_token_id=151669,
2143 + rope_parameters={"rope_type": "default", "rope_theta": 1000000.0},
2144 + )
2145 + draft_model_config = SimpleNamespace(
2146 + hf_config=draft_hf_config,
2147 + architectures=["Qwen3OmniDSparkModel"],
2148 + )
2149 + return target_model_config, draft_model_config
2150 +
2151 +
2152 +@pytest.mark.skip_global_cleanup
2153 +def test_qwen3_omni_dspark_checkpoint_contract_is_accepted():
2154 + target_config, draft_config = _make_qwen3_omni_dspark_configs()
2155 + _validate_qwen3_omni_dspark(target_config, draft_config, 7)
2156 +
2157 +
2158 +@pytest.mark.skip_global_cleanup
2159 +def test_qwen3_omni_dspark_rejects_generic_qwen3_architecture():
2160 + target_config, draft_config = _make_qwen3_omni_dspark_configs()
2161 + draft_config.hf_config.architectures = ["Qwen3DSparkModel"]
2162 + draft_config.architectures = ["Qwen3DSparkModel"]
2163 +
2164 + with pytest.raises(ValueError, match="must be converted first"):
2165 + _validate_qwen3_omni_dspark(target_config, draft_config, 7)
2166 +
2167 +
2168 +@pytest.mark.parametrize(
2169 + ("field", "value", "error"),
2170 + [
2171 + ("block_size", 5, "trained block_size"),
2172 + ("target_hidden_size", 4096, "target_hidden_size"),
2173 + ("hidden_size", 4096, "draft hidden_size"),
2174 + ("num_attention_heads", 16, "num_attention_heads"),
2175 + ("num_key_value_heads", 8, "num_key_value_heads"),
2176 + ("head_dim", 64, "head_dim"),
2177 + ("target_layer_ids", [7, 48], "zero-based text-layer"),
2178 + ("target_layer_ids", [23, 7], "strictly increasing"),
2179 + ("use_aux_hidden_state", False, "use_aux_hidden_state=true"),
2180 + ("markov_rank", 0, "markov_rank"),
2181 + ("markov_head_type", "gated", "markov_head_type='vanilla'"),
2182 + ("sample_from_anchor", False, "sample_from_anchor=true"),
2183 + ("dspark_bonus_anchor", True, "dspark_bonus_anchor=false"),
2184 + ("vocab_size", 0, "input vocab_size must be a positive integer"),
2185 + ],
2186 +)
2187 +@pytest.mark.skip_global_cleanup
2188 +def test_qwen3_omni_dspark_rejects_incompatible_checkpoint_fields(field, value, error):
2189 + target_config, draft_config = _make_qwen3_omni_dspark_configs()
2190 + setattr(draft_config.hf_config, field, value)
2191 +
2192 + with pytest.raises(ValueError, match=error):
2193 + _validate_qwen3_omni_dspark(target_config, draft_config, 7)
2194 +
2195 +
2196 +@pytest.mark.skip_global_cleanup
2197 +def test_qwen3_omni_dspark_rejects_mrope_draft_positions():
2198 + target_config, draft_config = _make_qwen3_omni_dspark_configs()
2199 + draft_config.hf_config.rope_parameters["mrope_section"] = [24, 20, 20]
2200 +
2201 + with pytest.raises(ValueError, match="logical 1-D RoPE"):
2202 + _validate_qwen3_omni_dspark(target_config, draft_config, 7)
2203 +
2204 +
2205 +@pytest.mark.skip_global_cleanup
2206 +def test_qwen3_omni_dspark_allows_draft_only_noise_token_row():
2207 + target_config, draft_config = _make_qwen3_omni_dspark_configs()
2208 + draft_config.hf_config.vocab_size = 152065
2209 + draft_config.hf_config.draft_vocab_size = 152064
2210 + draft_config.hf_config.mask_token_id = 152064
2211 +
2212 + _validate_qwen3_omni_dspark(target_config, draft_config, 7)
2213 +
2214 +
2215 +@pytest.mark.skip_global_cleanup
2216 +def test_qwen3_omni_dspark_allows_smaller_input_vocabulary():
2217 + target_config, draft_config = _make_qwen3_omni_dspark_configs()
2218 + draft_config.hf_config.vocab_size = 151936
2219 + draft_config.hf_config.mask_token_id = 151669
2220 +
2221 + _validate_qwen3_omni_dspark(target_config, draft_config, 7)
2222 +
2223 +
2224 def test_ir_op_priority_default():
2225 """Test that IR op priority defaults are set correctly."""
2226 from vllm.config.kernel import IrOpPriorityConfig
tests/transformers_utils/test_speculators_dspark_config.py new
+58
@@ -0,0 +1,58 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from vllm.transformers_utils.configs.speculators.base import SpeculatorsConfig
7 +
8 +pytestmark = pytest.mark.skip_global_cleanup
9 +
10 +
11 +def _make_dspark_config(architecture: str) -> dict:
12 + return {
13 + "speculators_model_type": "dspark",
14 + "architectures": [architecture],
15 + "sample_from_anchor": True,
16 + "use_aux_hidden_state": True,
17 + "draft_vocab_size": 32,
18 + "target_hidden_size": 16,
19 + "mask_token_id": 31,
20 + "markov_rank": 4,
21 + "markov_head_type": "vanilla",
22 + "block_size": 3,
23 + "aux_hidden_state_layer_ids": [2, 5],
24 + "transformer_layer_config": {
25 + "model_type": "qwen3",
26 + "hidden_size": 16,
27 + },
28 + }
29 +
30 +
31 +def test_dspark_updater_preserves_qwen3_omni_contract() -> None:
32 + config = SpeculatorsConfig.extract_transformers_pre_trained_config(
33 + _make_dspark_config("Qwen3OmniDSparkModel")
34 + )
35 +
36 + assert config["architectures"] == ["Qwen3OmniDSparkModel"]
37 + assert config["sample_from_anchor"] is True
38 + assert config["dspark_bonus_anchor"] is False
39 + assert config["use_aux_hidden_state"] is True
40 + assert config["target_layer_ids"] == [1, 4]
41 +
42 +
43 +def test_dspark_updater_keeps_legacy_qwen3_checkpoint_loadable() -> None:
44 + config = SpeculatorsConfig.extract_transformers_pre_trained_config(
45 + _make_dspark_config("DSparkSpeculator")
46 + )
47 +
48 + assert config["architectures"] == ["Qwen3DSparkModel"]
49 +
50 +
51 +def test_dspark_updater_maps_bonus_anchor_semantics() -> None:
52 + outer_config = _make_dspark_config("Qwen3DSparkModel")
53 + outer_config["sample_from_anchor"] = False
54 +
55 + config = SpeculatorsConfig.extract_transformers_pre_trained_config(outer_config)
56 +
57 + assert config["sample_from_anchor"] is False
58 + assert config["dspark_bonus_anchor"] is True
vllm/config/speculative.py
+302 -2
@@ -3,7 +3,7 @@
3
4 import copy
5 import functools
6 -from collections.abc import Callable
6 +from collections.abc import Callable, Mapping
7 from typing import TYPE_CHECKING, Any, Literal, get_args
8
9 from pydantic import Field, SkipValidation, field_validator, model_validator
@@ -80,6 +80,293 @@ SpeculativeMethod = Literal[
80 RejectionSampleMethod = Literal["standard", "synthetic", "block"]
81 DraftSampleMethod = Literal["greedy", "probabilistic"]
82
83 +_QWEN3_OMNI_TARGET_ARCHITECTURES = frozenset(
84 + {
85 + "Qwen3OmniMoeForConditionalGeneration",
86 + "Qwen3OmniMoeThinkerForConditionalGeneration",
87 + }
88 +)
89 +_QWEN3_OMNI_DSPARK_ARCHITECTURE = "Qwen3OmniDSparkModel"
90 +
91 +
92 +def _is_qwen3_omni_target(model_config: ModelConfig) -> bool:
93 + hf_config = model_config.hf_config
94 + architectures = set(getattr(model_config, "architectures", ()) or ())
95 + architectures.update(getattr(hf_config, "architectures", ()) or ())
96 + return getattr(hf_config, "model_type", None) == "qwen3_omni_moe" or bool(
97 + architectures & _QWEN3_OMNI_TARGET_ARCHITECTURES
98 + )
99 +
100 +
101 +def _get_qwen3_omni_text_config(model_config: ModelConfig) -> Any:
102 + thinker_config = getattr(model_config.hf_config, "thinker_config", None)
103 + text_config = getattr(thinker_config, "text_config", None)
104 + if text_config is None:
105 + text_config = getattr(model_config, "hf_text_config", None)
106 + return text_config
107 +
108 +
109 +def _get_attention_head_dim(config: Any) -> int | None:
110 + head_dim = getattr(config, "head_dim", None)
111 + if head_dim is not None:
112 + return head_dim
113 + hidden_size = getattr(config, "hidden_size", None)
114 + num_attention_heads = getattr(config, "num_attention_heads", None)
115 + if (
116 + isinstance(hidden_size, int)
117 + and isinstance(num_attention_heads, int)
118 + and num_attention_heads > 0
119 + and hidden_size % num_attention_heads == 0
120 + ):
121 + return hidden_size // num_attention_heads
122 + return None
123 +
124 +
125 +def _get_nested_config_value(config: Any, section: str, name: str) -> Any:
126 + nested_config = getattr(config, section, None)
127 + if isinstance(nested_config, Mapping):
128 + return nested_config.get(name)
129 + return getattr(nested_config, name, None)
130 +
131 +
132 +def _get_qwen3_dspark_value(config: Any, name: str) -> Any:
133 + # Match DFlashQwen3Model's precedence: dflash_config overrides
134 + # eagle_config, and top-level values are the modern fallback.
135 + value = _get_nested_config_value(config, "dflash_config", name)
136 + if value is None:
137 + value = _get_nested_config_value(config, "eagle_config", name)
138 + if value is None:
139 + value = getattr(config, name, None)
140 + return value
141 +
142 +
143 +def _validate_qwen3_omni_dspark(
144 + target_model_config: ModelConfig,
145 + draft_model_config: ModelConfig,
146 + num_speculative_tokens: int,
147 +) -> None:
148 + """Validate the standalone Qwen3-Omni DSpark checkpoint contract."""
149 + if not _is_qwen3_omni_target(target_model_config):
150 + return
151 +
152 + draft_hf_config = draft_model_config.hf_config
153 + draft_architectures = set(getattr(draft_model_config, "architectures", ()) or ())
154 + draft_architectures.update(getattr(draft_hf_config, "architectures", ()) or ())
155 + if _QWEN3_OMNI_DSPARK_ARCHITECTURE not in draft_architectures:
156 + raise ValueError(
157 + "Qwen3-Omni DSpark requires a standalone draft checkpoint with "
158 + f"architectures=['{_QWEN3_OMNI_DSPARK_ARCHITECTURE}']; generic "
159 + "Qwen3DSparkModel checkpoints must be converted first."
160 + )
161 +
162 + block_size = _get_qwen3_dspark_value(draft_hf_config, "block_size")
163 + if block_size is None:
164 + block_size = _get_qwen3_dspark_value(draft_hf_config, "dspark_block_size")
165 + if (
166 + not isinstance(block_size, int)
167 + or isinstance(block_size, bool)
168 + or block_size <= 0
169 + ):
170 + raise ValueError(
171 + "Qwen3-Omni DSpark requires a positive integer block_size in the "
172 + "draft config."
173 + )
174 + if num_speculative_tokens != block_size:
175 + raise ValueError(
176 + "Qwen3-Omni DSpark requires num_speculative_tokens to match the "
177 + f"trained block_size ({block_size}); got {num_speculative_tokens}."
178 + )
179 +
180 + target_text_config = _get_qwen3_omni_text_config(target_model_config)
181 + if target_text_config is None:
182 + raise ValueError(
183 + "Qwen3-Omni DSpark could not resolve the target thinker text_config."
184 + )
185 + target_hidden_size = getattr(
186 + target_text_config,
187 + "hidden_size",
188 + target_model_config.get_hidden_size(),
189 + )
190 + draft_target_hidden_size = getattr(draft_hf_config, "target_hidden_size", None)
191 + if draft_target_hidden_size != target_hidden_size:
192 + raise ValueError(
193 + "Qwen3-Omni DSpark draft target_hidden_size must match the target "
194 + f"text hidden size ({target_hidden_size}); got {draft_target_hidden_size}."
195 + )
196 + draft_hidden_size = getattr(draft_hf_config, "hidden_size", None)
197 + if draft_hidden_size != target_hidden_size:
198 + raise ValueError(
199 + "Qwen3-Omni DSpark draft hidden_size must match the target thinker "
200 + f"hidden size ({target_hidden_size}); got {draft_hidden_size}."
201 + )
202 +
203 + target_attention = {
204 + "num_attention_heads": getattr(target_text_config, "num_attention_heads", None),
205 + "num_key_value_heads": getattr(
206 + target_text_config,
207 + "num_key_value_heads",
208 + getattr(target_text_config, "num_attention_heads", None),
209 + ),
210 + "head_dim": _get_attention_head_dim(target_text_config),
211 + }
212 + draft_attention = {
213 + "num_attention_heads": getattr(draft_hf_config, "num_attention_heads", None),
214 + "num_key_value_heads": getattr(
215 + draft_hf_config,
216 + "num_key_value_heads",
217 + getattr(draft_hf_config, "num_attention_heads", None),
218 + ),
219 + "head_dim": _get_attention_head_dim(draft_hf_config),
220 + }
221 + for name, target_value in target_attention.items():
222 + if (
223 + not isinstance(target_value, int)
224 + or isinstance(target_value, bool)
225 + or target_value <= 0
226 + ):
227 + raise ValueError(
228 + "Qwen3-Omni DSpark requires a positive integer "
229 + f"{name} in the target thinker text_config; got {target_value}."
230 + )
231 + draft_value = draft_attention[name]
232 + if draft_value != target_value:
233 + raise ValueError(
234 + f"Qwen3-Omni DSpark draft {name} must match the target thinker "
235 + f"value ({target_value}); got {draft_value}."
236 + )
237 +
238 + target_layer_ids = _get_nested_config_value(
239 + draft_hf_config, "dflash_config", "target_layer_ids"
240 + )
241 + if target_layer_ids is None:
242 + target_layer_ids = getattr(draft_hf_config, "dspark_target_layer_ids", None)
243 + if target_layer_ids is None:
244 + target_layer_ids = getattr(draft_hf_config, "target_layer_ids", None)
245 + if not isinstance(target_layer_ids, (list, tuple)) or not target_layer_ids:
246 + raise ValueError(
247 + "Qwen3-Omni DSpark requires a non-empty target_layer_ids list in "
248 + "the draft config."
249 + )
250 + if any(
251 + not isinstance(layer_id, int) or isinstance(layer_id, bool)
252 + for layer_id in target_layer_ids
253 + ):
254 + raise ValueError(
255 + "Qwen3-Omni DSpark target_layer_ids must contain only integers."
256 + )
257 + if list(target_layer_ids) != sorted(set(target_layer_ids)):
258 + raise ValueError(
259 + "Qwen3-Omni DSpark target_layer_ids must be unique and strictly increasing."
260 + )
261 + target_num_layers = getattr(
262 + target_text_config,
263 + "num_hidden_layers",
264 + target_model_config.get_total_num_hidden_layers(),
265 + )
266 + if target_layer_ids[0] < 0 or target_layer_ids[-1] >= target_num_layers:
267 + raise ValueError(
268 + "Qwen3-Omni DSpark target_layer_ids must be zero-based text-layer "
269 + f"indices in [0, {target_num_layers - 1}]; got {target_layer_ids}."
270 + )
271 +
272 + if _get_qwen3_dspark_value(draft_hf_config, "use_aux_hidden_state") is not True:
273 + raise ValueError(
274 + "Qwen3-Omni DSpark requires use_aux_hidden_state=true because Omni "
275 + "conditioning is supplied through thinker hidden states."
276 + )
277 +
278 + markov_rank = getattr(draft_hf_config, "markov_rank", None)
279 + if (
280 + not isinstance(markov_rank, int)
281 + or isinstance(markov_rank, bool)
282 + or markov_rank <= 0
283 + ):
284 + raise ValueError(
285 + "Qwen3-Omni DSpark requires a positive integer markov_rank in the "
286 + "draft config."
287 + )
288 + markov_head_type = getattr(draft_hf_config, "markov_head_type", None)
289 + if markov_head_type != "vanilla":
290 + raise ValueError(
291 + "Qwen3-Omni DSpark currently requires markov_head_type='vanilla'; "
292 + f"got {markov_head_type!r}."
293 + )
294 +
295 + sample_from_anchor = getattr(draft_hf_config, "sample_from_anchor", True)
296 + bonus_anchor = getattr(draft_hf_config, "dspark_bonus_anchor", False)
297 + if sample_from_anchor is not True or bonus_anchor is not False:
298 + raise ValueError(
299 + "Qwen3-Omni DSpark requires sample_from_anchor=true and "
300 + "dspark_bonus_anchor=false."
301 + )
302 +
303 + target_vocab_size = getattr(
304 + target_text_config,
305 + "vocab_size",
306 + target_model_config.get_vocab_size(),
307 + )
308 + draft_input_vocab_size = getattr(draft_hf_config, "vocab_size", None)
309 + if (
310 + not isinstance(draft_input_vocab_size, int)
311 + or isinstance(draft_input_vocab_size, bool)
312 + or draft_input_vocab_size <= 0
313 + ):
314 + raise ValueError(
315 + "Qwen3-Omni DSpark input vocab_size must be a positive integer; "
316 + f"got {draft_input_vocab_size}."
317 + )
318 + draft_output_vocab_size = getattr(draft_hf_config, "draft_vocab_size", None)
319 + if draft_output_vocab_size is None:
320 + draft_output_vocab_size = draft_input_vocab_size
321 + if (
322 + not isinstance(draft_output_vocab_size, int)
323 + or isinstance(draft_output_vocab_size, bool)
324 + or not 0 < draft_output_vocab_size <= target_vocab_size
325 + ):
326 + raise ValueError(
327 + "Qwen3-Omni DSpark draft_vocab_size must be a positive integer no "
328 + f"larger than target vocab_size ({target_vocab_size}); got "
329 + f"{draft_output_vocab_size}."
330 + )
331 +
332 + noise_token_id = _get_nested_config_value(
333 + draft_hf_config, "dflash_config", "mask_token_id"
334 + )
335 + for name in (
336 + "mask_token_id",
337 + "dspark_noise_token_id",
338 + "pard_token",
339 + "ptd_token_id",
340 + ):
341 + if noise_token_id is None:
342 + noise_token_id = getattr(draft_hf_config, name, None)
343 + if noise_token_id is not None:
344 + break
345 + if (
346 + not isinstance(noise_token_id, int)
347 + or isinstance(noise_token_id, bool)
348 + or not 0 <= noise_token_id < draft_input_vocab_size
349 + ):
350 + raise ValueError(
351 + "Qwen3-Omni DSpark requires a valid mask/noise token id within the "
352 + f"draft input vocabulary [0, {draft_input_vocab_size - 1}]."
353 + )
354 +
355 + rope_configs = (
356 + getattr(draft_hf_config, "rope_parameters", None),
357 + getattr(draft_hf_config, "rope_scaling", None),
358 + )
359 + has_mrope = getattr(draft_hf_config, "mrope_section", None) is not None
360 + has_mrope = has_mrope or any(
361 + isinstance(rope_config, Mapping) and "mrope_section" in rope_config
362 + for rope_config in rope_configs
363 + )
364 + if has_mrope:
365 + raise ValueError(
366 + "Qwen3-Omni DSpark draft checkpoints must use logical 1-D RoPE and "
367 + "must not define mrope_section."
368 + )
369 +
370
371 @config
372 class SpeculativeConfig:
@@ -964,6 +1251,8 @@ class SpeculativeConfig:
1251 elif (
1252 "dspark" in self.draft_model_config.model.lower()
1253 or "Qwen3DSparkModel" in self.draft_model_config.architectures
1254 + or _QWEN3_OMNI_DSPARK_ARCHITECTURE
1255 + in self.draft_model_config.architectures
1256 or "Gemma4DSparkModel" in self.draft_model_config.architectures
1257 or (
1258 "DSparkDraftModel" in self.draft_model_config.architectures
@@ -1037,6 +1326,8 @@ class SpeculativeConfig:
1326 self.update_arch_()
1327 elif self.method == "dspark" and (
1328 "Qwen3DSparkModel" not in self.draft_model_config.architectures
1329 + and _QWEN3_OMNI_DSPARK_ARCHITECTURE
1330 + not in self.draft_model_config.architectures
1331 and "Gemma4DSparkModel" not in self.draft_model_config.architectures
1332 and "K3DSparkModel" not in self.draft_model_config.architectures
1333 ):
@@ -1125,13 +1416,22 @@ class SpeculativeConfig:
1416 if (
1417 "Qwen3DSparkModel"
1418 not in self.draft_model_config.architectures
1419 + and _QWEN3_OMNI_DSPARK_ARCHITECTURE
1420 + not in self.draft_model_config.architectures
1421 ):
1422 raise ValueError(
1423 "dspark_draft_topk is only supported by "
1131 - "Qwen3DSparkModel"
1424 + "Qwen3DSparkModel and Qwen3OmniDSparkModel"
1425 )
1426 hf_config.dspark_draft_topk = dspark_draft_topk
1427
1428 + assert self.target_model_config is not None
1429 + _validate_qwen3_omni_dspark(
1430 + self.target_model_config,
1431 + self.draft_model_config,
1432 + self.num_speculative_tokens,
1433 + )
1434 +
1435 self.draft_tensor_parallel_size = (
1436 SpeculativeConfig._verify_and_get_draft_tp(
1437 self.target_parallel_config,
vllm/model_executor/models/qwen3_dflash.py
+8 -8
@@ -377,8 +377,6 @@ class DFlashQwen3DecoderLayer(nn.Module):
377
378 @support_torch_compile
379 class DFlashQwen3Model(nn.Module):
380 - decoder_layer_cls = DFlashQwen3DecoderLayer
381 -
380 hf_to_vllm_mapper = WeightsMapper(
381 orig_to_new_substr={
382 "midlayer.": "layers.0.",
@@ -412,10 +410,10 @@ class DFlashQwen3Model(nn.Module):
410 drafter_config = getattr(self.config, "eagle_config", {})
411 drafter_config.update(getattr(self.config, "dflash_config", {}))
412
415 - if drafter_config is not None and "use_aux_hidden_state" in drafter_config:
416 - self.use_aux_hidden_state = drafter_config["use_aux_hidden_state"]
417 - else:
418 - self.use_aux_hidden_state = True
413 + self.use_aux_hidden_state = drafter_config.get(
414 + "use_aux_hidden_state",
415 + getattr(self.config, "use_aux_hidden_state", True),
416 + )
417
418 current_vllm_config = get_current_vllm_config()
419
@@ -430,7 +428,9 @@ class DFlashQwen3Model(nn.Module):
428 # at that slot id. Some checkpoints (XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash) ship
429 # with a separate mask embedding tensor to use instead. When present, we load it
430 # and substitute it for embed_tokens[mask_token_id] when computing embeddings.
433 - self.mask_token_id = drafter_config.get("mask_token_id")
431 + self.mask_token_id = drafter_config.get(
432 + "mask_token_id", getattr(self.config, "mask_token_id", None)
433 + )
434 self.mask_embedding = nn.Parameter(
435 torch.zeros(self.config.hidden_size, dtype=vllm_config.model_config.dtype),
436 requires_grad=False,
@@ -439,7 +439,7 @@ class DFlashQwen3Model(nn.Module):
439
440 self.layers = nn.ModuleList(
441 [
442 - self.decoder_layer_cls(
442 + DFlashQwen3DecoderLayer(
443 current_vllm_config,
444 config=self.config,
445 layer_idx=layer_idx,
vllm/model_executor/models/qwen3_dspark.py
+20 -2
@@ -200,8 +200,8 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
200 self.logits_processor = LogitsProcessor(
201 self.config.draft_vocab_size, scale=logit_scale
202 )
203 - target_vocab_size = vllm_config.model_config.get_vocab_size()
204 - if self.config.draft_vocab_size != target_vocab_size:
203 + self.target_vocab_size = vllm_config.model_config.get_vocab_size()
204 + if self.config.draft_vocab_size != self.target_vocab_size:
205 self.draft_id_to_target_id = nn.Parameter(
206 torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
207 requires_grad=False,
@@ -280,6 +280,24 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
280 # mask_embedding is an unused placeholder param; DSpark masks via the vocab row.
281 # embed_tokens / lm_head are optional; when omitted they are shared from
282 # the target by load_dspark_model, so skip the unloaded params here.
283 + uses_expanded_input_vocab = self.config.vocab_size > self.target_vocab_size
284 + uses_reduced_vocab = self.config.draft_vocab_size != self.target_vocab_size
285 + if uses_expanded_input_vocab and not includes_embed_tokens:
286 + raise ValueError(
287 + "Qwen3 DSpark checkpoints whose input vocab_size is larger than "
288 + "the target vocabulary must include embed_tokens weights."
289 + )
290 + if uses_reduced_vocab and not includes_lm_head:
291 + raise ValueError(
292 + "Reduced-vocabulary Qwen3 DSpark checkpoints must include "
293 + "lm_head weights; the full target lm_head cannot be shared."
294 + )
295 + if uses_reduced_vocab and not includes_draft_id_mapping:
296 + raise ValueError(
297 + "Reduced-vocabulary Qwen3 DSpark checkpoints must include a "
298 + "d2t mapping so sampled draft ids can be converted to target ids."
299 + )
300 +
301 orig_to_new_substr = {"mask_embedding": None}
302 if not includes_embed_tokens:
303 orig_to_new_substr["embed_tokens"] = None
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+11 -1
@@ -83,6 +83,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum
83
84 from .interfaces import (
85 MultiModalEmbeddings,
86 + SupportsEagle3,
87 SupportsMRoPE,
88 SupportsMultiModal,
89 SupportsPP,
@@ -1067,7 +1068,7 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
1068 intermediate_tensors: IntermediateTensors | None = None,
1069 inputs_embeds: torch.Tensor | None = None,
1070 deepstack_input_embeds: IntermediateTensors | None = None,
1070 - ) -> torch.Tensor | IntermediateTensors:
1071 + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
1072 if get_pp_group().is_first_rank:
1073 if inputs_embeds is not None:
1074 hidden_states = inputs_embeds
@@ -1078,6 +1079,9 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
1079 assert intermediate_tensors is not None
1080 hidden_states = intermediate_tensors["hidden_states"]
1081 residual = intermediate_tensors["residual"]
1082 + aux_hidden_states = self._maybe_add_hidden_state(
1083 + [], self.start_layer, hidden_states, residual
1084 + )
1085 for layer_idx, layer in enumerate(
1086 self.layers[self.start_layer : self.end_layer]
1087 ):
@@ -1096,12 +1100,17 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
1100 hidden_states
1101 + deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"]
1102 )
1103 + self._maybe_add_hidden_state(
1104 + aux_hidden_states, layer_idx + 1, hidden_states, residual
1105 + )
1106
1107 if not get_pp_group().is_last_rank:
1108 return IntermediateTensors(
1109 {"hidden_states": hidden_states, "residual": residual}
1110 )
1111 hidden_states, _ = self.norm(hidden_states, residual)
1112 + if aux_hidden_states:
1113 + return hidden_states, aux_hidden_states
1114 return hidden_states
1115
1116
@@ -1577,6 +1586,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1586 SupportsMultiModal,
1587 SupportsPP,
1588 SupportsMRoPE,
1589 + SupportsEagle3,
1590 Qwen3OmniMoeConditionalGenerationMixin,
1591 SupportsTranscription,
1592 ):
vllm/model_executor/models/registry.py
+1
@@ -643,6 +643,7 @@ _SPECULATIVE_DECODING_MODELS = {
643 "DFlashMuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
644 "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"),
645 "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
646 + "Qwen3OmniDSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
647 "K3DSparkModel": (
648 "vllm.models.kimi_k3.nvidia.dspark_mla",
649 "K3DSparkForCausalLM",
vllm/transformers_utils/configs/speculators/algos.py
+17 -4
@@ -154,10 +154,22 @@ def update_dspark(config_dict: dict, pre_trained_config: dict) -> None:
154 False (anchor is a bonus token, only mask tokens predict, yielding
155 block_size - 1 speculative tokens).
156 """
157 - pre_trained_config["architectures"] = ["Qwen3DSparkModel"]
158 - pre_trained_config["sample_from_anchor"] = config_dict.get(
159 - "sample_from_anchor", False
160 - )
157 + architectures = config_dict.get("architectures") or []
158 + supported_architectures = {
159 + "Qwen3DSparkModel",
160 + "Qwen3OmniDSparkModel",
161 + }
162 + selected_architectures = [
163 + architecture
164 + for architecture in architectures
165 + if architecture in supported_architectures
166 + ]
167 + # Legacy msModelSpec checkpoints use a training-only architecture name.
168 + pre_trained_config["architectures"] = selected_architectures or ["Qwen3DSparkModel"]
169 +
170 + sample_from_anchor = config_dict.get("sample_from_anchor", False)
171 + pre_trained_config["sample_from_anchor"] = sample_from_anchor
172 + pre_trained_config["dspark_bonus_anchor"] = not sample_from_anchor
173
174 aux_layer_ids = config_dict["aux_hidden_state_layer_ids"]
175 pre_trained_config["eagle_aux_hidden_state_layer_ids"] = aux_layer_ids
@@ -173,6 +185,7 @@ def update_dspark(config_dict: dict, pre_trained_config: dict) -> None:
185 "block_size",
186 "enable_confidence_head",
187 "confidence_head_with_markov",
188 + "use_aux_hidden_state",
189 ):
190 if config_dict.get(key) is not None:
191 pre_trained_config[key] = config_dict[key]
vllm/v1/worker/gpu/spec_decode/dspark/utils.py
+15 -4
@@ -86,11 +86,16 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
86 )
87 target_inner = target_language_model.model
88 draft_inner = draft_model.model
89 + target_vocab_size = vllm_config.model_config.get_vocab_size()
90
91 target_embed = getattr(target_inner, "embed_tokens", None)
92 draft_embed = getattr(draft_inner, "embed_tokens", None)
92 - if target_embed is not None and _should_share(
93 - draft_model, "has_own_embed_tokens", draft_embed, target_embed
93 + if (
94 + target_embed is not None
95 + and draft_model_config.get_vocab_size() <= target_vocab_size
96 + and _should_share(
97 + draft_model, "has_own_embed_tokens", draft_embed, target_embed
98 + )
99 ):
100 if draft_embed is not None:
101 del draft_inner.embed_tokens
@@ -98,8 +103,14 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
103
104 target_lm_head = get_target_lm_head(target_model, target_language_model)
105 draft_lm_head = getattr(draft_model, "lm_head", None)
101 - if target_lm_head is not None and _should_share(
102 - draft_model, "has_own_lm_head", draft_lm_head, target_lm_head
106 + draft_output_vocab_size = (
107 + getattr(draft_model_config.hf_config, "draft_vocab_size", None)
108 + or draft_model_config.get_vocab_size()
109 + )
110 + if (
111 + target_lm_head is not None
112 + and draft_output_vocab_size == target_vocab_size
113 + and _should_share(draft_model, "has_own_lm_head", draft_lm_head, target_lm_head)
114 ):
115 if draft_lm_head is not None:
116 del draft_model.lm_head