modeling_limite.py

"""Native Transformers implementation of the Limite causal language model.

BF16 SDPA is the portable attention path. The implementation also supports
Transformers' cache protocol so the same model can be used by ``generate``
without a separate decoding graph.
"""

from __future__ import annotations

from contextlib import nullcontext
from typing import Any

import torch
import torch.nn.functional as F
from torch import Tensor, nn
from torch.nn.attention import SDPBackend, sdpa_kernel
from transformers.cache_utils import Cache, DynamicCache, StaticCache
from transformers.generation import GenerationMixin
from transformers.masking_utils import (
    create_causal_mask,
    create_sliding_window_causal_mask,
)
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel

from .configuration_limite import LimiteConfig
from .registration import register_weight_converters

register_weight_converters()


def _uses_external_flash_attention(attn_implementation: str | None) -> bool:
    """Identify native and Hub-provided FlashAttention implementations."""
    normalized = str(attn_implementation or "").lower().replace("-", "_")
    return "flash_attention" in normalized or "flash_attn" in normalized


def _reject_external_flash_static_cache(attn_implementation: str | None) -> None:
    if _uses_external_flash_attention(attn_implementation):
        raise ValueError(
            "Limite does not support FlashAttention with StaticCache because "
            "that combination can produce incorrect logits. Use "
            "attn_implementation='sdpa' with StaticCache, or use "
            "DynamicCache with FlashAttention."
        )


def _validate_cache_attention_pair(
    *,
    attn_implementation: str | None,
    past_key_values: Cache | None,
) -> None:
    if isinstance(past_key_values, StaticCache):
        _reject_external_flash_static_cache(attn_implementation)


def rms_norm(x: Tensor) -> Tensor:
    """Gain-free RMS norm with PyTorch's dtype-dependent default epsilon."""
    return F.rms_norm(x, (x.size(-1),))


class LimiteRMSNorm(nn.Module):
    def forward(self, hidden_states: Tensor) -> Tensor:
        return rms_norm(hidden_states)


class LimiteRotaryEmbedding(nn.Module):
    """Checkpoint-exact rotary factors with the static frequencies cached."""

    def __init__(self, config: LimiteConfig) -> None:
        super().__init__()
        self.rope_base_local = float(config.rope_base_local)
        self.rope_n_pairs = int(config.rope_n_pairs)
        self.head_dim = int(config.head_dim)
        self.register_buffer(
            "frequency",
            self._build_frequency(),
            persistent=False,
        )

    def _build_frequency(self, device: torch.device | None = None) -> Tensor:
        frequency = (1.0 / self.rope_base_local) ** torch.linspace(
            0,
            1,
            steps=self.rope_n_pairs,
            dtype=torch.float32,
            device="cpu",
        )
        frequency = frequency.repeat_interleave(2)
        frequency = torch.cat(
            [frequency, frequency.new_zeros(self.head_dim - frequency.numel())]
        )
        return frequency if device is None else frequency.to(device=device)

    def _apply(self, fn: Any, recurse: bool = True) -> "LimiteRotaryEmbedding":
        super()._apply(fn, recurse=recurse)
        # Transformers applies ``dtype=...`` to buffers too. RoPE frequencies
        # are part of Limite's FP32 numerical contract, so reconstruct the
        # derived buffer from the CPU-FP32 formula on the destination device.
        self.frequency = self._build_frequency(device=self.frequency.device)
        return self

    def forward(self, position_ids: Tensor) -> tuple[Tensor, Tensor]:
        theta = position_ids.to(torch.float32).unsqueeze(-1) * self.frequency
        cosine = theta.cos().to(torch.bfloat16).unsqueeze(-2)
        sine = theta.sin().to(torch.bfloat16)
        sine[..., 1::2] *= -1
        return cosine, sine.unsqueeze(-2)


def apply_rotary(x: Tensor, cosine: Tensor, sine: Tensor) -> Tensor:
    paired = x.view(*x.shape[:-1], x.shape[-1] // 2, 2).flip(-1).view(x.shape)
    return cosine * x + sine * paired


def repeat_kv(hidden_states: Tensor, num_groups: int) -> Tensor:
    """Expand key/value heads for the eager attention oracle."""
    if num_groups == 1:
        return hidden_states
    batch_size, num_kv_heads, sequence_length, head_dim = hidden_states.shape
    hidden_states = hidden_states[:, :, None, :, :].expand(
        batch_size,
        num_kv_heads,
        num_groups,
        sequence_length,
        head_dim,
    )
    return hidden_states.reshape(
        batch_size,
        num_kv_heads * num_groups,
        sequence_length,
        head_dim,
    )


def eager_attention_forward(
    module: nn.Module,
    query: Tensor,
    key: Tensor,
    value: Tensor,
    attention_mask: Tensor | None,
    scaling: float,
    dropout: float = 0.0,
    **kwargs: Any,
) -> tuple[Tensor, Tensor]:
    """Reference attention used for backend parity checks."""
    del kwargs
    key = repeat_kv(key, module.num_key_value_groups)
    value = repeat_kv(value, module.num_key_value_groups)
    weights = torch.matmul(query, key.transpose(2, 3)) * scaling
    if attention_mask is not None:
        weights = weights + attention_mask
    probabilities = F.softmax(weights, dim=-1, dtype=torch.float32).to(query.dtype)
    probabilities = F.dropout(
        probabilities,
        p=dropout,
        training=module.training,
    )
    output = torch.matmul(probabilities, value).transpose(1, 2).contiguous()
    return output, probabilities


def _fp32_parameter(*shape: int, initial: float = 0.0) -> nn.Parameter:
    return nn.Parameter(torch.full(shape, initial, dtype=torch.float32))


class LimiteAttention(nn.Module):
    def __init__(self, config: LimiteConfig, layer_idx: int) -> None:
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.head_dim = int(config.head_dim)
        self.num_heads = int(config.num_attention_heads)
        self.num_kv_heads = int(config.num_key_value_heads)
        self.num_kv_groups = self.num_heads // self.num_kv_heads
        self.num_key_value_groups = self.num_kv_groups
        self.scaling = float(config.attention_softmax_scale)
        self.attention_dropout = float(config.attention_dropout)
        self.is_causal = True
        self.is_global = config.is_global_layer(layer_idx)
        self.window_span = None if self.is_global else int(config.sliding_window)
        self.applies_rope = not (self.is_global and bool(config.global_nope))
        self.has_ve = layer_idx in set(config.ve_layers)
        self.has_xsa = bool(config.xsa) and layer_idx in set(config.xsa_layers)
        self.attn_gate_channels = int(config.attn_gate_channels)

        hidden_size = int(config.hidden_size)
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.qkv_proj = nn.Linear(
            hidden_size,
            self.q_size + 2 * self.kv_size,
            bias=False,
        )
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, hidden_size, bias=False)

        self.qkv_scale = _fp32_parameter(initial=1.0)
        self.o_scale = _fp32_parameter(initial=1.0)
        self.register_buffer("_inference_qkv_weight", None, persistent=False)
        self.register_buffer("_inference_o_weight", None, persistent=False)
        self.register_buffer("_inference_xsa_alpha", None, persistent=False)
        self.register_buffer("_inference_ve_gate", None, persistent=False)
        self.register_buffer("_inference_attn_gate", None, persistent=False)
        if self.has_xsa:
            self.xsa_alpha = _fp32_parameter(self.num_heads)
        if self.has_ve:
            self.ve_gate = _fp32_parameter(
                int(config.ve_stored_heads), int(config.ve_gate_channels)
            )
        if self.attn_gate_channels:
            self.attn_gate = _fp32_parameter(self.num_heads, self.attn_gate_channels)

    @staticmethod
    def _scaled(weight: Tensor, scale: Tensor, dtype: torch.dtype) -> Tensor:
        return (scale.to(torch.float32).view(()) * weight).to(dtype)

    def train(self, mode: bool = True) -> "LimiteAttention":
        super().train(mode)
        if mode:
            self._inference_qkv_weight = None
            self._inference_o_weight = None
            self._inference_xsa_alpha = None
            self._inference_ve_gate = None
            self._inference_attn_gate = None
        else:
            dtype = self.qkv_proj.weight.dtype
            self._inference_qkv_weight = self._scaled(
                self.qkv_proj.weight,
                self.qkv_scale,
                dtype,
            ).detach()
            self._inference_o_weight = self._scaled(
                self.o_proj.weight, self.o_scale, self.o_proj.weight.dtype
            ).detach()
            if self.has_xsa:
                self._inference_xsa_alpha = torch.tanh(self.xsa_alpha.float()).detach()
            if self.has_ve:
                self._inference_ve_gate = (
                    self.ve_gate[: self.num_kv_heads].to(dtype).detach()
                )
            if self.attn_gate_channels:
                self._inference_attn_gate = self.attn_gate.to(dtype).detach()
        return self

    def _project_qkv(self, hidden_states: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        sizes = (self.q_size, self.kv_size, self.kv_size)
        if not self.training and self._inference_qkv_weight is not None:
            return F.linear(hidden_states, self._inference_qkv_weight).split(
                sizes, dim=-1
            )
        dtype = hidden_states.dtype
        return F.linear(
            hidden_states,
            self._scaled(self.qkv_proj.weight, self.qkv_scale, dtype),
        ).split(
            sizes,
            dim=-1,
        )

    def _apply_value_embeddings(
        self, hidden_states: Tensor, value_embeds: Tensor, value_states: Tensor
    ) -> Tensor:
        gate_weight = (
            self._inference_ve_gate
            if not self.training and self._inference_ve_gate is not None
            else self.ve_gate[: self.num_kv_heads].to(hidden_states.dtype)
        )
        gate = float(self.config.ve_gate_scale) * torch.sigmoid(
            F.linear(hidden_states[..., : gate_weight.size(-1)], gate_weight)
        )
        return value_states + gate.unsqueeze(-1) * value_embeds.to(value_states.dtype)

    def forward(
        self,
        hidden_states: Tensor,
        value_embeds: Tensor | None,
        cosine: Tensor,
        sine: Tensor,
        attention_mask: Tensor | None,
        past_key_values: Cache | None,
        use_cache: bool,
        output_attentions: bool,
    ) -> tuple[Tensor, Tensor | None]:
        batch_size, query_length, _ = hidden_states.shape
        query_states, key_states, value_states = self._project_qkv(hidden_states)
        query_states = query_states.view(
            batch_size, query_length, self.num_heads, self.head_dim
        )
        key_states = key_states.view(
            batch_size, query_length, self.num_kv_heads, self.head_dim
        )
        value_states = value_states.view(
            batch_size, query_length, self.num_kv_heads, self.head_dim
        )

        if self.has_ve and value_embeds is not None:
            value_states = self._apply_value_embeddings(
                hidden_states, value_embeds, value_states
            )
        current_values = value_states
        query_states, key_states = rms_norm(query_states), rms_norm(key_states)
        if self.applies_rope:
            query_states = apply_rotary(query_states, cosine, sine)
            key_states = apply_rotary(key_states, cosine, sine)

        if use_cache:
            if past_key_values is None:
                raise ValueError("use_cache=True requires a cache instance")
            cached_keys, cached_values = past_key_values.update(
                key_states.transpose(1, 2),
                value_states.transpose(1, 2),
                self.layer_idx,
            )
            key_states = cached_keys.transpose(1, 2)
            value_states = cached_values.transpose(1, 2)
        if (
            self.config._attn_implementation == "sdpa"
            and attention_mask is None
            and query_length == 1
        ):
            flash_decode = (
                query_states.is_cuda
                and query_states.dtype in (torch.float16, torch.bfloat16)
                and torch.cuda.get_device_capability(query_states.device)[0] >= 8
            )
            backend_context = (
                sdpa_kernel(SDPBackend.FLASH_ATTENTION)
                if flash_decode
                else nullcontext()
            )
            with backend_context:
                attention_output = (
                    F.scaled_dot_product_attention(
                        query_states.transpose(1, 2),
                        key_states.transpose(1, 2),
                        value_states.transpose(1, 2),
                        dropout_p=(
                            0.0 if not self.training else self.attention_dropout
                        ),
                        scale=self.scaling,
                        is_causal=False,
                        enable_gqa=True,
                    )
                    .transpose(1, 2)
                    .contiguous()
                )
            probabilities = None
        else:
            attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
                self.config._attn_implementation,
                eager_attention_forward,
            )
            attention_output, probabilities = attention_interface(
                self,
                query_states.transpose(1, 2),
                key_states.transpose(1, 2),
                value_states.transpose(1, 2),
                attention_mask,
                dropout=0.0 if not self.training else self.attention_dropout,
                scaling=self.scaling,
                sliding_window=self.window_span,
                output_attentions=output_attentions,
            )
        if self.has_xsa:
            alpha_values = (
                self._inference_xsa_alpha
                if not self.training and self._inference_xsa_alpha is not None
                else torch.tanh(self.xsa_alpha.float())
            )
            if not self.training:
                grouped_output = attention_output.view(
                    batch_size,
                    query_length,
                    self.num_kv_heads,
                    self.num_kv_groups,
                    self.head_dim,
                )
                value_direction = F.normalize(
                    current_values.float(),
                    dim=-1,
                    eps=float(self.config.xsa_normalize_eps),
                ).unsqueeze(3)
                projection = (grouped_output.float() * value_direction).sum(
                    -1, keepdim=True
                )
                alpha = alpha_values.view(
                    1, 1, self.num_kv_heads, self.num_kv_groups, 1
                )
                attention_output = (
                    grouped_output
                    - (alpha * projection * value_direction).to(grouped_output.dtype)
                ).reshape(batch_size, query_length, self.num_heads, self.head_dim)
            else:
                value_direction = F.normalize(
                    current_values.repeat_interleave(self.num_kv_groups, dim=2).float(),
                    dim=-1,
                    eps=float(self.config.xsa_normalize_eps),
                )
                projection = (attention_output.float() * value_direction).sum(
                    -1, keepdim=True
                )
                alpha = alpha_values.view(1, 1, self.num_heads, 1)
                attention_output = attention_output - (
                    alpha * projection * value_direction
                ).to(attention_output.dtype)
        if self.attn_gate_channels:
            gate_weight = (
                self._inference_attn_gate
                if not self.training and self._inference_attn_gate is not None
                else self.attn_gate.to(hidden_states.dtype)
            )
            gate = float(self.config.attn_gate_scale) * torch.sigmoid(
                F.linear(
                    hidden_states[..., : self.attn_gate_channels],
                    gate_weight,
                )
            )
            attention_output = attention_output * gate.to(
                attention_output.dtype
            ).unsqueeze(-1)

        attention_output = attention_output.reshape(batch_size, query_length, -1)
        if not self.training and self._inference_o_weight is not None:
            attention_output = F.linear(attention_output, self._inference_o_weight)
        else:
            attention_output = F.linear(
                attention_output,
                self._scaled(
                    self.o_proj.weight,
                    self.o_scale,
                    attention_output.dtype,
                ),
            )
        return attention_output, probabilities if output_attentions else None


class LimiteMLP(nn.Module):
    def __init__(self, config: LimiteConfig) -> None:
        super().__init__()
        hidden_size = int(config.hidden_size)
        intermediate_size = int(config.intermediate_size)
        self.intermediate_size = intermediate_size
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
        self.gate_up_proj = nn.Linear(
            hidden_size,
            2 * intermediate_size,
            bias=False,
        )

    def forward(self, hidden_states: Tensor) -> Tensor:
        gate, up = F.linear(
            hidden_states,
            self.gate_up_proj.weight,
        ).split(self.intermediate_size, dim=-1)
        activated = F.silu(gate) * up
        return self.down_proj(activated)


class LimiteDecoderLayer(nn.Module):
    def __init__(self, config: LimiteConfig, layer_idx: int) -> None:
        super().__init__()
        self.self_attn = LimiteAttention(config, layer_idx)
        self.mlp = LimiteMLP(config)
        self.resid_lambda_attn = _fp32_parameter(initial=1.0)
        self.post_lambda_attn = _fp32_parameter(initial=1.0)
        self.resid_lambda_mlp = _fp32_parameter(initial=1.0)
        self.post_lambda_mlp = _fp32_parameter(initial=1.0)
        self.register_buffer("_inference_residual_scales", None, persistent=False)

    def train(self, mode: bool = True) -> "LimiteDecoderLayer":
        super().train(mode)
        if mode:
            self._inference_residual_scales = None
        else:
            dtype = self.self_attn.qkv_proj.weight.dtype
            self._inference_residual_scales = (
                torch.stack(
                    (
                        self.resid_lambda_attn,
                        self.post_lambda_attn,
                        self.resid_lambda_mlp,
                        self.post_lambda_mlp,
                    )
                )
                .to(dtype)
                .detach()
            )
        return self

    def forward(
        self,
        hidden_states: Tensor,
        attention_input: Tensor,
        value_embeds: Tensor | None,
        cosine: Tensor,
        sine: Tensor,
        attention_mask: Tensor | None,
        past_key_values: Cache | None,
        use_cache: bool,
        output_attentions: bool,
        attention_residual: Tensor | None = None,
    ) -> tuple[Tensor, Tensor | None]:
        attention_output, probabilities = self.self_attn(
            attention_input,
            value_embeds,
            cosine,
            sine,
            attention_mask,
            past_key_values,
            use_cache,
            output_attentions,
        )
        residual_base = (
            hidden_states if attention_residual is None else attention_residual
        )
        use_constants = (
            not self.training and self._inference_residual_scales is not None
        )
        if use_constants:
            residual_scales = self._inference_residual_scales
            mixed = (
                residual_scales[0] * residual_base
                + residual_scales[1] * attention_output
            )
        else:
            mixed = (
                self.resid_lambda_attn.to(hidden_states.dtype) * residual_base
                + self.post_lambda_attn.to(attention_output.dtype) * attention_output
            )
        mlp_output = self.mlp(rms_norm(mixed))
        if use_constants:
            output = residual_scales[2] * mixed + residual_scales[3] * mlp_output
        else:
            output = (
                self.resid_lambda_mlp.to(mixed.dtype) * mixed
                + self.post_lambda_mlp.to(mlp_output.dtype) * mlp_output
            )
        return output, probabilities


class LimiteMudd(nn.Module):
    def __init__(self, config: LimiteConfig) -> None:
        super().__init__()
        self.dense1 = _fp32_parameter(int(config.mudd_inter), int(config.hidden_size))
        self.dense2 = _fp32_parameter(
            int(config.num_hidden_layers),
            int(config.mudd_taps),
            int(config.mudd_inter),
        )
        self.bias = _fp32_parameter(
            int(config.num_hidden_layers), int(config.mudd_taps)
        )
        self.register_buffer("_inference_dense1", None, persistent=False)
        self.register_buffer("_inference_dense2", None, persistent=False)
        self.register_buffer("_inference_bias", None, persistent=False)
        self.register_buffer("_inference_dense2_mlp", None, persistent=False)
        self.register_buffer("_inference_bias_mlp", None, persistent=False)
        self.uses_r_way = bool(config.mudd_mlp)
        if self.uses_r_way:
            self.dense2_mlp = _fp32_parameter(
                int(config.num_hidden_layers),
                int(config.mudd_taps),
                int(config.mudd_inter),
            )
            self.bias_mlp = _fp32_parameter(
                int(config.num_hidden_layers), int(config.mudd_taps)
            )

    def train(self, mode: bool = True) -> "LimiteMudd":
        super().train(mode)
        if mode:
            self._inference_dense1 = None
            self._inference_dense2 = None
            self._inference_bias = None
            self._inference_dense2_mlp = None
            self._inference_bias_mlp = None
        else:
            self._inference_dense1 = self.dense1.to(torch.bfloat16).detach()
            self._inference_dense2 = self.dense2.to(torch.bfloat16).detach()
            self._inference_bias = self.bias.to(torch.bfloat16).detach()
            if self.uses_r_way:
                self._inference_dense2_mlp = self.dense2_mlp.to(torch.bfloat16).detach()
                self._inference_bias_mlp = self.bias_mlp.to(torch.bfloat16).detach()
        return self

    def _inner(self, current: Tensor) -> Tensor:
        use_constants = not self.training and self._inference_dense1 is not None
        dense1 = (
            self._inference_dense1 if use_constants else self.dense1.to(current.dtype)
        )
        return F.gelu(F.linear(rms_norm(current), dense1))

    def _mix(
        self,
        taps: list[Tensor],
        inner: Tensor,
        layer_idx: int,
        *,
        r_way: bool,
    ) -> Tensor:
        count = len(taps)
        use_constants = not self.training and self._inference_dense1 is not None
        if r_way:
            if not self.uses_r_way:
                raise ValueError("R-way mixing requested without R-way weights")
            if use_constants:
                dense2, bias = (
                    self._inference_dense2_mlp,
                    self._inference_bias_mlp,
                )
            else:
                dense2, bias = self.dense2_mlp, self.bias_mlp
        elif use_constants:
            dense2, bias = self._inference_dense2, self._inference_bias
        else:
            dense2, bias = self.dense2, self.bias
        weights = torch.einsum(
            "btk,mk->btm",
            inner,
            dense2[layer_idx, :count].to(inner.dtype),
        )
        weights = weights + bias[layer_idx, :count].to(weights.dtype)
        output = weights[..., 0:1].type_as(taps[0]) * taps[0]
        for index in range(1, count):
            output = (
                output
                + weights[..., index : index + 1].type_as(taps[index]) * taps[index]
            )
        return output

    def forward(
        self,
        taps: list[Tensor],
        current: Tensor,
        layer_idx: int,
        *,
        r_way: bool = False,
    ) -> Tensor:
        return self._mix(
            taps,
            self._inner(current),
            layer_idx,
            r_way=r_way,
        )

    def forward_pair(
        self,
        taps: list[Tensor],
        current: Tensor,
        layer_idx: int,
    ) -> tuple[Tensor, Tensor]:
        if not self.uses_r_way:
            raise ValueError("paired MUDD mixing requires R-way weights")
        inner = self._inner(current)
        return (
            self._mix(taps, inner, layer_idx, r_way=False),
            self._mix(taps, inner, layer_idx, r_way=True),
        )


class LimitePreTrainedModel(PreTrainedModel):
    config_class = LimiteConfig
    base_model_prefix = "model"
    _no_split_modules = ["LimiteDecoderLayer"]
    _skip_keys_device_placement = ["past_key_values"]
    _supports_flash_attn = True
    _supports_sdpa = True
    _supports_attention_backend = True
    _can_compile_fullgraph = True

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | None,
        *model_args: Any,
        **kwargs: Any,
    ) -> LimitePreTrainedModel:
        requested_parallelism = [
            name
            for name in ("tp_plan", "tp_size", "distributed_config")
            if kwargs.get(name) is not None
        ]
        device_map = kwargs.get("device_map")
        if isinstance(device_map, str) and device_map in {
            "auto",
            "balanced",
            "balanced_low_0",
            "sequential",
        }:
            requested_parallelism.append(f"device_map={device_map!r}")
        elif isinstance(device_map, dict):
            placements = {str(device) for device in device_map.values()}
            if len(placements) > 1:
                requested_parallelism.append("multi-device device_map")

        if requested_parallelism:
            requested = ", ".join(requested_parallelism)
            raise NotImplementedError(
                "Limite supports one complete model replica per process; "
                "tensor parallelism, pipeline parallelism, and multi-device "
                f"model sharding are not supported (requested: {requested}). "
                "Use process-level data parallelism with one explicit device "
                "per replica."
            )
        return super().from_pretrained(
            pretrained_model_name_or_path,
            *model_args,
            **kwargs,
        )

    @torch.no_grad()
    def _init_weights(self, module: nn.Module) -> None:
        super()._init_weights(module)
        if isinstance(module, LimiteRotaryEmbedding):
            # Transformers materializes non-persistent buffers with
            # ``empty_like`` during low-memory/device-map loading. Restore this
            # derived FP32 buffer before the loaded model is returned.
            module.frequency.copy_(
                module._build_frequency(device=module.frequency.device)
            )


class LimiteModel(LimitePreTrainedModel):
    def __init__(self, config: LimiteConfig) -> None:
        super().__init__(config)
        self.vocab_size = int(config.vocab_size)
        self.embed_tokens = nn.Embedding(
            int(config.vocab_size), int(config.hidden_size)
        )
        self.value_embeds = nn.Embedding(
            int(config.vocab_size), int(config.ve_stored_heads) * int(config.ve_dim)
        )
        self.layers = nn.ModuleList(
            [
                LimiteDecoderLayer(config, layer_idx)
                for layer_idx in range(int(config.num_hidden_layers))
            ]
        )
        self.norm = LimiteRMSNorm()
        self.rotary_emb = LimiteRotaryEmbedding(config)
        self.mudd = LimiteMudd(config) if config.mudd else None
        self.retained_taps = sorted(
            {
                tap
                for layer, taps in config.mudd_tap_idx.items()
                for tap in taps
                if tap != int(layer)
            }
        )
        self.post_init()

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

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

    def _value_embeddings(self, input_ids: Tensor) -> Tensor | None:
        if not self.config.ve_layers:
            return None
        config = self.config
        embeddings = self.value_embeds(input_ids).view(
            *input_ids.shape, int(config.ve_stored_heads), int(config.ve_dim)
        )
        if config.ve_dim < config.head_dim:
            embeddings = F.pad(
                embeddings,
                (0, int(config.head_dim) - int(config.ve_dim)),
            )
        return embeddings[..., : int(config.num_key_value_heads), :].contiguous()

    def forward(
        self,
        input_ids: Tensor | None = None,
        attention_mask: Tensor | None = None,
        position_ids: Tensor | None = None,
        past_key_values: Cache | None = None,
        use_cache: bool | None = None,
        inputs_embeds: Tensor | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **kwargs: Any,
    ) -> BaseModelOutputWithPast | tuple[Tensor, ...]:
        del kwargs
        if input_ids is None or inputs_embeds is not None:
            raise ValueError(
                "Limite requires input_ids because value embeddings are a "
                "second token lookup"
            )
        use_cache = self.config.use_cache if use_cache is None else use_cache
        output_attentions = bool(output_attentions)
        output_hidden_states = bool(output_hidden_states)
        return_dict = (
            self.config.use_return_dict if return_dict is None else return_dict
        )
        if output_attentions:
            raise ValueError(
                "Limite does not materialize attention weights; "
                "output_attentions=True is unsupported"
            )

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)
        if use_cache:
            _validate_cache_attention_pair(
                attn_implementation=self.config._attn_implementation,
                past_key_values=past_key_values,
            )
        if (
            use_cache
            and isinstance(past_key_values, DynamicCache)
            and not hasattr(past_key_values, "_limite_unpadded")
        ):
            if attention_mask is None:
                past_key_values._limite_unpadded = True
            elif isinstance(attention_mask, Tensor):
                past_key_values._limite_unpadded = not bool(
                    (attention_mask == 0).any().item()
                )
            else:
                past_key_values._limite_unpadded = False
        if position_ids is None:
            if attention_mask is not None:
                position_ids = attention_mask.long().cumsum(-1) - 1
                position_ids.masked_fill_(attention_mask == 0, 0)
                position_ids = position_ids[:, -input_ids.shape[1] :]
            else:
                past_length = (
                    past_key_values.get_seq_length()
                    if past_key_values is not None
                    else 0
                )
                position_ids = (
                    torch.arange(
                        past_length,
                        past_length + input_ids.shape[1],
                        device=input_ids.device,
                    )
                    .unsqueeze(0)
                    .expand(input_ids.shape[0], -1)
                )

        cosine, sine = self.rotary_emb(position_ids)
        value_embeds = self._value_embeddings(input_ids)
        hidden_states = rms_norm(self.embed_tokens(input_ids))
        maskless_sdpa_decode = (
            self.config._attn_implementation == "sdpa"
            and isinstance(past_key_values, DynamicCache)
            and input_ids.shape[1] == 1
            and bool(getattr(past_key_values, "_limite_unpadded", False))
        )
        if isinstance(attention_mask, dict):
            causal_mask_mapping = attention_mask
        elif maskless_sdpa_decode:
            causal_mask_mapping = {
                "full_attention": None,
                "sliding_attention": None,
            }
        else:
            mask_kwargs = {
                "config": self.config,
                "inputs_embeds": hidden_states,
                "attention_mask": attention_mask,
                "past_key_values": past_key_values,
                "position_ids": position_ids,
            }
            causal_mask_mapping = {
                "full_attention": create_causal_mask(**mask_kwargs),
                "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs),
            }
        history: dict[int, Tensor] = (
            {0: hidden_states} if 0 in self.retained_taps else {}
        )
        all_hidden_states: tuple[Tensor, ...] = ()
        all_attentions: tuple[Tensor, ...] = ()

        for layer_idx, decoder_layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (hidden_states,)
            taps = self.config.tap_indices(layer_idx)
            if taps is not None:
                tap_values = [
                    history[tap] if tap != layer_idx else hidden_states for tap in taps
                ]
                if not self.training and self.mudd.uses_r_way:
                    attention_mix, residual_base = self.mudd.forward_pair(
                        tap_values, hidden_states, layer_idx
                    )
                    attention_input = rms_norm(attention_mix)
                else:
                    attention_input = rms_norm(
                        self.mudd(tap_values, hidden_states, layer_idx)
                    )
                    residual_base = (
                        self.mudd(
                            tap_values,
                            hidden_states,
                            layer_idx,
                            r_way=True,
                        )
                        if self.mudd.uses_r_way
                        else hidden_states
                    )
            else:
                attention_input = rms_norm(hidden_states)
                residual_base = hidden_states
            hidden_states, probabilities = decoder_layer(
                hidden_states,
                attention_input,
                value_embeds,
                cosine,
                sine,
                causal_mask_mapping[self.config.layer_types[layer_idx]],
                past_key_values,
                use_cache,
                output_attentions,
                residual_base,
            )
            if output_attentions:
                all_attentions += (probabilities,)
            if layer_idx + 1 in self.retained_taps:
                history[layer_idx + 1] = hidden_states

        hidden_states = self.norm(hidden_states)
        if output_hidden_states:
            all_hidden_states += (hidden_states,)
        if not return_dict:
            values: tuple[Tensor | Cache | tuple[Tensor, ...], ...] = (hidden_states,)
            if use_cache:
                values += (past_key_values,)
            if output_hidden_states:
                values += (all_hidden_states,)
            if output_attentions:
                values += (all_attentions,)
            return values
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
            hidden_states=all_hidden_states if output_hidden_states else None,
            attentions=all_attentions if output_attentions else None,
        )


class LimiteForCausalLM(LimitePreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}

    def __init__(self, config: LimiteConfig) -> None:
        super().__init__(config)
        self.model = LimiteModel(config)
        self.vocab_size = int(config.vocab_size)
        self.lm_head = nn.Linear(
            int(config.hidden_size), int(config.vocab_size), bias=False
        )
        softcap = dict(config.softcap_logits)
        self.softcap_a = float(softcap["a"])
        self.softcap_b = float(softcap["b"])
        self.softcap_c = float(softcap["c"])
        self.head_precision_mode = str(config.lm_head_precision_mode)
        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

    def get_decoder(self) -> LimiteModel:
        return self.model

    def set_decoder(self, decoder: LimiteModel) -> None:
        self.model = decoder

    def _prepare_static_cache(
        self,
        cache_implementation: str,
        batch_size: int,
        max_cache_len: int,
        model_kwargs: dict[str, Any],
    ) -> Cache:
        # GenerationMixin allocates its persistent static cache before the
        # first model forward. Reject the unsupported backend/cache pair here
        # so users receive the Limite contract error instead of failing inside
        # Transformers' cache preparation. SDPA remains entirely native.
        _reject_external_flash_static_cache(self.config._attn_implementation)
        return super()._prepare_static_cache(
            cache_implementation,
            batch_size,
            max_cache_len,
            model_kwargs,
        )

    def _softcapped_logits(self, hidden_states: Tensor) -> Tensor:
        if self.head_precision_mode == "oracle_exact":
            logits = F.linear(hidden_states, self.lm_head.weight).float()
        elif hidden_states.is_cuda and hidden_states.dtype in (
            torch.bfloat16,
            torch.float16,
        ):
            logits = torch.mm(
                hidden_states.reshape(-1, hidden_states.shape[-1]),
                self.lm_head.weight.t(),
                out_dtype=torch.float32,
            ).reshape(*hidden_states.shape[:-1], self.lm_head.weight.shape[0])
        else:
            logits = F.linear(hidden_states.float(), self.lm_head.weight.float())
        return self.softcap_a * torch.sigmoid(
            (logits + self.softcap_b) / self.softcap_c
        )

    def forward(
        self,
        input_ids: Tensor | None = None,
        attention_mask: Tensor | None = None,
        position_ids: Tensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: Tensor | None = None,
        labels: Tensor | None = None,
        use_cache: bool | None = None,
        logits_to_keep: int | Tensor = 0,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **kwargs: Any,
    ) -> CausalLMOutputWithPast | tuple[Tensor, ...]:
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
            **kwargs,
        )
        hidden_states = outputs.last_hidden_state
        indices = (
            slice(-logits_to_keep, None)
            if isinstance(logits_to_keep, int) and logits_to_keep > 0
            else logits_to_keep
            if isinstance(logits_to_keep, Tensor)
            else slice(None)
        )
        logits = self._softcapped_logits(hidden_states[:, indices, :])
        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1].contiguous().float()
            shift_labels = labels[:, 1:].contiguous().to(shift_logits.device)
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )
        if return_dict is False:
            result = (logits, outputs.past_key_values)
            return ((loss,) + result) if loss is not None else result
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


__all__ = [
    "LimiteDecoderLayer",
    "LimiteForCausalLM",
    "LimiteModel",
    "LimitePreTrainedModel",
]