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

[Model] Add AliceAI support (#57963)

Browse Source 
Signed-off-by: timka <timka.byzov@yandex.ru>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Timofey
2026-10-15 10:15:03 +00:00
committed by GitHub
co-authored by Codex
parent 7020600b1d
commit 716c5225ae
9 changed files with 1113 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/models_basic.yaml
+3 -1
View File
@@ -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
docs/models/supported_models.md
+1
View File
@@ -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. | ✅︎ | ✅︎ |
tests/models/registry.py
+4
View File
@@ -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),
tests/models/test_alice_ai.py
+231
View File
@@ -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]
tests/models/test_alice_ai_config.py
+93
View File
@@ -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)
vllm/config/vllm.py
+1
View File
@@ -73,6 +73,7 @@ logger = init_logger(__name__)
DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(
{
"AliceAIForCausalLM",
"DeepseekV32MTPModel",
"DeepseekV32ForCausalLM",
"DeepseekV4ForCausalLM",
vllm/model_executor/models/alice_ai.py
+733
View File
@@ -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
vllm/model_executor/models/config.py
+46
View File
@@ -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,
vllm/model_executor/models/registry.py
+1
View File
@@ -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"),