From 716c5225aefa7aae0628d88a6d58adddbca7c22c Mon Sep 17 00:00:00 2001 From: Timofey <37947728+timka-byzov@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:15:03 +0300 Subject: [PATCH] [Model] Add AliceAI support (#57963) Signed-off-by: timka Co-authored-by: Codex --- .buildkite/test_areas/models_basic.yaml | 4 +- docs/models/supported_models.md | 1 + tests/models/registry.py | 4 + tests/models/test_alice_ai.py | 231 ++++++++ tests/models/test_alice_ai_config.py | 93 +++ vllm/config/vllm.py | 1 + vllm/model_executor/models/alice_ai.py | 733 ++++++++++++++++++++++++ vllm/model_executor/models/config.py | 46 ++ vllm/model_executor/models/registry.py | 1 + 9 files changed, 1113 insertions(+), 1 deletion(-) create mode 100644 tests/models/test_alice_ai.py create mode 100644 tests/models/test_alice_ai_config.py create mode 100644 vllm/model_executor/models/alice_ai.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index e2c183b8e4..ebc17c1cce 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -254,6 +254,8 @@ steps: - tests/models/test_vision.py - tests/models/test_adapters.py - tests/models/test_qwen3_5_mtp_config.py + - tests/models/test_alice_ai.py + - tests/models/test_alice_ai_config.py - tests/models/transformers/fusers/ - tests/models/transformers/test_dynamic_rope.py - tests/model_executor/layers/test_activation.py @@ -264,7 +266,7 @@ steps: - tests/models/transformers/test_layer_registry.py device: cpu-small commands: - - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ models/transformers/test_dynamic_rope.py model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py + - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ models/transformers/test_dynamic_rope.py model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py models/test_alice_ai.py models/test_alice_ai_config.py - label: ':amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics' key: dsv41-mono-decode-numerics-amd diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 156210e434..0b69ac8f47 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -372,6 +372,7 @@ th { | Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | | ------------ | ------ | ----------------- | -------------------- | ------------------------- | | `AfmoeForCausalLM` | Afmoe | TBA | ✅︎ | ✅︎ | +| `AliceAIForCausalLM` | AliceAI | `yandex/AliceAI-Foundation-80B-A3B-Base` | ✅︎ | | | `ApertusForCausalLM` | Apertus | `swiss-ai/Apertus-8B-2509`, `swiss-ai/Apertus-70B-Instruct-2509`, etc. | ✅︎ | ✅︎ | | `AXK1ForCausalLM` | A.X-K1 | `skt/A.X-K1`, etc. | | ✅︎ | | `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | ✅︎ | ✅︎ | diff --git a/tests/models/registry.py b/tests/models/registry.py index 76e49aec33..febd9d1a83 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -197,6 +197,10 @@ class _HfExamplesInfo: _TEXT_GENERATION_EXAMPLE_MODELS = { # [Decoder-only] "AfmoeForCausalLM": _HfExamplesInfo("arcee-ai/Trinity-Nano-Preview"), + "AliceAIForCausalLM": _HfExamplesInfo( + "yandex/AliceAI-Foundation-80B-A3B-Base", + trust_remote_code=True, + ), "ApertusForCausalLM": _HfExamplesInfo("swiss-ai/Apertus-8B-Instruct-2509"), "ArceeForCausalLM": _HfExamplesInfo("arcee-ai/AFM-4.5B-Base"), "AXK1ForCausalLM": _HfExamplesInfo("skt/A.X-K1", trust_remote_code=True), diff --git a/tests/models/test_alice_ai.py b/tests/models/test_alice_ai.py new file mode 100644 index 0000000000..5d68e433de --- /dev/null +++ b/tests/models/test_alice_ai.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first +from vllm.model_executor.models import alice_ai as alice_ai_module +from vllm.model_executor.models.alice_ai import ( + AliceAIDecoderBlock, + AliceAIForCausalLM, + AliceAIModel, + _AttentionResidualFinal, + _load_split_kda_conv_weights, +) +from vllm.model_executor.offloader import PrefetchOffloader, UVAOffloader + + +@pytest.mark.parametrize( + "backend,error", + [ + ("uva", "requires UVA"), + ("prefetch", "does not support prefetch"), + ], +) +def test_unsupported_offloading_is_rejected(monkeypatch, backend, error): + from vllm.model_executor.offloader import uva + + monkeypatch.setattr(uva, "is_uva_available", lambda: False) + monkeypatch.setenv("VLLM_WEIGHT_OFFLOADING_DISABLE_UVA", "0") + if backend == "uva": + offloader = UVAOffloader(cpu_offload_max_bytes=1024) + else: + offloader = PrefetchOffloader.__new__(PrefetchOffloader) + monkeypatch.setattr(alice_ai_module, "get_offloader", lambda: offloader) + config = SimpleNamespace(lora_config=None) + + with pytest.raises(NotImplementedError, match=error): + AliceAIForCausalLM(vllm_config=config, prefix="target") + + +@pytest.mark.parametrize( + "source,target,shard_id", + [ + ("linear_attn.q_proj.weight", "linear_attn.in_proj_qkvgfab.weight", 0), + ("linear_attn.k_proj.weight", "linear_attn.in_proj_qkvgfab.weight", 1), + ("linear_attn.v_proj.weight", "linear_attn.in_proj_qkvgfab.weight", 2), + ("linear_attn.b_proj.weight", "linear_attn.in_proj_qkvgfab.weight", 3), + ("linear_attn.f_a_proj.weight", "linear_attn.in_proj_qkvgfab.weight", 4), + ("self_attn.q_proj.weight", "self_attn.qkv_proj.weight", "q"), + ("linear_attn.a_log_bias", "linear_attn.A_log", None), + ], +) +def test_weight_mapping_preserves_shard_ids(source, target, shard_id): + tensor = torch.empty(1) + [(mapped, mapped_tensor)] = list( + AliceAIModel.hf_to_vllm_mapper.apply([(f"model.layers.0.{source}", tensor)]) + ) + assert mapped == f"model.layers.0.{target}" + assert getattr(mapped_tensor, "shard_id", None) == shard_id + + +def test_split_kda_conv_weights_use_fused_loader_shard_ids() -> None: + calls: list[tuple[torch.Tensor, int]] = [] + param = torch.nn.Parameter(torch.empty(3, 1, 1)) + + def weight_loader( + target: torch.Tensor, + weight: torch.Tensor, + shard_id: int, + ) -> None: + assert target is param + calls.append((weight, shard_id)) + + param.weight_loader = weight_loader + target_name = "model.layers.0.linear_attn.conv1d.weight" + weights = [ + ("model.layers.0.linear_attn.q_conv1d.weight", torch.full((1, 1), 1.0)), + ("model.layers.0.linear_attn.k_conv1d.weight", torch.full((1, 1), 2.0)), + ("model.layers.0.linear_attn.v_conv1d.weight", torch.full((1, 1), 3.0)), + ("model.embed_tokens.weight", torch.full((1, 1), 4.0)), + ] + loaded: set[str] = set() + + remaining = list( + _load_split_kda_conv_weights(weights, {target_name: param}, loaded) + ) + + assert loaded == {target_name} + assert [shard_id for _, shard_id in calls] == [0, 1, 2] + assert [weight.item() for weight, _ in calls] == [1.0, 2.0, 3.0] + assert remaining == [weights[-1]] + + +def test_kda_cache_contract_is_fp32_recurrent() -> None: + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + dtype=torch.bfloat16, + hf_text_config=SimpleNamespace( + linear_num_key_heads=32, + linear_key_head_dim=128, + linear_conv_kernel_dim=4, + ), + ), + cache_config=SimpleNamespace(mamba_cache_dtype="auto"), + parallel_config=SimpleNamespace(tensor_parallel_size=1), + speculative_config=None, + ) + + assert AliceAIForCausalLM.get_mamba_state_dtype_from_config(vllm_config) == ( + torch.bfloat16, + torch.float32, + ) + conv_shape = (12288, 3) if is_conv_state_dim_first() else (3, 12288) + assert AliceAIForCausalLM.get_mamba_state_shape_from_config(vllm_config) == ( + conv_shape, + (32, 128, 128), + ) + + +class _FixedOutput(torch.nn.Module): + def __init__(self, value: float) -> None: + super().__init__() + self.value = value + self.inputs: list[torch.Tensor] = [] + + def forward( + self, hidden_states: torch.Tensor, positions: torch.Tensor | None = None + ) -> torch.Tensor: + del positions + self.inputs.append(hidden_states.detach().clone()) + return torch.full_like(hidden_states, self.value) + + +def _make_state_layer( + layer_idx: int, + block_size: int, + attention_value: float, + mlp_value: float, +) -> AliceAIDecoderBlock: + layer = AliceAIDecoderBlock.__new__(AliceAIDecoderBlock) + torch.nn.Module.__init__(layer) + layer.layer_idx = layer_idx + layer.layer_type = "linear_attention" + layer.layer_scale = False + layer.attn_res_block_size = block_size + layer.attn_res_rms_norm_eps = 1e-6 + layer.input_layernorm = torch.nn.Identity() + layer.post_attention_layernorm = torch.nn.Identity() + layer.linear_attn = _FixedOutput(attention_value) + layer.mlp = _FixedOutput(mlp_value) + if layer_idx > 0: + layer.attn_res_proj = torch.nn.Linear(1, 1, bias=False) + layer.attn_res_norm_weight = torch.nn.Parameter(torch.ones(1)) + layer.mlp_res_proj = torch.nn.Linear(1, 1, bias=False) + layer.mlp_res_norm_weight = torch.nn.Parameter(torch.ones(1)) + return layer + + +def _make_final_attn_res() -> _AttentionResidualFinal: + final = _AttentionResidualFinal.__new__(_AttentionResidualFinal) + torch.nn.Module.__init__(final) + final.rms_norm_eps = 1e-6 + final.res_proj = torch.nn.Linear(1, 1, bias=False) + final.res_norm_weight = torch.nn.Parameter(torch.ones(1)) + return final + + +def test_attn_res_block_grouping_and_final_norm( + monkeypatch, +) -> None: + calls: list[list[torch.Tensor]] = [] + + def reference_attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + block_write_idx, + eps, + output_norm_eps, + ): + del delta, norm_weight, qk_weight, output_norm_weight + del block_write_idx, eps, output_norm_eps + sources = [blocks[:, index] for index in range(num_blocks)] + [prefix] + calls.append([source.detach().clone() for source in sources]) + return torch.stack(sources).sum(dim=0) + + monkeypatch.setattr(alice_ai_module, "attn_res", reference_attn_res) + monkeypatch.setattr( + alice_ai_module, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + + raw_outputs = [(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)] + model = AliceAIModel.__new__(AliceAIModel) + torch.nn.Module.__init__(model) + model.attn_res_block_size = 4 + model.start_layer = 0 + model.end_layer = len(raw_outputs) + model.layers = torch.nn.ModuleList( + [ + _make_state_layer(index, 4, attention, mlp) + for index, (attention, mlp) in enumerate(raw_outputs) + ] + ) + model.config = SimpleNamespace(num_hidden_layers=len(raw_outputs)) + model.attnres_final = _make_final_attn_res() + model.norm = torch.nn.Linear(1, 1, bias=False) + with torch.no_grad(): + model.norm.weight.fill_(2) + + output = AliceAIModel.forward( + model, + input_ids=None, + positions=torch.arange(2), + inputs_embeds=torch.full((2, 1), 10.0), + ) + + torch.testing.assert_close(output, torch.full((2, 1), 130.0)) + torch.testing.assert_close( + model.layers[0].linear_attn.inputs[0], torch.full((2, 1), 10.0) + ) + assert [len(sources) for sources in calls] == [2] * 8 + [3, 3] + assert [source[0, 0].item() for source in calls[-1]] == [10, 36, 19] diff --git a/tests/models/test_alice_ai_config.py b/tests/models/test_alice_ai_config.py new file mode 100644 index 0000000000..ede3b59622 --- /dev/null +++ b/tests/models/test_alice_ai_config.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU-only tests for AliceAI configuration.""" + +from types import SimpleNamespace + +import pytest + +from vllm.model_executor.models.config import AliceAIForCausalLMConfig + + +def _make_config(**overrides) -> SimpleNamespace: + return SimpleNamespace( + **{ + "kda_allow_negative_eigenvalues": False, + "router_score_function": "sigmoid", + "router_bias_correction": True, + **overrides, + } + ) + + +def test_model_config_warns_when_mtp_is_ignored(caplog) -> None: + config = _make_config(mtp_num_hidden_layers=1) + + with caplog.at_level("WARNING"): + AliceAIForCausalLMConfig.verify_and_update_model_config( + SimpleNamespace(hf_config=config) + ) + assert "ignores the MTP configuration and weights" in caplog.text + + +@pytest.mark.parametrize( + "overrides,error_type,error", + [ + ({}, None, None), + ( + {"kda_allow_negative_eigenvalues": True}, + NotImplementedError, + "kda_allow_negative_eigenvalues=false", + ), + ( + {"router_score_function": "softmax"}, + ValueError, + "router_score_function='sigmoid'", + ), + ( + {"router_bias_correction": False}, + ValueError, + "router_bias_correction=true", + ), + ], +) +def test_model_config_rejects_unsupported_alice_contract( + overrides, error_type, error +) -> None: + model_config = SimpleNamespace(hf_config=_make_config(**overrides)) + + if error is None: + AliceAIForCausalLMConfig.verify_and_update_model_config(model_config) + else: + with pytest.raises(error_type, match=error): + AliceAIForCausalLMConfig.verify_and_update_model_config(model_config) + + +@pytest.mark.parametrize( + "pipeline_parallel_size,speculative_method,error", + [ + (1, None, None), + (2, None, "pipeline_parallel_size=1"), + (1, "eagle3", "eagle3"), + (1, "dflash", "dflash"), + (1, "dspark", "dspark"), + (1, "extract_hidden_states", "extract_hidden_states"), + ], +) +def test_runtime_config_rejects_unsupported_alice_modes( + pipeline_parallel_size, speculative_method, error +) -> None: + vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(pipeline_parallel_size=pipeline_parallel_size), + speculative_config=( + None + if speculative_method is None + else SimpleNamespace(method=speculative_method) + ), + ) + + if error is None: + AliceAIForCausalLMConfig.verify_and_update_config(vllm_config) + else: + with pytest.raises(NotImplementedError, match=error): + AliceAIForCausalLMConfig.verify_and_update_config(vllm_config) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index f587ffac5b..837bf1e80a 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -73,6 +73,7 @@ logger = init_logger(__name__) DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset( { + "AliceAIForCausalLM", "DeepseekV32MTPModel", "DeepseekV32ForCausalLM", "DeepseekV4ForCausalLM", diff --git a/vllm/model_executor/models/alice_ai.py b/vllm/model_executor/models/alice_ai.py new file mode 100644 index 0000000000..a0c5f53ecf --- /dev/null +++ b/vllm/model_executor/models/alice_ai.py @@ -0,0 +1,733 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only AliceAI model.""" + +from collections.abc import Iterable +from copy import copy +from itertools import islice +from typing import cast + +import torch +from torch import nn +from transformers import PreTrainedConfig, Qwen3NextConfig + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import ( + get_ep_group, + get_pp_group, + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory, GateLinear +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, + resolve_layer_fused_shared_expert, +) +from vllm.model_executor.layers.layernorm import GemmaRMSNorm +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import ( + KimiGatedDeltaNetAttention, +) +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, + MambaStateDtypeCalculator, + MambaStateShapeCalculator, +) +from vllm.model_executor.layers.quantization.utils.config_utils import ( + get_quark_ocp_mx_group_size, +) +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.interfaces import ( + HasInnerState, + IsHybrid, + MixtureOfExperts, + SupportsLoRA, +) +from vllm.model_executor.models.qwen2_moe import Qwen2MoeMLP +from vllm.model_executor.models.qwen3_next import ( + Qwen3NextAttention, + _should_replicate_misaligned_shared_expert, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + extract_layer_index, + make_layers, + maybe_fuse_shared_experts, + maybe_prefix, +) +from vllm.model_executor.offloader import ( + PrefetchOffloader, + UVAOffloader, + get_offloader, +) +from vllm.models.kimi_k3.nvidia.ops import attn_res +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig + +_SPLIT_KDA_CONV_SHARDS = { + ".linear_attn.q_conv1d": (".linear_attn.conv1d", 0), + ".linear_attn.k_conv1d": (".linear_attn.conv1d", 1), + ".linear_attn.v_conv1d": (".linear_attn.conv1d", 2), +} + + +def _load_split_kda_conv_weights( + weights: Iterable[tuple[str, torch.Tensor]], + params: dict[str, nn.Parameter], + loaded: set[str], +) -> Iterable[tuple[str, torch.Tensor]]: + """Load split AliceAI conv weights through the upstream fused-conv ABI.""" + for name, weight in weights: + for source_name, (target_name, shard_id) in _SPLIT_KDA_CONV_SHARDS.items(): + if source_name not in name: + continue + mapped_name = name.replace(source_name, target_name, 1) + param = params.get(mapped_name) + if param is None: + raise ValueError(f"No target parameter for {name!r}: {mapped_name!r}") + weight_loader = getattr(param, "weight_loader", None) + if weight_loader is None: + raise ValueError( + f"Target parameter {mapped_name!r} has no weight loader" + ) + weight_loader(param, weight, shard_id) + loaded.add(mapped_name) + break + else: + yield name, weight + + +def get_attn_res_block_size(hf_config: object) -> int: + """Return the architecture-rebase AttnRes grouping, not legacy metadata.""" + block_size = getattr(hf_config, "block_attn_res_block_size", None) + if isinstance(block_size, bool) or not isinstance(block_size, int): + raise ValueError( + "AliceAI requires a positive integer block_attn_res_block_size" + ) + if block_size <= 0: + raise ValueError("block_attn_res_block_size must be positive") + return block_size + + +def _make_kimi_kda_config(config: PreTrainedConfig) -> KimiLinearConfig: + """Adapt the flat AliceAI KDA geometry to the existing upstream Kimi layer.""" + num_heads = config.linear_num_key_heads + num_value_heads = config.linear_num_value_heads + head_dim = config.linear_key_head_dim + value_head_dim = config.linear_value_head_dim + if num_heads != num_value_heads or head_dim != value_head_dim: + raise ValueError( + "AliceAI KDA requires matching key/value head counts and dimensions" + ) + + adapted = copy(config) + adapted.linear_attn_config = { + "head_dim": head_dim, + "num_heads": num_heads, + "short_conv_kernel_size": config.linear_conv_kernel_dim, + "use_full_rank_gate": False, + "gate_lower_bound": None, + } + return cast(KimiLinearConfig, adapted) + + +def _make_alice_ai_full_attention_config( + config: PreTrainedConfig, +) -> Qwen3NextConfig: + adapted = copy(config) + adapted.attn_output_gate = True + adapted.qkv_bias = False + adapted.dual_chunk_attention_config = None + adapted.is_causal = True + return cast(Qwen3NextConfig, adapted) + + +class _AttentionResidualFinal(nn.Module): + def __init__(self, hidden_size: int, rms_norm_eps: float, prefix: str) -> None: + super().__init__() + self.rms_norm_eps = rms_norm_eps + self.res_proj = ReplicatedLinear( + input_size=hidden_size, + output_size=1, + bias=False, + prefix=f"{prefix}.res_proj", + ) + self.res_norm_weight = nn.Parameter(torch.ones(hidden_size)) + + def forward( + self, + residual_bank: torch.Tensor, + num_completed_sources: int, + partial: torch.Tensor | None, + ) -> torch.Tensor: + if partial is None: + if num_completed_sources < 1: + raise RuntimeError("AttnRes requires at least one residual source") + prefix = residual_bank[:, num_completed_sources - 1] + num_blocks = num_completed_sources - 1 + else: + prefix = partial + num_blocks = num_completed_sources + return attn_res( + prefix=prefix, + delta=None, + blocks=residual_bank, + norm_weight=self.res_norm_weight, + qk_weight=self.res_proj.weight.squeeze(0), + output_norm_weight=None, + num_blocks=num_blocks, + block_write_idx=-1, + eps=self.rms_norm_eps, + output_norm_eps=0.0, + ) + + +class AliceAISparseMoeBlock(nn.Module): + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + + config = vllm_config.model_config.hf_text_config + parallel_config = vllm_config.parallel_config + quant_config = vllm_config.quant_config + + self.tp_size = get_tensor_model_parallel_world_size() + self.ep_size = get_ep_group().device_group.size() + self.n_routed_experts = config.num_experts + if self.tp_size > self.n_routed_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {self.n_routed_experts}." + ) + + self.is_fused_shared_expert_enabled = False + if config.shared_expert_intermediate_size > 0: + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert( + quant_config=quant_config, + prefix=prefix, + shared_expert_name="shared_expert", + ) + + self.replicate_shared_expert = False + if not self.is_fused_shared_expert_enabled: + shared_expert_group_size = get_quark_ocp_mx_group_size( + quant_config=quant_config, + layer_name=f"{prefix}.shared_expert.down_proj", + ) + self.replicate_shared_expert = _should_replicate_misaligned_shared_expert( + intermediate_size=config.shared_expert_intermediate_size, + tp_size=self.tp_size, + group_size=shared_expert_group_size, + enable_expert_parallel=parallel_config.enable_expert_parallel, + is_sequence_parallel=False, + ) + + self.enable_eplb = parallel_config.enable_eplb + self.n_logical_experts = self.n_routed_experts + self.n_redundant_experts = parallel_config.eplb_config.num_redundant_experts + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts + self.n_local_physical_experts = self.n_physical_experts // self.ep_size + + self.gate = GateLinear( + input_size=config.hidden_size, + output_size=config.num_experts, + prefix=f"{prefix}.gate", + ) + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(config.num_experts, dtype=torch.float32), + requires_grad=False, + ) + self.shared_expert_gate = ReplicatedLinear( + input_size=config.hidden_size, + output_size=1, + bias=False, + prefix=f"{prefix}.shared_expert_gate", + ) + + if ( + self.is_fused_shared_expert_enabled + or config.shared_expert_intermediate_size <= 0 + ): + self.shared_expert = None + else: + self.shared_expert = Qwen2MoeMLP( + hidden_size=config.hidden_size, + intermediate_size=config.shared_expert_intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + expert_gate=self.shared_expert_gate, + disable_tp=self.replicate_shared_expert, + prefix=f"{prefix}.shared_expert", + ) + + self.experts = FusedMoEFactory( + shared_experts=( + None if self.replicate_shared_expert else self.shared_expert + ), + gate=self.gate, + num_experts=self.n_routed_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=getattr(config, "norm_topk_prob", True), + scoring_func="sigmoid", + e_score_correction_bias=self.gate.e_score_correction_bias, + quant_config=quant_config, + prefix=f"{prefix}.experts", + enable_eplb=self.enable_eplb, + num_redundant_experts=self.n_redundant_experts, + n_shared_experts=1 if self.shared_expert is None else None, + fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_expert", + shared_expert_gate=self.shared_expert_gate + if self.shared_expert is None + else None, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + orig_shape = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + shared_output = ( + self.shared_expert(hidden_states) + if self.replicate_shared_expert and self.shared_expert is not None + else None + ) + output = self.experts(hidden_states=hidden_states, router_logits=hidden_states) + if shared_output is not None: + output += shared_output + return output.view(orig_shape) + + +class AliceAIDecoderBlock(nn.Module): + """AliceAI topology with upstream Qwen, Kimi KDA, MoE and AttnRes operations.""" + + def __init__( + self, + vllm_config: VllmConfig, + layer_type: str, + prefix: str = "", + ) -> None: + super().__init__() + + config = vllm_config.model_config.hf_text_config + model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + parallel_config = vllm_config.parallel_config + quant_config = vllm_config.quant_config + + if config.num_experts <= 0: + raise ValueError("AliceAI requires sparse experts") + if parallel_config.use_sequence_parallel_moe: + raise NotImplementedError("AliceAI does not support sequence-parallel MoE") + + self.layer_type = layer_type + self.layer_idx = extract_layer_index(prefix) + self.attn_res_block_size = get_attn_res_block_size(config) + self.attn_res_rms_norm_eps = float(config.rms_norm_eps) + + if layer_type == "linear_attention": + self.linear_attn = KimiGatedDeltaNetAttention( + config=_make_kimi_kda_config(config), + vllm_config=vllm_config, + prefix=f"{prefix}.linear_attn", + ) + self.linear_attn.o_norm.eps = float(config.rms_norm_eps) + elif layer_type == "full_attention": + self.self_attn = Qwen3NextAttention( + config=_make_alice_ai_full_attention_config(config), + model_config=model_config, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + else: + raise ValueError(f"Invalid layer_type {layer_type}") + + self.mlp = AliceAISparseMoeBlock( + vllm_config=vllm_config, + prefix=f"{prefix}.mlp", + ) + self.input_layernorm = GemmaRMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = GemmaRMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps + ) + + if self.layer_idx > 0: + self.attn_res_proj = ReplicatedLinear( + input_size=config.hidden_size, + output_size=1, + bias=False, + prefix=f"{prefix}.attn_res_proj", + ) + self.attn_res_norm_weight = nn.Parameter(torch.ones(config.hidden_size)) + self.mlp_res_proj = ReplicatedLinear( + input_size=config.hidden_size, + output_size=1, + bias=False, + prefix=f"{prefix}.mlp_res_proj", + ) + self.mlp_res_norm_weight = nn.Parameter(torch.ones(config.hidden_size)) + + def _apply_attention_residual( + self, + residual_bank: torch.Tensor, + num_completed_sources: int, + partial: torch.Tensor | None, + *, + projection: ReplicatedLinear, + norm_weight_name: str, + ) -> torch.Tensor: + if partial is None: + if num_completed_sources < 1: + raise RuntimeError("AttnRes requires at least one residual source") + prefix = residual_bank[:, num_completed_sources - 1] + num_blocks = num_completed_sources - 1 + else: + prefix = partial + num_blocks = num_completed_sources + return attn_res( + prefix=prefix, + delta=None, + blocks=residual_bank, + norm_weight=getattr(self, norm_weight_name), + qk_weight=projection.weight.squeeze(0), + output_norm_weight=None, + num_blocks=num_blocks, + block_write_idx=-1, + eps=self.attn_res_rms_norm_eps, + output_norm_eps=0.0, + ) + + def _apply_attention( + self, hidden_states: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + if self.layer_type == "linear_attention": + return self.linear_attn(hidden_states=hidden_states, positions=positions) + return self.self_attn(hidden_states=hidden_states, positions=positions) + + def forward_attn_res( + self, + positions: torch.Tensor, + residual_bank: torch.Tensor, + num_completed_sources: int, + partial: torch.Tensor | None, + ) -> tuple[int, torch.Tensor]: + if self.layer_idx > 0 and self.layer_idx % self.attn_res_block_size == 0: + if partial is None: + raise RuntimeError("cannot finalize an empty AttnRes block") + residual_bank[:, num_completed_sources].copy_(partial) + num_completed_sources += 1 + partial = None + + if self.layer_idx == 0: + hidden_states = residual_bank[:, 0] + else: + hidden_states = self._apply_attention_residual( + residual_bank=residual_bank, + num_completed_sources=num_completed_sources, + partial=partial, + projection=self.attn_res_proj, + norm_weight_name="attn_res_norm_weight", + ) + hidden_states = self.input_layernorm(hidden_states) + attention_output = self._apply_attention( + hidden_states=hidden_states, positions=positions + ) + partial = attention_output if partial is None else partial + attention_output + + hidden_states = self._apply_attention_residual( + residual_bank=residual_bank, + num_completed_sources=num_completed_sources, + partial=partial, + projection=self.mlp_res_proj, + norm_weight_name="mlp_res_norm_weight", + ) + hidden_states = self.post_attention_layernorm(hidden_states) + mlp_output = self.mlp(hidden_states) + return num_completed_sources, partial + mlp_output + + +@support_torch_compile +class AliceAIModel(nn.Module): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + ".linear_attn.a_log_bias": ".linear_attn.A_log", + }, + orig_to_new_stacked={ + ".linear_attn.q_proj": (".linear_attn.in_proj_qkvgfab", 0), + ".linear_attn.k_proj": (".linear_attn.in_proj_qkvgfab", 1), + ".linear_attn.v_proj": (".linear_attn.in_proj_qkvgfab", 2), + ".linear_attn.b_proj": (".linear_attn.in_proj_qkvgfab", 3), + ".linear_attn.f_a_proj": (".linear_attn.in_proj_qkvgfab", 4), + ".self_attn.q_proj": (".self_attn.qkv_proj", "q"), + ".self_attn.k_proj": (".self_attn.qkv_proj", "k"), + ".self_attn.v_proj": (".self_attn.qkv_proj", "v"), + ".shared_expert.gate_proj": (".shared_expert.gate_up_proj", 0), + ".shared_expert.up_proj": (".shared_expert.gate_up_proj", 1), + }, + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config = vllm_config.model_config.hf_text_config + self.config = config + self.vocab_size = config.vocab_size + self.num_redundant_experts = ( + vllm_config.parallel_config.eplb_config.num_redundant_experts + ) + self.attn_res_block_size = get_attn_res_block_size(config) + self.embed_tokens = VocabParallelEmbedding( + num_embeddings=config.vocab_size, embedding_dim=config.hidden_size + ) + + def get_layer(prefix: str) -> AliceAIDecoderBlock: + return AliceAIDecoderBlock( + vllm_config=vllm_config, + layer_type=config.layer_types[extract_layer_index(prefix)], + prefix=prefix, + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + num_hidden_layers=config.num_hidden_layers, + layer_fn=get_layer, + prefix=f"{prefix}.layers", + ) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + layers=self.layers, moe_cls=AliceAISparseMoeBlock, moe_name="mlp" + ) + self.norm = GemmaRMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps + ) + self.attnres_final = _AttentionResidualFinal( + hidden_size=config.hidden_size, + rms_norm_eps=float(config.rms_norm_eps), + prefix=f"{prefix}.attnres_final", + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def _init_residual_bank(self, hidden_states: torch.Tensor) -> torch.Tensor: + residual_bank = hidden_states.new_empty( + ( + hidden_states.shape[0], + (self.config.num_hidden_layers + self.attn_res_block_size - 1) + // self.attn_res_block_size, + hidden_states.shape[-1], + ) + ) + residual_bank[:, 0].copy_(hidden_states) + return residual_bank + + 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 intermediate_tensors is not None: + raise NotImplementedError("AliceAI does not support pipeline parallelism") + if not get_pp_group().is_first_rank or not get_pp_group().is_last_rank: + raise NotImplementedError("AliceAI requires pipeline_parallel_size=1") + + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + assert input_ids is not None + hidden_states = self.embed_input_ids(input_ids) + residual_bank = self._init_residual_bank(hidden_states) + + num_completed_sources = 1 + partial = None + for layer in islice(self.layers, self.start_layer, self.end_layer): + num_completed_sources, partial = layer.forward_attn_res( + positions=positions, + residual_bank=residual_bank, + num_completed_sources=num_completed_sources, + partial=partial, + ) + if partial is None: + raise RuntimeError("AttnRes produced an empty final block") + hidden_states = self.attnres_final( + residual_bank=residual_bank, + num_completed_sources=num_completed_sources, + partial=partial, + ) + return self.norm(hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + weights = maybe_fuse_shared_experts( + weights=weights, + enabled=self.is_fused_shared_expert_enabled, + n_routed_experts=self.config.num_experts, + n_shared_experts=1, + ckpt_prefix="mlp.shared_expert", + ) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights=weights, mapper=self.hf_to_vllm_mapper) + + +class AliceAIForCausalLM( + nn.Module, HasInnerState, IsHybrid, SupportsLoRA, MixtureOfExperts +): + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"mtp.": None}) + + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + "in_proj_qkvgfab": [ + "q_proj", + "k_proj", + "v_proj", + "b_proj", + "f_a_proj", + ], + "conv1d": ["q_conv1d", "k_conv1d", "v_conv1d"], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + # AttnRes calls bypass the per-layer forward hooks used by offloading. + offloader = get_offloader() + if isinstance(offloader, PrefetchOffloader): + raise NotImplementedError("AliceAI does not support prefetch offloading") + if ( + isinstance(offloader, UVAOffloader) + and offloader.cpu_offload_max_bytes > 0 + and not offloader.uva_offloading + ): + raise NotImplementedError("AliceAI CPU offloading requires UVA") + + super().__init__() + config = vllm_config.model_config.hf_text_config + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config + self.scheduler_config = vllm_config.scheduler_config + self.quant_config = vllm_config.quant_config + self.config = config + self.model = AliceAIModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.lm_head = ParallelLMHead( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + self.set_moe_parameters() + + def set_moe_parameters(self) -> None: + self.moe_layers = [layer.mlp.experts for layer in self.model.layers] + moe = self.model.layers[self.model.start_layer].mlp + self.num_moe_layers = len(self.moe_layers) + self.num_expert_groups = 1 + self.num_shared_experts = 0 + self.num_logical_experts = moe.n_logical_experts + self.num_physical_experts = moe.n_physical_experts + self.num_local_physical_experts = moe.n_local_physical_experts + self.num_routed_experts = moe.n_routed_experts + self.num_redundant_experts = moe.n_redundant_experts + + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for layer in self.model.layers: + moe = layer.mlp + moe.n_local_physical_experts = num_local_physical_experts + moe.n_physical_experts = num_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() + + 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: object, + ) -> torch.Tensor | IntermediateTensors: + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(lm_head=self.lm_head, hidden_states=hidden_states) + + @classmethod + def get_mamba_state_dtype_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.kda_state_dtype( + model_dtype=vllm_config.model_config.dtype, + mamba_cache_dtype=vllm_config.cache_config.mamba_cache_dtype, + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + config = vllm_config.model_config.hf_text_config + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + return MambaStateShapeCalculator.kda_state_shape( + tp_world_size=vllm_config.parallel_config.tensor_parallel_size, + num_heads=config.linear_num_key_heads, + head_dim=config.linear_key_head_dim, + conv_kernel_size=config.linear_conv_kernel_dim, + num_spec=num_spec, + ) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + weights = self.hf_to_vllm_mapper.apply(weights) + manually_loaded: set[str] = set() + weights = _load_split_kda_conv_weights( + weights=weights, + params=dict(self.named_parameters()), + loaded=manually_loaded, + ) + loader = AutoWeightsLoader(self) + loaded = loader.load_weights(weights=weights) + loaded.update(manually_loaded) + expected_target_weights = { + name + for name, _ in self.named_parameters() + if ".attnres_final." in name + or ".attn_res_" in name + or ".mlp_res_" in name + or name.endswith(".e_score_correction_bias") + } + if missing := expected_target_weights - loaded: + raise ValueError( + "AliceAI checkpoint is missing target-specific weights: " + + ", ".join(sorted(missing)) + ) + return loaded diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 4aec389062..f796bcf615 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -25,6 +25,51 @@ class VerifyAndUpdateConfig: return +class AliceAIForCausalLMConfig(VerifyAndUpdateConfig): + @staticmethod + def verify_and_update_model_config(model_config: "ModelConfig") -> None: + hf_config = model_config.hf_config + if getattr(hf_config, "kda_allow_negative_eigenvalues", None) is not False: + raise NotImplementedError( + "AliceAI requires kda_allow_negative_eigenvalues=false" + ) + if getattr(hf_config, "router_score_function", None) != "sigmoid": + raise ValueError("AliceAI requires router_score_function='sigmoid'") + if getattr(hf_config, "router_bias_correction", None) is not True: + raise ValueError("AliceAI requires router_bias_correction=true") + mtp_layers = getattr(hf_config, "mtp_num_hidden_layers", 0) + if mtp_layers: + logger.warning( + "AliceAI checkpoint declares %s MTP layer(s), but this " + "implementation ignores the MTP configuration and weights; " + "MTP speculative decoding is unavailable.", + mtp_layers, + ) + + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + if vllm_config.parallel_config.pipeline_parallel_size != 1: + raise NotImplementedError( + "AliceAI does not support pipeline parallelism; " + "pipeline_parallel_size=1 is required" + ) + + speculative_config = vllm_config.speculative_config + speculative_method = ( + speculative_config.method if speculative_config is not None else None + ) + if speculative_method in { + "dflash", + "dspark", + "eagle3", + "extract_hidden_states", + }: + raise NotImplementedError( + "AliceAI does not support speculative method " + f"{speculative_method!r} because it requires auxiliary hidden states" + ) + + class DeepseekV32ForCausalLM(VerifyAndUpdateConfig): @classmethod def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: @@ -1135,6 +1180,7 @@ class LongcatFlashNgramForCausalLMConfig(VerifyAndUpdateConfig): MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { + "AliceAIForCausalLM": AliceAIForCausalLMConfig, "ColBERTJinaRobertaModel": JinaRobertaModelConfig, "ColQwen3_5": ColQwen3_5Config, "D3ForSequenceClassification": D3ForSequenceClassificationConfig, diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 9bb811f203..23e67c9f50 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -71,6 +71,7 @@ logger = init_logger(__name__) _TEXT_GENERATION_MODELS = { # [Decoder-only] "AfmoeForCausalLM": ("afmoe", "AfmoeForCausalLM"), + "AliceAIForCausalLM": ("alice_ai", "AliceAIForCausalLM"), "ApertusForCausalLM": ("apertus", "ApertusForCausalLM"), "AXK1ForCausalLM": ("AXK1", "AXK1ForCausalLM"), "BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"),