diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py index 1ceaba04aa..54d2f2672f 100644 --- a/tests/model_executor/test_qwen3_omni.py +++ b/tests/model_executor/test_qwen3_omni.py @@ -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"]) diff --git a/tests/models/registry.py b/tests/models/registry.py index 24a4869399..712f160e3d 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -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", diff --git a/tests/test_config.py b/tests/test_config.py index eac6a95195..40ab2fa426 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -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 diff --git a/tests/transformers_utils/test_speculators_dspark_config.py b/tests/transformers_utils/test_speculators_dspark_config.py new file mode 100644 index 0000000000..ce32b6e5df --- /dev/null +++ b/tests/transformers_utils/test_speculators_dspark_config.py @@ -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 diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index f0eee88429..d3b4af1cf0 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -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, diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 45f0bf5344..b7b5b29d41 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -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, diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index 250d54d9ca..2cee016124 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/vllm/model_executor/models/qwen3_dspark.py @@ -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 diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 0e2aab634c..2bcc07a276 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -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, ): diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 24fc66a089..c6043716a8 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -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", diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py index fbd2c45181..8ec978d735 100644 --- a/vllm/transformers_utils/configs/speculators/algos.py +++ b/vllm/transformers_utils/configs/speculators/algos.py @@ -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] diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index a33fa259c3..177060747a 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -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