modeling_alice_ai.py

from __future__ import annotations

import math
from typing import Any, ClassVar

import torch
import torch.nn.functional as functional
from torch import nn
from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation import GenerationMixin
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.modeling_outputs import (
    MoeCausalLMOutputWithPast,
    MoeModelOutputWithPast,
)
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
from transformers.utils import can_return_tuple
from transformers.utils.generic import merge_with_config_defaults
from transformers.utils.output_capturing import capture_outputs

from .configuration_alice_ai import AliceAIConfig


def _rms_norm(
    hidden_states: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    *,
    zero_centered: bool = False,
) -> torch.Tensor:
    dtype = hidden_states.dtype
    normalized = hidden_states.float() * torch.rsqrt(
        hidden_states.float().square().mean(dim=-1, keepdim=True) + eps
    )
    scale = 1.0 + weight.float() if zero_centered else weight.float()
    return (normalized * scale).to(dtype)


class AliceAIRMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float, *, zero_centered: bool) -> None:
        super().__init__()
        initial_value = 0.0 if zero_centered else 1.0
        self.weight = nn.Parameter(torch.full((hidden_size,), initial_value))
        self.eps = eps
        self.zero_centered = zero_centered

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return _rms_norm(
            hidden_states,
            self.weight,
            self.eps,
            zero_centered=self.zero_centered,
        )


def _depth_softmax_mix(
    sources: list[torch.Tensor],
    query_weight: torch.Tensor,
    norm_weight: torch.Tensor,
    eps: float,
) -> torch.Tensor:
    stacked = torch.stack(sources, dim=0)
    keys = _rms_norm(stacked, norm_weight, eps)
    scores = functional.linear(keys.float(), query_weight.float()).squeeze(-1)
    weights = scores.softmax(dim=0).to(stacked.dtype).unsqueeze(-1)
    return (stacked * weights).sum(dim=0)


class AliceAIFinalBlockAttnRes(nn.Module):
    def __init__(self, config: AliceAIConfig) -> None:
        super().__init__()
        self.res_proj = nn.Linear(config.hidden_size, 1, bias=False)
        self.res_norm_weight = nn.Parameter(torch.ones(config.hidden_size))
        self.eps = config.rms_norm_eps

    def forward(
        self, completed: list[torch.Tensor], partial: torch.Tensor | None = None
    ) -> torch.Tensor:
        sources = completed if partial is None else [*completed, partial]
        return _depth_softmax_mix(
            sources, self.res_proj.weight, self.res_norm_weight, self.eps
        )


class AliceAIRotaryEmbedding(nn.Module):
    def __init__(self, config: AliceAIConfig) -> None:
        super().__init__()
        rotary_dim = int(config.head_dim * config.partial_rotary_factor)
        if rotary_dim % 2:
            raise ValueError("The rotary dimension must be even")
        inv_freq = 1.0 / (
            config.rope_theta
            ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
        )
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(
        self, position_ids: torch.LongTensor, dtype: torch.dtype
    ) -> tuple[torch.Tensor, torch.Tensor]:
        frequencies = torch.einsum(
            "bi,j->bij", position_ids.float(), self.inv_freq.float()
        )
        embeddings = torch.cat((frequencies, frequencies), dim=-1)
        return embeddings.cos().to(dtype), embeddings.sin().to(dtype)


def _rotate_half(hidden_states: torch.Tensor) -> torch.Tensor:
    first, second = hidden_states.chunk(2, dim=-1)
    return torch.cat((-second, first), dim=-1)


def _apply_rotary(
    hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> torch.Tensor:
    rotary_dim = cos.shape[-1]
    rotary, remainder = hidden_states[..., :rotary_dim], hidden_states[..., rotary_dim:]
    cos = cos.unsqueeze(1)
    sin = sin.unsqueeze(1)
    rotary = rotary * cos + _rotate_half(rotary) * sin
    return torch.cat((rotary, remainder), dim=-1)


def eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    dropout: float = 0.0,
    scaling: float | None = None,
    **_: Any,
) -> tuple[torch.Tensor, torch.Tensor]:
    key = key.repeat_interleave(module.num_key_value_groups, dim=1)
    value = value.repeat_interleave(module.num_key_value_groups, dim=1)
    scores = torch.matmul(query.float(), key.float().transpose(-1, -2))
    scores = scores * (module.scaling if scaling is None else scaling)
    if attention_mask is not None:
        if attention_mask.dtype == torch.bool:
            scores = scores.masked_fill(~attention_mask, torch.finfo(scores.dtype).min)
        else:
            scores = scores + attention_mask
    probabilities = scores.softmax(dim=-1).to(query.dtype)
    probabilities = functional.dropout(
        probabilities, p=dropout, training=module.training
    )
    output = torch.matmul(probabilities, value).transpose(1, 2).contiguous()
    return output, probabilities


class AliceAIAttention(nn.Module):
    def __init__(self, config: AliceAIConfig, layer_idx: int) -> None:
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_kv_heads
        self.head_dim = config.head_dim
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout
        self.is_causal = True
        self.q_proj = nn.Linear(
            config.hidden_size, 2 * self.num_heads * self.head_dim, bias=False
        )
        self.k_proj = nn.Linear(
            config.hidden_size, self.num_kv_heads * self.head_dim, bias=False
        )
        self.v_proj = nn.Linear(
            config.hidden_size, self.num_kv_heads * self.head_dim, bias=False
        )
        self.o_proj = nn.Linear(
            self.num_heads * self.head_dim, config.hidden_size, bias=False
        )
        self.q_norm = AliceAIRMSNorm(
            self.head_dim, config.rms_norm_eps, zero_centered=True
        )
        self.k_norm = AliceAIRMSNorm(
            self.head_dim, config.rms_norm_eps, zero_centered=True
        )

    def _backend_mask(
        self,
        attention_mask: torch.Tensor | None,
        cache_position: torch.LongTensor,
        key_length: int,
        implementation: str,
    ) -> torch.Tensor | None:
        if attention_mask is not None and attention_mask.ndim == 4:
            return attention_mask[..., :key_length]
        padding_mask = (
            attention_mask[..., :key_length].to(torch.bool)
            if attention_mask is not None
            else None
        )
        if padding_mask is not None and bool(padding_mask.all()):
            padding_mask = None
        if implementation.startswith("flash_attention"):
            return padding_mask

        query_length = cache_position.shape[0]
        is_plain_prefill = query_length == key_length and cache_position[0] == 0
        is_single_token_decode = query_length == 1
        if (
            implementation == "sdpa"
            and padding_mask is None
            and (is_plain_prefill or is_single_token_decode)
        ):
            return None

        key_positions = torch.arange(key_length, device=cache_position.device)
        causal_mask = key_positions.view(1, 1, 1, -1) <= cache_position.view(
            1, 1, -1, 1
        )
        if padding_mask is not None:
            causal_mask = causal_mask & padding_mask[:, None, None, :]
        return causal_mask

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None,
        cache: Cache | None,
        cache_position: torch.LongTensor,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        batch_size, sequence_length, _ = hidden_states.shape
        query, output_gate = (
            self.q_proj(hidden_states)
            .view(batch_size, sequence_length, self.num_heads, 2 * self.head_dim)
            .chunk(2, dim=-1)
        )
        output_gate = output_gate.reshape(
            batch_size, sequence_length, self.num_heads * self.head_dim
        )
        key = self.k_proj(hidden_states).view(
            batch_size, sequence_length, self.num_kv_heads, self.head_dim
        )
        value = self.v_proj(hidden_states).view(
            batch_size, sequence_length, self.num_kv_heads, self.head_dim
        )
        query = self.q_norm(query).transpose(1, 2)
        key = self.k_norm(key).transpose(1, 2)
        value = value.transpose(1, 2)
        cos, sin = position_embeddings
        query = _apply_rotary(query, cos, sin)
        key = _apply_rotary(key, cos, sin)
        if cache is not None:
            key, value = cache.update(key, value, self.layer_idx)

        implementation = self.config._attn_implementation
        backend_mask = self._backend_mask(
            attention_mask, cache_position, key.shape[-2], implementation
        )
        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
            implementation, eager_attention_forward
        )
        output, attention_weights = attention_interface(
            self,
            query,
            key,
            value,
            backend_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            is_causal=True,
        )
        output = output.reshape(
            batch_size, sequence_length, self.num_heads * self.head_dim
        )
        output = output * torch.sigmoid(output_gate)
        return self.o_proj(output), attention_weights


class AliceAIKDA(nn.Module):
    def __init__(self, config: AliceAIConfig, layer_idx: int) -> None:
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.num_k_heads = config.linear_num_key_heads
        self.num_v_heads = config.linear_num_value_heads
        self.head_k_dim = config.linear_key_head_dim
        self.head_v_dim = config.linear_value_head_dim
        self.key_dim = self.num_k_heads * self.head_k_dim
        self.value_dim = self.num_v_heads * self.head_v_dim
        self.conv_kernel_size = config.linear_conv_kernel_dim
        self.act_fn = ACT2FN[config.hidden_act]
        conv_kwargs = {
            "kernel_size": self.conv_kernel_size,
            "padding": self.conv_kernel_size - 1,
            "bias": False,
        }
        self.q_conv1d = nn.Conv1d(
            self.key_dim, self.key_dim, groups=self.key_dim, **conv_kwargs
        )
        self.k_conv1d = nn.Conv1d(
            self.key_dim, self.key_dim, groups=self.key_dim, **conv_kwargs
        )
        self.v_conv1d = nn.Conv1d(
            self.value_dim, self.value_dim, groups=self.value_dim, **conv_kwargs
        )
        self.q_proj = nn.Linear(config.hidden_size, self.key_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.key_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.value_dim, bias=False)
        self.f_a_proj = nn.Linear(config.hidden_size, self.head_v_dim, bias=False)
        self.f_b_proj = nn.Linear(self.head_v_dim, self.key_dim, bias=False)
        self.b_proj = nn.Linear(config.hidden_size, self.num_k_heads, bias=False)
        self.g_a_proj = nn.Linear(config.hidden_size, self.head_v_dim, bias=False)
        self.g_b_proj = nn.Linear(self.head_v_dim, self.value_dim, bias=False)
        self.o_norm = AliceAIRMSNorm(
            self.head_v_dim, config.rms_norm_eps, zero_centered=False
        )
        self.o_proj = nn.Linear(self.value_dim, config.hidden_size, bias=False)
        self.a_log_bias = nn.Parameter(torch.empty(self.num_k_heads))
        self.dt_bias = nn.Parameter(torch.empty(self.key_dim))

    def _causal_conv(
        self,
        hidden_states: torch.Tensor,
        convolution: nn.Conv1d,
        cache: Cache | None,
        state_idx: int,
    ) -> torch.Tensor:
        channels_first = hidden_states.transpose(1, 2)
        sequence_length = channels_first.shape[-1]
        if cache is not None:
            channels_first = cache.update_conv_state(
                channels_first,
                self.layer_idx,
                state_idx,
                conv_kernel_size=self.conv_kernel_size,
            )
        output = self.act_fn(
            convolution(channels_first)[..., : channels_first.shape[-1]]
        )
        return output[..., -sequence_length:].transpose(1, 2)

    def _torch_kda(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        alpha: torch.Tensor,
        beta: torch.Tensor,
        initial_state: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        query = functional.normalize(query.float(), dim=-1) * self.head_k_dim**-0.5
        key = functional.normalize(key.float(), dim=-1)
        gate = -self.a_log_bias.float().exp().view(
            1, 1, self.num_k_heads, 1
        ) * functional.softplus(
            alpha.float()
            + self.dt_bias.float().view(1, 1, self.num_k_heads, self.head_k_dim)
        )
        beta = beta.float().sigmoid()
        if self.config.kda_allow_negative_eigenvalues:
            beta = beta * 2.0
        state = (
            initial_state.float()
            if initial_state is not None
            else query.new_zeros(
                query.shape[0],
                self.num_v_heads,
                self.head_k_dim,
                self.head_v_dim,
            )
        )
        repeat = self.num_v_heads // self.num_k_heads
        if repeat > 1:
            query = query.repeat_interleave(repeat, dim=2)
            key = key.repeat_interleave(repeat, dim=2)
            gate = gate.repeat_interleave(repeat, dim=2)
            beta = beta.repeat_interleave(repeat, dim=2)
        outputs = []
        for token_idx in range(query.shape[1]):
            query_i = query[:, token_idx]
            key_i = key[:, token_idx]
            value_i = value[:, token_idx].float()
            state = state * gate[:, token_idx].exp().unsqueeze(-1)
            prediction = torch.einsum("bhk,bhkv->bhv", key_i, state)
            delta = (value_i - prediction) * beta[:, token_idx].unsqueeze(-1)
            state = state + torch.einsum("bhk,bhv->bhkv", key_i, delta)
            outputs.append(torch.einsum("bhk,bhkv->bhv", query_i, state))
        return torch.stack(outputs, dim=1).to(value.dtype), state

    def _kda(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        alpha: torch.Tensor,
        beta: torch.Tensor,
        initial_state: torch.Tensor | None,
        use_cache: bool,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if query.is_cuda:
            try:
                from fla.ops.kda import chunk_kda, fused_recurrent_kda  # noqa: PLC0415

                kernel = (
                    fused_recurrent_kda
                    if query.shape[1] == 1 and initial_state is not None
                    else chunk_kda
                )
                kernel_beta = beta.float().sigmoid()
                if self.config.kda_allow_negative_eigenvalues:
                    kernel_beta = kernel_beta * 2.0
                return kernel(
                    q=query,
                    k=key,
                    v=value,
                    g=alpha,
                    beta=kernel_beta,
                    A_log=self.a_log_bias.float(),
                    dt_bias=self.dt_bias.float(),
                    initial_state=(
                        initial_state.float() if initial_state is not None else None
                    ),
                    output_final_state=use_cache,
                    use_qk_l2norm_in_kernel=True,
                    use_gate_in_kernel=True,
                )
            except ImportError as error:
                raise ImportError(
                    "CUDA execution requires flash-linear-attention with KDA support; "
                    "install flash-linear-attention>=0.5.0"
                ) from error
        output, state = self._torch_kda(query, key, value, alpha, beta, initial_state)
        return output, state if use_cache else None

    def forward(
        self,
        hidden_states: torch.Tensor,
        cache: Cache | None,
        attention_mask: torch.Tensor | None,
    ) -> torch.Tensor:
        if attention_mask is not None:
            hidden_states = (
                hidden_states * attention_mask[..., -hidden_states.shape[1] :, None]
            )
        query = self.q_proj(hidden_states)
        key = self.k_proj(hidden_states)
        value = self.v_proj(hidden_states)
        output_gate = self.g_b_proj(self.g_a_proj(hidden_states))
        alpha = self.f_b_proj(self.f_a_proj(hidden_states))
        beta = self.b_proj(hidden_states)
        has_previous_state = cache is not None and cache.has_previous_state(
            self.layer_idx
        )
        query = self._causal_conv(query, self.q_conv1d, cache, 0)
        key = self._causal_conv(key, self.k_conv1d, cache, 1)
        value = self._causal_conv(value, self.v_conv1d, cache, 2)
        query = query.view(*query.shape[:2], self.num_k_heads, self.head_k_dim)
        key = key.view(*key.shape[:2], self.num_k_heads, self.head_k_dim)
        value = value.view(*value.shape[:2], self.num_v_heads, self.head_v_dim)
        alpha = alpha.view(*alpha.shape[:2], self.num_k_heads, self.head_k_dim)
        initial_state = (
            cache.layers[self.layer_idx].recurrent_states[0]
            if has_previous_state
            else None
        )
        output, final_state = self._kda(
            query, key, value, alpha, beta, initial_state, cache is not None
        )
        if cache is not None:
            cache.update_recurrent_state(final_state, self.layer_idx)
        output = self.o_norm(output)
        output = output * torch.sigmoid(output_gate.view_as(output))
        return self.o_proj(output.flatten(-2))


class AliceAIMLP(nn.Module):
    def __init__(
        self, hidden_size: int, intermediate_size: int, hidden_act: str
    ) -> None:
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
        self.act_fn = ACT2FN[hidden_act]

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.down_proj(
            self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)
        )


class AliceAIExperts(nn.Module):
    def __init__(self, config: AliceAIConfig) -> None:
        super().__init__()
        self.num_experts = config.num_experts
        self.act_fn = ACT2FN[config.hidden_act]
        self.gate_up_proj = nn.Parameter(
            torch.empty(
                config.num_experts,
                2 * config.moe_intermediate_size,
                config.hidden_size,
            )
        )
        self.down_proj = nn.Parameter(
            torch.empty(
                config.num_experts,
                config.hidden_size,
                config.moe_intermediate_size,
            )
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        expert_indices: torch.LongTensor,
        expert_weights: torch.Tensor,
    ) -> torch.Tensor:
        output = torch.zeros_like(hidden_states)
        for expert_idx in range(self.num_experts):
            token_indices, slots = torch.where(expert_indices == expert_idx)
            if token_indices.numel() == 0:
                continue
            expert_input = hidden_states[token_indices]
            gate_up = functional.linear(expert_input, self.gate_up_proj[expert_idx])
            gate, up = gate_up.chunk(2, dim=-1)
            expert_output = functional.linear(
                self.act_fn(gate) * up, self.down_proj[expert_idx]
            )
            expert_output = expert_output * expert_weights[token_indices, slots, None]
            output.index_add_(0, token_indices, expert_output.to(output.dtype))
        return output


class AliceAISigmoidTopKRouter(nn.Module):
    def __init__(self, config: AliceAIConfig) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.empty(config.num_experts, config.hidden_size))
        self.top_k = config.num_experts_per_tok
        if config.router_bias_correction:
            self.register_buffer(
                "e_score_correction_bias", torch.zeros(config.num_experts)
            )
        else:
            self.e_score_correction_bias = None

    def forward(
        self, hidden_states: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
        router_logits = functional.linear(hidden_states.float(), self.weight.float())
        scores = router_logits.sigmoid()
        selection_scores = (
            scores
            if self.e_score_correction_bias is None
            else scores + self.e_score_correction_bias
        )
        indices = selection_scores.topk(self.top_k, dim=-1).indices
        weights = scores.gather(-1, indices)
        weights = weights / weights.sum(dim=-1, keepdim=True)
        return router_logits, weights.to(hidden_states.dtype), indices


class AliceAISparseMoEBlock(nn.Module):
    def __init__(self, config: AliceAIConfig) -> None:
        super().__init__()
        self.gate = AliceAISigmoidTopKRouter(config)
        self.experts = AliceAIExperts(config)
        self.shared_expert = AliceAIMLP(
            config.hidden_size,
            config.shared_expert_intermediate_size,
            config.hidden_act,
        )
        self.shared_expert_gate = nn.Linear(config.hidden_size, 1, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        original_shape = hidden_states.shape
        flattened = hidden_states.reshape(-1, original_shape[-1])
        router_logits, weights, indices = self.gate(flattened)
        routed = self.experts(flattened, indices, weights)
        shared = self.shared_expert(flattened) * torch.sigmoid(
            self.shared_expert_gate(flattened)
        )
        return (routed + shared).reshape(original_shape), router_logits


class AliceAIDecoderLayer(GradientCheckpointingLayer):
    def __init__(self, config: AliceAIConfig, layer_idx: int) -> None:
        super().__init__()
        self.layer_idx = layer_idx
        self.layer_type = config.layer_types[layer_idx]
        if self.layer_type == "linear_attention":
            self.linear_attn = AliceAIKDA(config, layer_idx)
        else:
            self.self_attn = AliceAIAttention(config, layer_idx)
        self.input_layernorm = AliceAIRMSNorm(
            config.hidden_size, config.rms_norm_eps, zero_centered=True
        )
        self.post_attention_layernorm = AliceAIRMSNorm(
            config.hidden_size, config.rms_norm_eps, zero_centered=True
        )
        self.mlp = AliceAISparseMoEBlock(config)
        if layer_idx != 0:
            self.attn_res_proj = nn.Linear(config.hidden_size, 1, bias=False)
            self.attn_res_norm_weight = nn.Parameter(torch.ones(config.hidden_size))
        self.mlp_res_proj = nn.Linear(config.hidden_size, 1, bias=False)
        self.mlp_res_norm_weight = nn.Parameter(torch.ones(config.hidden_size))
        self.eps = config.rms_norm_eps

    def _mix(
        self,
        completed: list[torch.Tensor],
        partial: torch.Tensor | None,
        projection: nn.Linear,
        norm_weight: torch.Tensor,
    ) -> torch.Tensor:
        sources = completed if partial is None else [*completed, partial]
        return _depth_softmax_mix(sources, projection.weight, norm_weight, self.eps)

    def forward(
        self,
        embeddings: torch.Tensor,
        completed_blocks: tuple[torch.Tensor, ...],
        partial: torch.Tensor | None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None,
        past_key_values: Cache | None,
        cache_position: torch.LongTensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if self.layer_idx == 0:
            mixed = embeddings
        else:
            mixed = self._mix(
                completed_blocks, partial, self.attn_res_proj, self.attn_res_norm_weight
            )
        mixed = self.input_layernorm(mixed)
        if self.layer_type == "linear_attention":
            attention_output = self.linear_attn(
                mixed, past_key_values, attention_mask
            )
        else:
            attention_output, _ = self.self_attn(
                mixed,
                position_embeddings,
                attention_mask,
                past_key_values,
                cache_position,
            )
        partial = attention_output if partial is None else partial + attention_output
        mixed = self._mix(
            completed_blocks, partial, self.mlp_res_proj, self.mlp_res_norm_weight
        )
        moe_output, router_logits = self.mlp(self.post_attention_layernorm(mixed))
        return partial + moe_output, router_logits


class AliceAIPreTrainedModel(PreTrainedModel):
    config_class = AliceAIConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _is_stateful = True
    _no_split_modules: ClassVar[list[str]] = ["AliceAIDecoderLayer"]
    _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [r"^mtp\."]
    _supports_sdpa = True
    _supports_flash_attn = True
    _can_record_outputs: ClassVar[dict[str, type[nn.Module]]] = {
        "hidden_states": AliceAIDecoderLayer,
        "attentions": AliceAIAttention,
    }

    @staticmethod
    def _needs_initialization(parameter: torch.Tensor) -> bool:
        return not getattr(parameter, "_is_hf_initialized", False)

    @torch.no_grad()
    def _init_weights(self, module: nn.Module) -> None:  # noqa: PLR0912
        if isinstance(module, (nn.Linear, nn.Conv1d)):
            module.weight.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight[module.padding_idx].zero_()
        elif isinstance(module, AliceAIExperts):
            if self._needs_initialization(module.gate_up_proj):
                module.gate_up_proj.normal_(mean=0.0, std=self.config.initializer_range)
            if self._needs_initialization(module.down_proj):
                module.down_proj.normal_(mean=0.0, std=self.config.initializer_range)
        elif isinstance(module, AliceAISigmoidTopKRouter):
            if self._needs_initialization(module.weight):
                module.weight.normal_(mean=0.0, std=self.config.initializer_range)
        elif isinstance(module, AliceAIKDA):
            if self._needs_initialization(module.a_log_bias):
                module.a_log_bias.uniform_(0, 16).log_()
            if self._needs_initialization(module.dt_bias):
                dt = torch.exp(
                    torch.rand_like(module.dt_bias) * (math.log(0.1) - math.log(0.001))
                    + math.log(0.001)
                ).clamp_min(1e-4)
                module.dt_bias.copy_(dt + torch.log(-torch.expm1(-dt)))
        elif isinstance(module, AliceAIFinalBlockAttnRes):
            if self._needs_initialization(module.res_proj.weight):
                module.res_proj.weight.zero_()
        elif isinstance(module, AliceAIDecoderLayer):
            if hasattr(module, "attn_res_proj"):
                if self._needs_initialization(module.attn_res_proj.weight):
                    module.attn_res_proj.weight.zero_()
            if self._needs_initialization(module.mlp_res_proj.weight):
                module.mlp_res_proj.weight.zero_()


class AliceAIModel(AliceAIPreTrainedModel):
    def __init__(self, config: AliceAIConfig) -> None:
        super().__init__(config)
        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size, config.pad_token_id
        )
        self.layers = nn.ModuleList(
            [
                AliceAIDecoderLayer(config, index)
                for index in range(config.num_hidden_layers)
            ]
        )
        self.attnres_final = AliceAIFinalBlockAttnRes(config)
        self.norm = AliceAIRMSNorm(
            config.hidden_size, config.rms_norm_eps, zero_centered=True
        )
        self.rotary_emb = AliceAIRotaryEmbedding(config)
        self.gradient_checkpointing = False
        self.post_init()

    @merge_with_config_defaults
    @capture_outputs
    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        use_cache: bool | None = None,
        output_router_logits: bool | None = None,
        **_: Any,
    ) -> MoeModelOutputWithPast:
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Specify exactly one of input_ids or inputs_embeds")
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        use_cache = self.config.use_cache if use_cache is None else use_cache
        output_router_logits = (
            self.config.output_router_logits
            if output_router_logits is None
            else output_router_logits
        )
        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)
        past_length = (
            past_key_values.get_seq_length() if past_key_values is not None else 0
        )
        cache_position = torch.arange(
            past_length,
            past_length + inputs_embeds.shape[1],
            device=inputs_embeds.device,
        )
        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)
        position_embeddings = self.rotary_emb(position_ids, inputs_embeds.dtype)
        if isinstance(attention_mask, dict):
            linear_mask = attention_mask["linear_attention"]
            full_attention_mask = attention_mask["full_attention"]
        else:
            linear_mask = attention_mask
            full_attention_mask = attention_mask
            if cache_position[0] > 0 or (
                attention_mask is not None and torch.all(attention_mask == 1)
            ):
                linear_mask = None

        completed_blocks = [inputs_embeds]
        partial = None
        router_logits = []
        for layer_idx, layer in enumerate(self.layers):
            if layer_idx > 0 and layer_idx % self.config.block_attn_res_block_size == 0:
                completed_blocks.append(partial)
                partial = None
            layer_mask = (
                linear_mask
                if layer.layer_type == "linear_attention"
                else full_attention_mask
            )
            partial, layer_router_logits = layer(
                inputs_embeds,
                completed_blocks=tuple(completed_blocks),
                partial=partial,
                position_embeddings=position_embeddings,
                attention_mask=layer_mask,
                past_key_values=past_key_values,
                cache_position=cache_position,
            )
            if output_router_logits:
                router_logits.append(layer_router_logits)
        hidden_states = self.norm(self.attnres_final(completed_blocks, partial))
        return MoeModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values,
            router_logits=tuple(router_logits) if output_router_logits else None,
        )


class AliceAIForCausalLM(AliceAIPreTrainedModel, GenerationMixin):
    _tied_weights_keys: ClassVar[dict[str, str]] = {
        "lm_head.weight": "model.embed_tokens.weight"
    }

    def __init__(self, config: AliceAIConfig) -> None:
        super().__init__(config)
        self.model = AliceAIModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.vocab_size = config.vocab_size
        self.post_init()

    def get_input_embeddings(self) -> nn.Module:
        return self.model.embed_tokens

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.model.embed_tokens = value

    def get_output_embeddings(self) -> nn.Module:
        return self.lm_head

    def set_output_embeddings(self, value: nn.Module) -> None:
        self.lm_head = value

    @can_return_tuple
    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        labels: torch.LongTensor | None = None,
        use_cache: bool | None = None,
        output_router_logits: bool | None = None,
        logits_to_keep: int | torch.Tensor = 0,
        **kwargs: Any,
    ) -> MoeCausalLMOutputWithPast:
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_router_logits=output_router_logits,
            **kwargs,
        )
        indices = (
            slice(-logits_to_keep, None)
            if isinstance(logits_to_keep, int)
            else logits_to_keep
        )
        logits = self.lm_head(outputs.last_hidden_state[:, indices, :])
        loss = (
            self.loss_function(logits, labels, self.vocab_size, **kwargs)
            if labels is not None
            else None
        )
        return MoeCausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            router_logits=outputs.router_logits,
        )


AliceAIConfig.register_for_auto_class()
AliceAIModel.register_for_auto_class("AutoModel")
AliceAIForCausalLM.register_for_auto_class("AutoModelForCausalLM")