iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Zhou248 <630563665@qq.com>
This commit is contained in:
Zhou248
2026-08-22 15:35:22 -07:00
committed by GitHub
parent a014e35f38
commit 2f55ef254c
11 changed files with 731 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/model_executor/test_qwen3_omni.py
+171 -1
View File
@@ -1,9 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import Mock
from contextlib import nullcontext
from types import SimpleNamespace
from unittest.mock import Mock, patch
import pytest
import torch
import torch.nn as nn
from transformers import PretrainedConfig
from vllm.multimodal.processing import InputProcessingContext
@@ -218,5 +222,171 @@ def test_qwen3_omni_get_updates_use_audio_in_video(
)
@pytest.mark.skip_global_cleanup
def test_qwen3_omni_exposes_eagle3_to_its_text_backbone():
from vllm.model_executor.models.interfaces import EagleModelMixin, supports_eagle3
from vllm.model_executor.models.qwen3_omni_moe_thinker import (
Qwen3OmniMoeThinkerForConditionalGeneration,
)
class DummyBackbone(nn.Module, EagleModelMixin):
def __init__(self):
super().__init__()
self.layers = nn.ModuleList([nn.Identity(), nn.Identity()])
class DummyLanguageModel(nn.Module):
def __init__(self):
super().__init__()
self.model = DummyBackbone()
def embed_input_ids(self, input_ids):
return input_ids
model = Qwen3OmniMoeThinkerForConditionalGeneration.__new__(
Qwen3OmniMoeThinkerForConditionalGeneration
)
nn.Module.__init__(model)
model.language_model = DummyLanguageModel()
assert supports_eagle3(model)
model.set_aux_hidden_state_layers((1, 2))
assert model.language_model.model.aux_hidden_state_layers == (1, 2)
@pytest.mark.skip_global_cleanup
def test_qwen3_omni_text_model_collects_post_deepstack_aux_hidden_states():
from vllm.model_executor.models.qwen3_omni_moe_thinker import Qwen3MoeLLMModel
class DummyLayer(nn.Module):
def forward(self, positions, hidden_states, residual):
return hidden_states + 1, torch.full_like(hidden_states, 10)
class DummyNorm(nn.Module):
def forward(self, hidden_states, residual):
return hidden_states + residual, None
model = Qwen3MoeLLMModel.__new__(Qwen3MoeLLMModel)
nn.Module.__init__(model)
model.start_layer = 0
model.end_layer = 1
model.layers = nn.ModuleList([DummyLayer()])
model.norm = DummyNorm()
model.aux_hidden_state_layers = (1,)
pp_group = Mock(is_first_rank=True, is_last_rank=True)
inputs_embeds = torch.tensor([[1.0]])
deepstack_inputs = {"deepstack_input_embeds_0": torch.tensor([[3.0]])}
with patch(
"vllm.model_executor.models.qwen3_omni_moe_thinker.get_pp_group",
return_value=pp_group,
):
output, aux_hidden_states = model.forward(
input_ids=None,
positions=torch.tensor([0]),
inputs_embeds=inputs_embeds,
deepstack_input_embeds=deepstack_inputs,
)
torch.testing.assert_close(output, torch.tensor([[15.0]]))
assert len(aux_hidden_states) == 1
torch.testing.assert_close(aux_hidden_states[0], torch.tensor([[15.0]]))
@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize(
("input_vocab_size", "draft_vocab_size", "weights", "error"),
[
(101, 100, [], "must include embed_tokens weights"),
(99, 99, [], "must include lm_head weights"),
(100, 40, [], "must include lm_head weights"),
(
100,
40,
[("lm_head.weight", torch.empty(40, 8))],
"must include a d2t mapping",
),
],
)
def test_qwen3_dspark_rejects_incomplete_vocab_weights(
input_vocab_size, draft_vocab_size, weights, error
):
from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM
model = Qwen3DSparkForCausalLM.__new__(Qwen3DSparkForCausalLM)
nn.Module.__init__(model)
object.__setattr__(
model,
"config",
SimpleNamespace(
vocab_size=input_vocab_size,
draft_vocab_size=draft_vocab_size,
),
)
object.__setattr__(model, "target_vocab_size", 100)
with pytest.raises(ValueError, match=error):
model.load_weights(weights)
@pytest.mark.skip_global_cleanup
def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary():
from vllm.v1.worker.gpu.spec_decode.dspark import utils as dspark_utils
target_embedding = nn.Embedding(100, 8)
draft_embedding = nn.Embedding(99, 8)
target_model = SimpleNamespace(model=SimpleNamespace(embed_tokens=target_embedding))
draft_model = SimpleNamespace(
model=SimpleNamespace(embed_tokens=draft_embedding),
has_own_embed_tokens=False,
)
draft_model_config = SimpleNamespace(
hf_config=SimpleNamespace(model_type="qwen3"),
get_vocab_size=Mock(return_value=99),
)
vllm_config = SimpleNamespace(
speculative_config=SimpleNamespace(
draft_model_config=draft_model_config,
attention_backend=None,
kv_cache_dtype=None,
),
attention_config=SimpleNamespace(backend=None),
cache_config=SimpleNamespace(),
model_config=SimpleNamespace(get_vocab_size=Mock(return_value=100)),
)
def fake_replace(config, **changes):
values = vars(config).copy()
values.update(changes)
return SimpleNamespace(**values)
with (
patch.object(dspark_utils, "replace", side_effect=fake_replace),
patch.object(
dspark_utils,
"get_pp_group",
return_value=SimpleNamespace(world_size=1),
),
patch(
"vllm.compilation.backends.set_model_tag",
return_value=nullcontext(),
),
patch(
"vllm.model_executor.model_loader.get_model",
return_value=draft_model,
),
patch(
"vllm.model_executor.models.qwen3_dflash.dflash_has_any_non_causal",
return_value=False,
),
patch(
"vllm.model_executor.models.utils.get_draft_quant_config",
return_value=None,
),
):
loaded_model = dspark_utils.load_dspark_model(target_model, vllm_config)
assert loaded_model.model.embed_tokens is target_embedding
if __name__ == "__main__":
pytest.main([__file__, "-v"])
tests/models/registry.py
+7
View File
@@ -1557,6 +1557,13 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
is_available_online=False,
use_original_num_layers=True, # DSpark backbone requires all layers
),
"Qwen3OmniDSparkModel": _HfExamplesInfo(
"Qwen/Qwen3-Omni-30B-A3B-Instruct",
speculative_model="Qwen/Qwen3-Omni-30B-A3B-DSpark",
is_available_online=False,
use_original_num_layers=True,
min_transformers_version="4.57",
),
"Gemma4DSparkModel": _HfExamplesInfo(
"google/gemma-4-12B-it",
speculative_model="deepseek-ai/dspark_gemma4_12b_block7",
tests/test_config.py
+121
View File
@@ -32,6 +32,7 @@ from vllm.config.compilation import CompilationMode, CUDAGraphMode
from vllm.config.kernel import IrOpPriorityConfig
from vllm.config.load import LoadConfig
from vllm.config.mamba import MambaBackendEnum
from vllm.config.speculative import _validate_qwen3_omni_dspark
from vllm.config.utils import get_field
from vllm.config.vllm import OPTIMIZATION_LEVEL_TO_CONFIG, OptimizationLevel
from vllm.platforms import current_platform
@@ -2100,6 +2101,126 @@ def test_mtp_draft_uses_model_weights_not_local_cache(mock_model_config_cls):
assert actual_model == s3_url
def _make_qwen3_omni_dspark_configs():
text_config = SimpleNamespace(
hidden_size=2048,
num_hidden_layers=48,
num_attention_heads=32,
num_key_value_heads=4,
head_dim=128,
vocab_size=152064,
)
target_model_config = SimpleNamespace(
hf_config=SimpleNamespace(
model_type="qwen3_omni_moe",
architectures=["Qwen3OmniMoeForConditionalGeneration"],
thinker_config=SimpleNamespace(text_config=text_config),
),
hf_text_config=text_config,
architectures=["Qwen3OmniMoeForConditionalGeneration"],
get_hidden_size=lambda: 2048,
get_total_num_hidden_layers=lambda: 48,
get_vocab_size=lambda: 152064,
)
draft_hf_config = SimpleNamespace(
model_type="qwen3",
architectures=["Qwen3OmniDSparkModel"],
hidden_size=2048,
num_attention_heads=32,
num_key_value_heads=4,
head_dim=128,
block_size=7,
target_hidden_size=2048,
target_layer_ids=[1, 9, 17, 25, 33],
use_aux_hidden_state=True,
markov_rank=64,
markov_head_type="vanilla",
sample_from_anchor=True,
dspark_bonus_anchor=False,
vocab_size=152064,
draft_vocab_size=32000,
mask_token_id=151669,
rope_parameters={"rope_type": "default", "rope_theta": 1000000.0},
)
draft_model_config = SimpleNamespace(
hf_config=draft_hf_config,
architectures=["Qwen3OmniDSparkModel"],
)
return target_model_config, draft_model_config
@pytest.mark.skip_global_cleanup
def test_qwen3_omni_dspark_checkpoint_contract_is_accepted():
target_config, draft_config = _make_qwen3_omni_dspark_configs()
_validate_qwen3_omni_dspark(target_config, draft_config, 7)
@pytest.mark.skip_global_cleanup
def test_qwen3_omni_dspark_rejects_generic_qwen3_architecture():
target_config, draft_config = _make_qwen3_omni_dspark_configs()
draft_config.hf_config.architectures = ["Qwen3DSparkModel"]
draft_config.architectures = ["Qwen3DSparkModel"]
with pytest.raises(ValueError, match="must be converted first"):
_validate_qwen3_omni_dspark(target_config, draft_config, 7)
@pytest.mark.parametrize(
("field", "value", "error"),
[
("block_size", 5, "trained block_size"),
("target_hidden_size", 4096, "target_hidden_size"),
("hidden_size", 4096, "draft hidden_size"),
("num_attention_heads", 16, "num_attention_heads"),
("num_key_value_heads", 8, "num_key_value_heads"),
("head_dim", 64, "head_dim"),
("target_layer_ids", [7, 48], "zero-based text-layer"),
("target_layer_ids", [23, 7], "strictly increasing"),
("use_aux_hidden_state", False, "use_aux_hidden_state=true"),
("markov_rank", 0, "markov_rank"),
("markov_head_type", "gated", "markov_head_type='vanilla'"),
("sample_from_anchor", False, "sample_from_anchor=true"),
("dspark_bonus_anchor", True, "dspark_bonus_anchor=false"),
("vocab_size", 0, "input vocab_size must be a positive integer"),
],
)
@pytest.mark.skip_global_cleanup
def test_qwen3_omni_dspark_rejects_incompatible_checkpoint_fields(field, value, error):
target_config, draft_config = _make_qwen3_omni_dspark_configs()
setattr(draft_config.hf_config, field, value)
with pytest.raises(ValueError, match=error):
_validate_qwen3_omni_dspark(target_config, draft_config, 7)
@pytest.mark.skip_global_cleanup
def test_qwen3_omni_dspark_rejects_mrope_draft_positions():
target_config, draft_config = _make_qwen3_omni_dspark_configs()
draft_config.hf_config.rope_parameters["mrope_section"] = [24, 20, 20]
with pytest.raises(ValueError, match="logical 1-D RoPE"):
_validate_qwen3_omni_dspark(target_config, draft_config, 7)
@pytest.mark.skip_global_cleanup
def test_qwen3_omni_dspark_allows_draft_only_noise_token_row():
target_config, draft_config = _make_qwen3_omni_dspark_configs()
draft_config.hf_config.vocab_size = 152065
draft_config.hf_config.draft_vocab_size = 152064
draft_config.hf_config.mask_token_id = 152064
_validate_qwen3_omni_dspark(target_config, draft_config, 7)
@pytest.mark.skip_global_cleanup
def test_qwen3_omni_dspark_allows_smaller_input_vocabulary():
target_config, draft_config = _make_qwen3_omni_dspark_configs()
draft_config.hf_config.vocab_size = 151936
draft_config.hf_config.mask_token_id = 151669
_validate_qwen3_omni_dspark(target_config, draft_config, 7)
def test_ir_op_priority_default():
"""Test that IR op priority defaults are set correctly."""
from vllm.config.kernel import IrOpPriorityConfig
tests/transformers_utils/test_speculators_dspark_config.py
+58
View File
@@ -0,0 +1,58 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm.transformers_utils.configs.speculators.base import SpeculatorsConfig
pytestmark = pytest.mark.skip_global_cleanup
def _make_dspark_config(architecture: str) -> dict:
return {
"speculators_model_type": "dspark",
"architectures": [architecture],
"sample_from_anchor": True,
"use_aux_hidden_state": True,
"draft_vocab_size": 32,
"target_hidden_size": 16,
"mask_token_id": 31,
"markov_rank": 4,
"markov_head_type": "vanilla",
"block_size": 3,
"aux_hidden_state_layer_ids": [2, 5],
"transformer_layer_config": {
"model_type": "qwen3",
"hidden_size": 16,
},
}
def test_dspark_updater_preserves_qwen3_omni_contract() -> None:
config = SpeculatorsConfig.extract_transformers_pre_trained_config(
_make_dspark_config("Qwen3OmniDSparkModel")
)
assert config["architectures"] == ["Qwen3OmniDSparkModel"]
assert config["sample_from_anchor"] is True
assert config["dspark_bonus_anchor"] is False
assert config["use_aux_hidden_state"] is True
assert config["target_layer_ids"] == [1, 4]
def test_dspark_updater_keeps_legacy_qwen3_checkpoint_loadable() -> None:
config = SpeculatorsConfig.extract_transformers_pre_trained_config(
_make_dspark_config("DSparkSpeculator")
)
assert config["architectures"] == ["Qwen3DSparkModel"]
def test_dspark_updater_maps_bonus_anchor_semantics() -> None:
outer_config = _make_dspark_config("Qwen3DSparkModel")
outer_config["sample_from_anchor"] = False
config = SpeculatorsConfig.extract_transformers_pre_trained_config(outer_config)
assert config["sample_from_anchor"] is False
assert config["dspark_bonus_anchor"] is True
vllm/config/speculative.py
+302 -2
View File
@@ -3,7 +3,7 @@
import copy
import functools
from collections.abc import Callable
from collections.abc import Callable, Mapping
from typing import TYPE_CHECKING, Any, Literal, get_args
from pydantic import Field, SkipValidation, field_validator, model_validator
@@ -80,6 +80,293 @@ SpeculativeMethod = Literal[
RejectionSampleMethod = Literal["standard", "synthetic", "block"]
DraftSampleMethod = Literal["greedy", "probabilistic"]
_QWEN3_OMNI_TARGET_ARCHITECTURES = frozenset(
{
"Qwen3OmniMoeForConditionalGeneration",
"Qwen3OmniMoeThinkerForConditionalGeneration",
}
)
_QWEN3_OMNI_DSPARK_ARCHITECTURE = "Qwen3OmniDSparkModel"
def _is_qwen3_omni_target(model_config: ModelConfig) -> bool:
hf_config = model_config.hf_config
architectures = set(getattr(model_config, "architectures", ()) or ())
architectures.update(getattr(hf_config, "architectures", ()) or ())
return getattr(hf_config, "model_type", None) == "qwen3_omni_moe" or bool(
architectures & _QWEN3_OMNI_TARGET_ARCHITECTURES
)
def _get_qwen3_omni_text_config(model_config: ModelConfig) -> Any:
thinker_config = getattr(model_config.hf_config, "thinker_config", None)
text_config = getattr(thinker_config, "text_config", None)
if text_config is None:
text_config = getattr(model_config, "hf_text_config", None)
return text_config
def _get_attention_head_dim(config: Any) -> int | None:
head_dim = getattr(config, "head_dim", None)
if head_dim is not None:
return head_dim
hidden_size = getattr(config, "hidden_size", None)
num_attention_heads = getattr(config, "num_attention_heads", None)
if (
isinstance(hidden_size, int)
and isinstance(num_attention_heads, int)
and num_attention_heads > 0
and hidden_size % num_attention_heads == 0
):
return hidden_size // num_attention_heads
return None
def _get_nested_config_value(config: Any, section: str, name: str) -> Any:
nested_config = getattr(config, section, None)
if isinstance(nested_config, Mapping):
return nested_config.get(name)
return getattr(nested_config, name, None)
def _get_qwen3_dspark_value(config: Any, name: str) -> Any:
# Match DFlashQwen3Model's precedence: dflash_config overrides
# eagle_config, and top-level values are the modern fallback.
value = _get_nested_config_value(config, "dflash_config", name)
if value is None:
value = _get_nested_config_value(config, "eagle_config", name)
if value is None:
value = getattr(config, name, None)
return value
def _validate_qwen3_omni_dspark(
target_model_config: ModelConfig,
draft_model_config: ModelConfig,
num_speculative_tokens: int,
) -> None:
"""Validate the standalone Qwen3-Omni DSpark checkpoint contract."""
if not _is_qwen3_omni_target(target_model_config):
return
draft_hf_config = draft_model_config.hf_config
draft_architectures = set(getattr(draft_model_config, "architectures", ()) or ())
draft_architectures.update(getattr(draft_hf_config, "architectures", ()) or ())
if _QWEN3_OMNI_DSPARK_ARCHITECTURE not in draft_architectures:
raise ValueError(
"Qwen3-Omni DSpark requires a standalone draft checkpoint with "
f"architectures=['{_QWEN3_OMNI_DSPARK_ARCHITECTURE}']; generic "
"Qwen3DSparkModel checkpoints must be converted first."
)
block_size = _get_qwen3_dspark_value(draft_hf_config, "block_size")
if block_size is None:
block_size = _get_qwen3_dspark_value(draft_hf_config, "dspark_block_size")
if (
not isinstance(block_size, int)
or isinstance(block_size, bool)
or block_size <= 0
):
raise ValueError(
"Qwen3-Omni DSpark requires a positive integer block_size in the "
"draft config."
)
if num_speculative_tokens != block_size:
raise ValueError(
"Qwen3-Omni DSpark requires num_speculative_tokens to match the "
f"trained block_size ({block_size}); got {num_speculative_tokens}."
)
target_text_config = _get_qwen3_omni_text_config(target_model_config)
if target_text_config is None:
raise ValueError(
"Qwen3-Omni DSpark could not resolve the target thinker text_config."
)
target_hidden_size = getattr(
target_text_config,
"hidden_size",
target_model_config.get_hidden_size(),
)
draft_target_hidden_size = getattr(draft_hf_config, "target_hidden_size", None)
if draft_target_hidden_size != target_hidden_size:
raise ValueError(
"Qwen3-Omni DSpark draft target_hidden_size must match the target "
f"text hidden size ({target_hidden_size}); got {draft_target_hidden_size}."
)
draft_hidden_size = getattr(draft_hf_config, "hidden_size", None)
if draft_hidden_size != target_hidden_size:
raise ValueError(
"Qwen3-Omni DSpark draft hidden_size must match the target thinker "
f"hidden size ({target_hidden_size}); got {draft_hidden_size}."
)
target_attention = {
"num_attention_heads": getattr(target_text_config, "num_attention_heads", None),
"num_key_value_heads": getattr(
target_text_config,
"num_key_value_heads",
getattr(target_text_config, "num_attention_heads", None),
),
"head_dim": _get_attention_head_dim(target_text_config),
}
draft_attention = {
"num_attention_heads": getattr(draft_hf_config, "num_attention_heads", None),
"num_key_value_heads": getattr(
draft_hf_config,
"num_key_value_heads",
getattr(draft_hf_config, "num_attention_heads", None),
),
"head_dim": _get_attention_head_dim(draft_hf_config),
}
for name, target_value in target_attention.items():
if (
not isinstance(target_value, int)
or isinstance(target_value, bool)
or target_value <= 0
):
raise ValueError(
"Qwen3-Omni DSpark requires a positive integer "
f"{name} in the target thinker text_config; got {target_value}."
)
draft_value = draft_attention[name]
if draft_value != target_value:
raise ValueError(
f"Qwen3-Omni DSpark draft {name} must match the target thinker "
f"value ({target_value}); got {draft_value}."
)
target_layer_ids = _get_nested_config_value(
draft_hf_config, "dflash_config", "target_layer_ids"
)
if target_layer_ids is None:
target_layer_ids = getattr(draft_hf_config, "dspark_target_layer_ids", None)
if target_layer_ids is None:
target_layer_ids = getattr(draft_hf_config, "target_layer_ids", None)
if not isinstance(target_layer_ids, (list, tuple)) or not target_layer_ids:
raise ValueError(
"Qwen3-Omni DSpark requires a non-empty target_layer_ids list in "
"the draft config."
)
if any(
not isinstance(layer_id, int) or isinstance(layer_id, bool)
for layer_id in target_layer_ids
):
raise ValueError(
"Qwen3-Omni DSpark target_layer_ids must contain only integers."
)
if list(target_layer_ids) != sorted(set(target_layer_ids)):
raise ValueError(
"Qwen3-Omni DSpark target_layer_ids must be unique and strictly increasing."
)
target_num_layers = getattr(
target_text_config,
"num_hidden_layers",
target_model_config.get_total_num_hidden_layers(),
)
if target_layer_ids[0] < 0 or target_layer_ids[-1] >= target_num_layers:
raise ValueError(
"Qwen3-Omni DSpark target_layer_ids must be zero-based text-layer "
f"indices in [0, {target_num_layers - 1}]; got {target_layer_ids}."
)
if _get_qwen3_dspark_value(draft_hf_config, "use_aux_hidden_state") is not True:
raise ValueError(
"Qwen3-Omni DSpark requires use_aux_hidden_state=true because Omni "
"conditioning is supplied through thinker hidden states."
)
markov_rank = getattr(draft_hf_config, "markov_rank", None)
if (
not isinstance(markov_rank, int)
or isinstance(markov_rank, bool)
or markov_rank <= 0
):
raise ValueError(
"Qwen3-Omni DSpark requires a positive integer markov_rank in the "
"draft config."
)
markov_head_type = getattr(draft_hf_config, "markov_head_type", None)
if markov_head_type != "vanilla":
raise ValueError(
"Qwen3-Omni DSpark currently requires markov_head_type='vanilla'; "
f"got {markov_head_type!r}."
)
sample_from_anchor = getattr(draft_hf_config, "sample_from_anchor", True)
bonus_anchor = getattr(draft_hf_config, "dspark_bonus_anchor", False)
if sample_from_anchor is not True or bonus_anchor is not False:
raise ValueError(
"Qwen3-Omni DSpark requires sample_from_anchor=true and "
"dspark_bonus_anchor=false."
)
target_vocab_size = getattr(
target_text_config,
"vocab_size",
target_model_config.get_vocab_size(),
)
draft_input_vocab_size = getattr(draft_hf_config, "vocab_size", None)
if (
not isinstance(draft_input_vocab_size, int)
or isinstance(draft_input_vocab_size, bool)
or draft_input_vocab_size <= 0
):
raise ValueError(
"Qwen3-Omni DSpark input vocab_size must be a positive integer; "
f"got {draft_input_vocab_size}."
)
draft_output_vocab_size = getattr(draft_hf_config, "draft_vocab_size", None)
if draft_output_vocab_size is None:
draft_output_vocab_size = draft_input_vocab_size
if (
not isinstance(draft_output_vocab_size, int)
or isinstance(draft_output_vocab_size, bool)
or not 0 < draft_output_vocab_size <= target_vocab_size
):
raise ValueError(
"Qwen3-Omni DSpark draft_vocab_size must be a positive integer no "
f"larger than target vocab_size ({target_vocab_size}); got "
f"{draft_output_vocab_size}."
)
noise_token_id = _get_nested_config_value(
draft_hf_config, "dflash_config", "mask_token_id"
)
for name in (
"mask_token_id",
"dspark_noise_token_id",
"pard_token",
"ptd_token_id",
):
if noise_token_id is None:
noise_token_id = getattr(draft_hf_config, name, None)
if noise_token_id is not None:
break
if (
not isinstance(noise_token_id, int)
or isinstance(noise_token_id, bool)
or not 0 <= noise_token_id < draft_input_vocab_size
):
raise ValueError(
"Qwen3-Omni DSpark requires a valid mask/noise token id within the "
f"draft input vocabulary [0, {draft_input_vocab_size - 1}]."
)
rope_configs = (
getattr(draft_hf_config, "rope_parameters", None),
getattr(draft_hf_config, "rope_scaling", None),
)
has_mrope = getattr(draft_hf_config, "mrope_section", None) is not None
has_mrope = has_mrope or any(
isinstance(rope_config, Mapping) and "mrope_section" in rope_config
for rope_config in rope_configs
)
if has_mrope:
raise ValueError(
"Qwen3-Omni DSpark draft checkpoints must use logical 1-D RoPE and "
"must not define mrope_section."
)
@config
class SpeculativeConfig:
@@ -964,6 +1251,8 @@ class SpeculativeConfig:
elif (
"dspark" in self.draft_model_config.model.lower()
or "Qwen3DSparkModel" in self.draft_model_config.architectures
or _QWEN3_OMNI_DSPARK_ARCHITECTURE
in self.draft_model_config.architectures
or "Gemma4DSparkModel" in self.draft_model_config.architectures
or (
"DSparkDraftModel" in self.draft_model_config.architectures
@@ -1037,6 +1326,8 @@ class SpeculativeConfig:
self.update_arch_()
elif self.method == "dspark" and (
"Qwen3DSparkModel" not in self.draft_model_config.architectures
and _QWEN3_OMNI_DSPARK_ARCHITECTURE
not in self.draft_model_config.architectures
and "Gemma4DSparkModel" not in self.draft_model_config.architectures
and "K3DSparkModel" not in self.draft_model_config.architectures
):
@@ -1125,13 +1416,22 @@ class SpeculativeConfig:
if (
"Qwen3DSparkModel"
not in self.draft_model_config.architectures
and _QWEN3_OMNI_DSPARK_ARCHITECTURE
not in self.draft_model_config.architectures
):
raise ValueError(
"dspark_draft_topk is only supported by "
"Qwen3DSparkModel"
"Qwen3DSparkModel and Qwen3OmniDSparkModel"
)
hf_config.dspark_draft_topk = dspark_draft_topk
assert self.target_model_config is not None
_validate_qwen3_omni_dspark(
self.target_model_config,
self.draft_model_config,
self.num_speculative_tokens,
)
self.draft_tensor_parallel_size = (
SpeculativeConfig._verify_and_get_draft_tp(
self.target_parallel_config,
vllm/model_executor/models/qwen3_dflash.py
+8 -8
View File
@@ -377,8 +377,6 @@ class DFlashQwen3DecoderLayer(nn.Module):
@support_torch_compile
class DFlashQwen3Model(nn.Module):
decoder_layer_cls = DFlashQwen3DecoderLayer
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_substr={
"midlayer.": "layers.0.",
@@ -412,10 +410,10 @@ class DFlashQwen3Model(nn.Module):
drafter_config = getattr(self.config, "eagle_config", {})
drafter_config.update(getattr(self.config, "dflash_config", {}))
if drafter_config is not None and "use_aux_hidden_state" in drafter_config:
self.use_aux_hidden_state = drafter_config["use_aux_hidden_state"]
else:
self.use_aux_hidden_state = True
self.use_aux_hidden_state = drafter_config.get(
"use_aux_hidden_state",
getattr(self.config, "use_aux_hidden_state", True),
)
current_vllm_config = get_current_vllm_config()
@@ -430,7 +428,9 @@ class DFlashQwen3Model(nn.Module):
# at that slot id. Some checkpoints (XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash) ship
# with a separate mask embedding tensor to use instead. When present, we load it
# and substitute it for embed_tokens[mask_token_id] when computing embeddings.
self.mask_token_id = drafter_config.get("mask_token_id")
self.mask_token_id = drafter_config.get(
"mask_token_id", getattr(self.config, "mask_token_id", None)
)
self.mask_embedding = nn.Parameter(
torch.zeros(self.config.hidden_size, dtype=vllm_config.model_config.dtype),
requires_grad=False,
@@ -439,7 +439,7 @@ class DFlashQwen3Model(nn.Module):
self.layers = nn.ModuleList(
[
self.decoder_layer_cls(
DFlashQwen3DecoderLayer(
current_vllm_config,
config=self.config,
layer_idx=layer_idx,
vllm/model_executor/models/qwen3_dspark.py
+20 -2
View File
@@ -200,8 +200,8 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
self.logits_processor = LogitsProcessor(
self.config.draft_vocab_size, scale=logit_scale
)
target_vocab_size = vllm_config.model_config.get_vocab_size()
if self.config.draft_vocab_size != target_vocab_size:
self.target_vocab_size = vllm_config.model_config.get_vocab_size()
if self.config.draft_vocab_size != self.target_vocab_size:
self.draft_id_to_target_id = nn.Parameter(
torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
requires_grad=False,
@@ -280,6 +280,24 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
# mask_embedding is an unused placeholder param; DSpark masks via the vocab row.
# embed_tokens / lm_head are optional; when omitted they are shared from
# the target by load_dspark_model, so skip the unloaded params here.
uses_expanded_input_vocab = self.config.vocab_size > self.target_vocab_size
uses_reduced_vocab = self.config.draft_vocab_size != self.target_vocab_size
if uses_expanded_input_vocab and not includes_embed_tokens:
raise ValueError(
"Qwen3 DSpark checkpoints whose input vocab_size is larger than "
"the target vocabulary must include embed_tokens weights."
)
if uses_reduced_vocab and not includes_lm_head:
raise ValueError(
"Reduced-vocabulary Qwen3 DSpark checkpoints must include "
"lm_head weights; the full target lm_head cannot be shared."
)
if uses_reduced_vocab and not includes_draft_id_mapping:
raise ValueError(
"Reduced-vocabulary Qwen3 DSpark checkpoints must include a "
"d2t mapping so sampled draft ids can be converted to target ids."
)
orig_to_new_substr = {"mask_embedding": None}
if not includes_embed_tokens:
orig_to_new_substr["embed_tokens"] = None
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+11 -1
View File
@@ -83,6 +83,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum
from .interfaces import (
MultiModalEmbeddings,
SupportsEagle3,
SupportsMRoPE,
SupportsMultiModal,
SupportsPP,
@@ -1067,7 +1068,7 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
deepstack_input_embeds: IntermediateTensors | None = None,
) -> torch.Tensor | IntermediateTensors:
) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
if get_pp_group().is_first_rank:
if inputs_embeds is not None:
hidden_states = inputs_embeds
@@ -1078,6 +1079,9 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
assert intermediate_tensors is not None
hidden_states = intermediate_tensors["hidden_states"]
residual = intermediate_tensors["residual"]
aux_hidden_states = self._maybe_add_hidden_state(
[], self.start_layer, hidden_states, residual
)
for layer_idx, layer in enumerate(
self.layers[self.start_layer : self.end_layer]
):
@@ -1096,12 +1100,17 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
hidden_states
+ deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"]
)
self._maybe_add_hidden_state(
aux_hidden_states, layer_idx + 1, hidden_states, residual
)
if not get_pp_group().is_last_rank:
return IntermediateTensors(
{"hidden_states": hidden_states, "residual": residual}
)
hidden_states, _ = self.norm(hidden_states, residual)
if aux_hidden_states:
return hidden_states, aux_hidden_states
return hidden_states
@@ -1577,6 +1586,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
SupportsMultiModal,
SupportsPP,
SupportsMRoPE,
SupportsEagle3,
Qwen3OmniMoeConditionalGenerationMixin,
SupportsTranscription,
):
vllm/model_executor/models/registry.py
+1
View File
@@ -643,6 +643,7 @@ _SPECULATIVE_DECODING_MODELS = {
"DFlashMuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
"DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"),
"Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
"Qwen3OmniDSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
"K3DSparkModel": (
"vllm.models.kimi_k3.nvidia.dspark_mla",
"K3DSparkForCausalLM",
vllm/transformers_utils/configs/speculators/algos.py
+17 -4
View File
@@ -154,10 +154,22 @@ def update_dspark(config_dict: dict, pre_trained_config: dict) -> None:
False (anchor is a bonus token, only mask tokens predict, yielding
block_size - 1 speculative tokens).
"""
pre_trained_config["architectures"] = ["Qwen3DSparkModel"]
pre_trained_config["sample_from_anchor"] = config_dict.get(
"sample_from_anchor", False
)
architectures = config_dict.get("architectures") or []
supported_architectures = {
"Qwen3DSparkModel",
"Qwen3OmniDSparkModel",
}
selected_architectures = [
architecture
for architecture in architectures
if architecture in supported_architectures
]
# Legacy msModelSpec checkpoints use a training-only architecture name.
pre_trained_config["architectures"] = selected_architectures or ["Qwen3DSparkModel"]
sample_from_anchor = config_dict.get("sample_from_anchor", False)
pre_trained_config["sample_from_anchor"] = sample_from_anchor
pre_trained_config["dspark_bonus_anchor"] = not sample_from_anchor
aux_layer_ids = config_dict["aux_hidden_state_layer_ids"]
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:
"block_size",
"enable_confidence_head",
"confidence_head_with_markov",
"use_aux_hidden_state",
):
if config_dict.get(key) is not None:
pre_trained_config[key] = config_dict[key]
vllm/v1/worker/gpu/spec_decode/dspark/utils.py
+15 -4
View File
@@ -86,11 +86,16 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
)
target_inner = target_language_model.model
draft_inner = draft_model.model
target_vocab_size = vllm_config.model_config.get_vocab_size()
target_embed = getattr(target_inner, "embed_tokens", None)
draft_embed = getattr(draft_inner, "embed_tokens", None)
if target_embed is not None and _should_share(
draft_model, "has_own_embed_tokens", draft_embed, target_embed
if (
target_embed is not None
and draft_model_config.get_vocab_size() <= target_vocab_size
and _should_share(
draft_model, "has_own_embed_tokens", draft_embed, target_embed
)
):
if draft_embed is not None:
del draft_inner.embed_tokens
@@ -98,8 +103,14 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
target_lm_head = get_target_lm_head(target_model, target_language_model)
draft_lm_head = getattr(draft_model, "lm_head", None)
if target_lm_head is not None and _should_share(
draft_model, "has_own_lm_head", draft_lm_head, target_lm_head
draft_output_vocab_size = (
getattr(draft_model_config.hf_config, "draft_vocab_size", None)
or draft_model_config.get_vocab_size()
)
if (
target_lm_head is not None
and draft_output_vocab_size == target_vocab_size
and _should_share(draft_model, "has_own_lm_head", draft_lm_head, target_lm_head)
):
if draft_lm_head is not None:
del draft_model.lm_head