vllm.model_executor.models.alice_ai

Inference-only AliceAI model.

Classes:

• AliceAIDecoderBlock –

  AliceAI topology with upstream Qwen, Kimi KDA, MoE and AttnRes operations.

Functions:

• get_attn_res_block_size –

  Return the architecture-rebase AttnRes grouping, not legacy metadata.

AliceAIDecoderBlock

Bases: Module

AliceAI topology with upstream Qwen, Kimi KDA, MoE and AttnRes operations.

Source code in vllm/model_executor/models/alice_ai.py

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

_load_split_kda_conv_weights(weights, params, loaded)

Load split AliceAI conv weights through the upstream fused-conv ABI.

Source code in vllm/model_executor/models/alice_ai.py

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

_make_kimi_kda_config(config)

Adapt the flat AliceAI KDA geometry to the existing upstream Kimi layer.

Source code in vllm/model_executor/models/alice_ai.py

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)

get_attn_res_block_size(hf_config)

Return the architecture-rebase AttnRes grouping, not legacy metadata.

Source code in vllm/model_executor/models/alice_ai.py

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