diff --git a/tests/model_executor/models/spark2_5/__init__.py b/tests/model_executor/models/spark2_5/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_executor/models/spark2_5/test_spark2_5.py b/tests/model_executor/models/spark2_5/test_spark2_5.py new file mode 100644 index 000000000..2620c7931 --- /dev/null +++ b/tests/model_executor/models/spark2_5/test_spark2_5.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for the Spark-X2.5 model definition. + +These cover the three places where Spark departs from a Llama-style decoder +and where a port is most likely to go wrong: + +* the fused ``q_k_v_proj`` checkpoint layout has to be split before vLLM's + sharded ``QKVParallelLinear`` can load it, +* ``layer_types`` must map onto per-layer sliding windows, and +* each layer type carries its own RoPE theta and partial-rotary factor. + +End-to-end numerical parity against the reference implementation is measured +separately (see analysis/spark_edge_support.md); it needs the real 3.2 GiB +checkpoint and so does not belong in the unit suite. +""" + +import pytest +import torch + +from vllm_omni.model_executor.models.spark2_5.configuration_spark2_5 import ( + Spark2_5Config, +) +from vllm_omni.model_executor.models.spark2_5.spark2_5 import _split_fused_qkv + + +def _config(**overrides) -> Spark2_5Config: + kwargs = dict( + vocab_size=256, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + hidden_act="gelu", + sliding_window=16, + headwise_attn_output_gate=True, + layer_types=[ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ], + rope_parameters={ + "full_attention": {"rope_theta": 5000000, "partial_rotary_factor": 0.25}, + "sliding_attention": {"rope_theta": 10000, "partial_rotary_factor": 1.0}, + }, + tie_word_embeddings=True, + ) + kwargs.update(overrides) + return Spark2_5Config(**kwargs) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_split_fused_qkv_recovers_q_k_v_shards(): + config = _config() + q_dim = config.num_attention_heads * config.head_dim + kv_dim = config.num_key_value_heads * config.head_dim + + # Rows are tagged by value so a mis-ordered split is visible. + fused = torch.cat( + [ + torch.full((q_dim, config.hidden_size), 1.0), + torch.full((kv_dim, config.hidden_size), 2.0), + torch.full((kv_dim, config.hidden_size), 3.0), + ], + dim=0, + ) + out = dict( + _split_fused_qkv( + [("model.layers.0.self_attn.q_k_v_proj.weight", fused)], config + ) + ) + + assert set(out) == { + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.k_proj.weight", + "model.layers.0.self_attn.v_proj.weight", + } + assert out["model.layers.0.self_attn.q_proj.weight"].shape == ( + q_dim, + config.hidden_size, + ) + assert out["model.layers.0.self_attn.k_proj.weight"].shape == ( + kv_dim, + config.hidden_size, + ) + assert torch.all(out["model.layers.0.self_attn.q_proj.weight"] == 1.0) + assert torch.all(out["model.layers.0.self_attn.k_proj.weight"] == 2.0) + assert torch.all(out["model.layers.0.self_attn.v_proj.weight"] == 3.0) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_split_fused_qkv_passes_other_weights_through(): + config = _config() + mlp = torch.zeros(config.intermediate_size, config.hidden_size) + out = list(_split_fused_qkv([("model.layers.0.mlp.gate_proj.weight", mlp)], config)) + assert len(out) == 1 + assert out[0][0] == "model.layers.0.mlp.gate_proj.weight" + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_config_exposes_per_layer_type_rope(): + config = _config() + # Sliding layers rotate every head dim with a small theta; full-attention + # layers rotate only the first quarter with a large one. + assert config.get_partial_rotary_factor("sliding_attention") == 1.0 + assert config.get_rope_theta("sliding_attention") == 10000 + assert config.get_partial_rotary_factor("full_attention") == 0.25 + assert config.get_rope_theta("full_attention") == 5000000 + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_config_rejects_layer_types_of_the_wrong_length(): + with pytest.raises(ValueError, match="layer_types length"): + _config(layer_types=["full_attention"]) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_headwise_gate_scales_each_head_independently(): + """The gate is one sigmoid scalar per head, broadcast over head_dim.""" + config = _config() + num_heads, head_dim = config.num_attention_heads, config.head_dim + attn_output = torch.ones(3, num_heads * head_dim) + gate_score = torch.tensor([[-10.0, 0.0, 10.0, 0.0]]).repeat(3, 1) + + gate = torch.sigmoid(gate_score.float()).to(attn_output.dtype) + gated = ( + attn_output.view(-1, num_heads, head_dim) * gate.view(-1, num_heads, 1) + ).view(-1, num_heads * head_dim) + + per_head = gated.view(-1, num_heads, head_dim) + # A strongly negative score nearly zeroes its head, a zero score halves it, + # and every dim inside one head gets the identical factor. + assert per_head[:, 0].max() < 1e-3 + assert torch.allclose(per_head[:, 1], torch.full_like(per_head[:, 1], 0.5)) + assert per_head[:, 2].min() > 0.999 + assert torch.allclose( + per_head[:, 0], per_head[:, 0, :1].expand_as(per_head[:, 0]) + ) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_registered_in_the_omni_registry(): + from vllm_omni.model_executor.models.registry import _OMNI_MODELS + + assert _OMNI_MODELS["Spark2_5ForCausalLM"] == ( + "spark2_5", + "spark2_5", + "Spark2_5ForCausalLM", + ) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index b3d894eb4..3fc39dc28 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -75,6 +75,9 @@ def _register_omni_hf_configs() -> None: from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( Qwen3TTSConfig, ) + from vllm_omni.model_executor.models.spark2_5.configuration_spark2_5 import ( + Spark2_5Config, + ) from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config from vllm_omni.transformers_utils.configs.glm_tts import GLMTTSConfig from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig @@ -99,6 +102,7 @@ def _register_omni_hf_configs() -> None: ("moss_tts_local", MossTTSLocalConfig), ("moss_tts_realtime", MossTTSRealtimeConfig), ("qwen3_tts", Qwen3TTSConfig), + ("spark2_5", Spark2_5Config), ("personaplex", PersonaPlexConfig), ("nemotron_voicechat", NemotronVoiceChatConfig), ("cosyvoice3", CosyVoice3Config), diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index cb5167597..8a7e5c50f 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -9,6 +9,11 @@ from vllm.model_executor.models.registry import ( ) _OMNI_MODELS = { + "Spark2_5ForCausalLM": ( + "spark2_5", + "spark2_5", + "Spark2_5ForCausalLM", + ), "Qwen2_5OmniForConditionalGeneration": ( "qwen2_5_omni", "qwen2_5_omni", diff --git a/vllm_omni/model_executor/models/spark2_5/__init__.py b/vllm_omni/model_executor/models/spark2_5/__init__.py new file mode 100644 index 000000000..48601c834 --- /dev/null +++ b/vllm_omni/model_executor/models/spark2_5/__init__.py @@ -0,0 +1,4 @@ +from .configuration_spark2_5 import Spark2_5Config +from .spark2_5 import Spark2_5ForCausalLM + +__all__ = ["Spark2_5Config", "Spark2_5ForCausalLM"] diff --git a/vllm_omni/model_executor/models/spark2_5/configuration_spark2_5.py b/vllm_omni/model_executor/models/spark2_5/configuration_spark2_5.py new file mode 100644 index 000000000..4b463e5b5 --- /dev/null +++ b/vllm_omni/model_executor/models/spark2_5/configuration_spark2_5.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 The XHToken team and the HuggingFace Inc. team. +"""Spark-X2.5 configuration, vendored so the model loads without +``trust_remote_code``. + +Mirrors ``configuration_spark.py`` shipped with XHToken/Spark-X2.5-*. +""" + +from transformers import PretrainedConfig + + +class Spark2_5Config(PretrainedConfig): + model_type = "spark2_5" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=32000, + hidden_size=4096, + intermediate_size=11008, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + hidden_act="gelu", + max_position_embeddings=2048, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=None, + bos_token_id=1, + eos_token_id=2, + tie_word_embeddings=False, + rope_parameters=None, + attention_bias=False, + attention_dropout=0.0, + mlp_bias=False, + head_dim=None, + headwise_attn_output_gate=False, + gate_attn_act_mode="sigmoid", + sliding_window=None, + layer_types=None, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + if num_attention_heads % num_key_value_heads != 0: + raise ValueError( + f"num_attention_heads ({num_attention_heads}) must be divisible " + f"by num_key_value_heads ({num_key_value_heads})" + ) + self.num_key_value_heads = num_key_value_heads + + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.mlp_bias = mlp_bias + self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads + self.headwise_attn_output_gate = headwise_attn_output_gate + self.gate_attn_act_mode = gate_attn_act_mode + self.sliding_window = sliding_window + self.rope_parameters = rope_parameters + + if layer_types is None: + layer_types = ["full_attention"] * num_hidden_layers + if len(layer_types) != num_hidden_layers: + raise ValueError( + f"layer_types length ({len(layer_types)}) must match " + f"num_hidden_layers ({num_hidden_layers})" + ) + self.layer_types = layer_types + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + def get_rope_theta(self, layer_type): + return (self.rope_parameters or {}).get(layer_type, {}).get("rope_theta", 10000) + + def get_partial_rotary_factor(self, layer_type): + return (self.rope_parameters or {}).get(layer_type, {}).get( + "partial_rotary_factor", 1.0 + ) + + +__all__ = ["Spark2_5Config"] diff --git a/vllm_omni/model_executor/models/spark2_5/spark2_5.py b/vllm_omni/model_executor/models/spark2_5/spark2_5.py new file mode 100644 index 000000000..dbcccedff --- /dev/null +++ b/vllm_omni/model_executor/models/spark2_5/spark2_5.py @@ -0,0 +1,407 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright 2026 The XHToken team. All rights reserved. +"""Inference-only Spark-X2.5 model compatible with HuggingFace weights. + +Spark-X2.5 (XHToken/iFlytek) is a small on-device model built for agentic +tool use. Three traits separate it from a Llama-style decoder and drive the +implementation below: + +* **Hybrid attention.** ``layer_types`` interleaves three 512-token + ``sliding_attention`` layers with one ``full_attention`` layer, which is what + keeps the 1M-token context affordable on a phone. +* **Per-layer-type RoPE.** Sliding layers rotate every one of the 256 head + dims with theta 1e4; full layers rotate only the first quarter with theta + 5e6. Both come out of the nested ``rope_parameters`` dict. +* **Head-wise attention output gate.** A ``g_proj`` of shape + ``[hidden, num_heads]`` reads the *same* post-norm input as QKV and scales + each head's attention output by a sigmoid before ``out_proj``. + +The checkpoint stores QKV pre-fused as a single ``q_k_v_proj``; load_weights +splits it so vLLM's sharded ``QKVParallelLinear`` can take it. +""" + +from collections.abc import Iterable + +import torch +from torch import nn + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size +from vllm.model_executor.layers.activation import get_act_and_mul_fn +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + extract_layer_index, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) +from vllm.sequence import IntermediateTensors + +from .configuration_spark2_5 import Spark2_5Config + + +class Spark2_5MLP(nn.Module): + def __init__( + self, + config: Spark2_5Config, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + config.hidden_size, + [config.intermediate_size] * 2, + bias=config.mlp_bias, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + config.intermediate_size, + config.hidden_size, + bias=config.mlp_bias, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + # Spark's config validator rejects anything but exact GELU. + if config.hidden_act != "gelu": + raise ValueError( + f"Spark-X2.5 only supports hidden_act='gelu', got {config.hidden_act!r}" + ) + self.act_fn = get_act_and_mul_fn(config.hidden_act) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class Spark2_5Attention(nn.Module): + def __init__( + self, + config: Spark2_5Config, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + tp_size = get_tensor_model_parallel_world_size() + + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + self.qkv_proj = QKVParallelLinear( + config.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.out_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.out_proj", + ) + + # One scalar gate per attention head. Sharded column-wise so each rank + # gets exactly the gates for the query heads it owns. + self.headwise_attn_output_gate = config.headwise_attn_output_gate + if self.headwise_attn_output_gate: + self.g_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_heads, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.g_proj", + ) + if config.gate_attn_act_mode not in ("sigmoid", "silu"): + raise ValueError( + f"Unsupported gate_attn_act_mode: {config.gate_attn_act_mode}" + ) + self.gate_attn_act_mode = config.gate_attn_act_mode + else: + self.g_proj = None + + layer_idx = extract_layer_index(prefix) + layer_type = config.layer_types[layer_idx] + self.is_sliding = layer_type == "sliding_attention" + sliding_window = config.sliding_window if self.is_sliding else None + + # Sliding and full layers are trained with different theta and + # different partial-rotary factors; both live under rope_parameters. + rope_parameters = dict(config.rope_parameters[layer_type]) + rope_parameters.setdefault("rope_type", "default") + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + rope_parameters=rope_parameters, + is_neox_style=True, + ) + + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + per_layer_sliding_window=sliding_window, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v) + + if self.g_proj is not None: + gate_score, _ = self.g_proj(hidden_states) + if self.gate_attn_act_mode == "sigmoid": + gate = torch.sigmoid(gate_score.float()) + else: + gate = nn.functional.silu(gate_score.float()) + gate = gate.to(attn_output.dtype) + attn_output = ( + attn_output.view(-1, self.num_heads, self.head_dim) + * gate.view(-1, self.num_heads, 1) + ).view(-1, self.num_heads * self.head_dim) + + output, _ = self.out_proj(attn_output) + return output + + +class Spark2_5DecoderLayer(nn.Module): + def __init__( + self, + config: Spark2_5Config, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.self_attn = Spark2_5Attention( + config, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + self.mlp = Spark2_5MLP(config, quant_config=quant_config, prefix=f"{prefix}.mlp") + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) + + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +@support_torch_compile +class Spark2_5Model(nn.Module): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={"model.embedding.": "model.embed_tokens."}, + orig_to_new_stacked={ + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + }, + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: Spark2_5DecoderLayer( + config, cache_config, quant_config, prefix=prefix + ), + prefix=f"{prefix}.layers", + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + for layer in self.layers[self.start_layer : self.end_layer]: + hidden_states, residual = layer(positions, 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) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + +def _split_fused_qkv( + weights: Iterable[tuple[str, torch.Tensor]], + config: Spark2_5Config, +) -> Iterable[tuple[str, torch.Tensor]]: + """Split the checkpoint's fused ``q_k_v_proj`` into q/k/v. + + vLLM's ``QKVParallelLinear`` loads one shard at a time so that tensor + parallelism can slice heads; the Spark checkpoint ships the three + projections already concatenated, so undo that here. + """ + q_dim = config.num_attention_heads * config.head_dim + kv_dim = config.num_key_value_heads * config.head_dim + for name, loaded in weights: + if ".self_attn.q_k_v_proj." in name: + q, k, v = loaded.split([q_dim, kv_dim, kv_dim], dim=0) + for shard, tensor in (("q_proj", q), ("k_proj", k), ("v_proj", v)): + yield name.replace("q_k_v_proj", shard), tensor + else: + yield name, loaded + + +class Spark2_5ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + hf_to_vllm_mapper = Spark2_5Model.hf_to_vllm_mapper + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + + self.model = Spark2_5Model( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + if config.tie_word_embeddings: + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) + + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors: + return self.model( + input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), + ) + return loader.load_weights( + _split_fused_qkv(weights, self.config), mapper=self.hf_to_vllm_mapper + )