modeling_raon1_mini.py

# SPDX-License-Identifier: Apache-2.0
"""Hugging Face ``transformers`` implementation of raon-1-mini.

Architecture (released 28.8B / 4.3B-active checkpoint): 25 decoder layers on a 4-stream Raon_mHC
residual, 18 Raon Delta Attention (RDA) layers + 7 NoPE gated
Multi-head Latent Attention (MLA) layers, a dense SiTU-GLU MLP in layer 0 and a 128-expert
top-6 sigmoid-routed MoE with one shared expert everywhere else, and an untied output head.

Every equation follows the training implementation's plain-PyTorch reference:

* Raon_mHC selector, router, SiTU nonlinearity, attention softmax, RDA gates and the RDA recurrent
  state are computed in fp32; GEMM activations live in the model dtype.
* RDA prefill/decode use a pure-PyTorch recurrence; no custom kernel package is required.
* The parameter names are the release checkpoint's tensor names.
* The released checkpoint stores routed experts as MXFP8 (``quant_method: mxfp8``) and loads them
  as E4M3 weights + E8M0 scales on every device. On sm_100/sm_103 the routed experts run as native
  grouped MXFP8 GEMMs (activations quantized per call with the production rule, W8A8); elsewhere
  ``routed_linear`` expands each used expert to the activation dtype per call (W8A16). Passing
  ``dequantize=True`` in the quantization config instead dequantizes exactly to BF16 at load.
"""

from __future__ import annotations

import contextlib
import copy
import math
from typing import Any, ClassVar

import torch
import torch.nn.functional as F
from torch import nn
from transformers import initialization as init
from transformers.generation import GenerationMixin
from transformers.integrations.finegrained_fp8 import FP8Linear
from transformers.utils.quantization_config import FineGrainedFP8Config
from transformers.masking_utils import create_causal_mask
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import logging

from .configuration_raon1_mini import Raon1MiniConfig

logger = logging.get_logger(__name__)

# ---------------------------------------------------------------------------------------------
# primitives
# ---------------------------------------------------------------------------------------------
class Raon1MiniRMSNorm(nn.Module):
    """RMSNorm with fp32 statistics, plain (non zero-centered) gamma, no bias."""

    def __init__(self, dim: int, eps: float) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, hidden: torch.Tensor, *, keep_fp32: bool = False) -> torch.Tensor:
        value = hidden.float()
        normed = (
            value
            * torch.rsqrt(value.pow(2).mean(-1, keepdim=True) + self.eps)
            * self.weight.float()
        )
        return normed if keep_fp32 else normed.to(hidden.dtype)

    def extra_repr(self) -> str:
        return f"{tuple(self.weight.shape)}, eps={self.eps}"


def situ_glu(
    gate: torch.Tensor, value: torch.Tensor, gate_beta: float, value_beta: float
) -> torch.Tensor:
    """SiTU-GLU in fp32: ``gate_beta*tanh(g/gate_beta)*sigmoid(g) * value_beta*tanh(v/value_beta)``."""
    g = gate.float()
    v = value.float()
    return (gate_beta * torch.tanh(g / gate_beta) * torch.sigmoid(g)) * (
        value_beta * torch.tanh(v / value_beta)
    )


def helmert_basis(n: int) -> torch.Tensor:
    """Orthonormal ``Q`` in ``R^{n x (n-1)}`` with ``Q^T Q = I`` and ``Q^T 1 = 0``."""
    q = torch.zeros(n, n - 1, dtype=torch.float32)
    for k in range(1, n):
        q[:k, k - 1] = 1.0
        q[k, k - 1] = -float(k)
        q[:, k - 1] /= math.sqrt(k * (k + 1.0))
    return q


_RAON_MHC_CONSTANTS: dict[
    tuple[int, torch.device], tuple[torch.Tensor, torch.Tensor, torch.Tensor]
] = {}


def _raon_mhc_constants(
    n: int, device: torch.device
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """``(helmert Q [n, n-1], I [n, n], B = (I + 11^T/n)/2 [n, n])`` in fp32 on ``device``.

    Built on demand instead of as module buffers: ``from_pretrained`` materialises non-persistent buffers
    on the meta device without values.
    """
    key = (n, device)
    if key not in _RAON_MHC_CONSTANTS:
        eye = torch.eye(n, dtype=torch.float32, device=device)
        _RAON_MHC_CONSTANTS[key] = (
            helmert_basis(n).to(device),
            eye,
            0.5 * (eye + torch.full_like(eye, 1.0 / n)),
        )
    return _RAON_MHC_CONSTANTS[key]


# ---------------------------------------------------------------------------------------------
# Raon_mHC multi-stream residual
# ---------------------------------------------------------------------------------------------
class RaonMHCStreamMix(nn.Module):
    """The ``res.*`` parameters of one Raon_mHC site: the stream-to-stream residual chart."""

    def __init__(self, num_streams: int, hidden_size: int) -> None:
        super().__init__()
        m = (num_streams - 1) ** 2
        self.phi_z = nn.Parameter(torch.zeros(m, num_streams * hidden_size))
        self.alpha_z = nn.Parameter(torch.ones(()))
        self.b_z = nn.Parameter(torch.zeros(m))
        self.eta = nn.Parameter(torch.zeros((), dtype=torch.float32))

    def forward(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        return self.phi_z, self.alpha_z, self.b_z, self.eta


class RaonMHC(nn.Module):
    """One Raon_mHC site (attention or MLP): selects the branch input from the streams and writes back.

    ``route`` returns fp32 coefficients ``h_pre [N, n]``, ``h_post [N, n]`` and ``h_res [N, n, n]``
    for ``N = batch*seq`` rows (Raon_mHC chart:
    ``E = Q Z Q^T``, ``rho = max(0, max(-E/B))``, ``Phi = B + E/(1+rho)``, ``H = (1-gamma) I + gamma Phi``
    with ``B = (I + 11^T/n)/2`` and ``gamma = sigmoid(eta)``).
    """

    def __init__(self, config: Raon1MiniConfig) -> None:
        super().__init__()
        n, c = config.raon_mhc_streams, config.hidden_size
        self.num_streams = n
        self.eps = config.raon_mhc_selector_eps
        self.post_factor = config.raon_mhc_post_factor
        self.phi_pre = nn.Parameter(torch.zeros(n, n * c))
        self.phi_post = nn.Parameter(torch.zeros(n, n * c))
        self.alpha_pre = nn.Parameter(torch.ones(()))
        self.alpha_post = nn.Parameter(torch.ones(()))
        self.b_pre = nn.Parameter(torch.zeros(n))
        self.b_post = nn.Parameter(torch.zeros(n))
        self.res = RaonMHCStreamMix(n, c)

    def route(
        self, stream: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        return self(stream)

    def forward(
        self, stream: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        phi_z, alpha_z, b_z, eta = self.res()
        n = self.num_streams
        rows = stream.shape[0] * stream.shape[1]
        flat = stream.reshape(rows, n * stream.shape[-1])
        weight = torch.cat([self.phi_pre, self.phi_post, phi_z], dim=0)
        # RMS scale applied AFTER the (model-dtype) selector GEMM, as in training.
        inv = torch.rsqrt(flat.float().pow(2).mean(-1, keepdim=True) + self.eps)
        z = F.linear(flat, weight.to(flat.dtype)).float() * inv
        t_pre = self.alpha_pre.float() * z[:, :n] + self.b_pre.float()
        t_post = self.alpha_post.float() * z[:, n : 2 * n] + self.b_post.float()
        t_z = alpha_z.float() * z[:, 2 * n :] + b_z.float()
        h_pre = torch.sigmoid(t_pre)
        h_post = self.post_factor * torch.sigmoid(t_post)
        nm = n - 1
        helmert, eye, chart_b = _raon_mhc_constants(n, z.device)
        e = torch.einsum("im,bmk,jk->bij", helmert, t_z.reshape(rows, nm, nm), helmert)
        rho = (-e / chart_b).amax(dim=(-2, -1)).clamp_min(0.0)
        phi = chart_b + e / (1.0 + rho)[:, None, None]
        gamma = torch.sigmoid(eta.float())
        h_res = (1.0 - gamma) * eye + gamma * phi
        return h_pre, h_post, h_res

    @staticmethod
    def aggregate(
        stream_f32: torch.Tensor, h_pre: torch.Tensor, dtype: torch.dtype
    ) -> torch.Tensor:
        """``[N, n, C]`` fp32 streams -> ``[N, C]`` branch input in the model dtype."""
        return torch.einsum("bn,bnc->bc", h_pre, stream_f32).to(dtype)

    @staticmethod
    def mix(
        stream_f32: torch.Tensor,
        branch_out: torch.Tensor,
        h_post: torch.Tensor,
        h_res: torch.Tensor,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        """Residual mixing plus the branch write-back, rounded once to the model dtype."""
        mixed = torch.einsum("bij,bjc->bic", h_res, stream_f32)
        return (mixed + h_post.unsqueeze(-1) * branch_out.float().unsqueeze(1)).to(
            dtype
        )


# ---------------------------------------------------------------------------------------------
# RDA (Raon Delta Attention)
# ---------------------------------------------------------------------------------------------
class Raon1MiniShortConv(nn.Module):
    """Depthwise causal conv (no bias) whose weight is stored squeezed as ``[channels, kernel]``."""

    def __init__(self, channels: int, kernel_size: int) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(channels, kernel_size))
        self.kernel_size = kernel_size

    def forward(
        self, x: torch.Tensor, state: torch.Tensor | None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """``x [B, T, C]`` raw rows -> silu(conv) ``[B, T, C]`` in x's dtype, plus the new raw window ``[B, C, K-1]``."""
        channels = x.shape[-1]
        k = self.kernel_size
        signal = x.transpose(1, 2)  # [B, C, T]
        padded = (
            F.pad(signal, (k - 1, 0))
            if state is None
            else torch.cat([state.to(signal.dtype), signal], dim=-1)
        )
        # fp32 accumulation and silu, one rounding on the output (causal_conv1d semantics).
        out = F.conv1d(
            padded.float(), self.weight.float().unsqueeze(1), groups=channels
        )
        new_state = padded[..., -(k - 1) :]
        return F.silu(out).to(x.dtype).transpose(1, 2), new_state


def torch_recurrent_rda(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    log_alpha: torch.Tensor,
    beta: torch.Tensor,
    state: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """V-first delta rule, token serial, fp32 state ``[B, H, V, K]``. Reference for the fused kernels.

    ``q/k/v [B, T, H, D]`` (q already l2-normalised and scaled, k l2-normalised), ``log_alpha [B, T, H, K]``
    fp32 log decay, ``beta [B, T, H]``. Returns ``out [B, T, H, V]`` in v's dtype and the final state.
    """
    batch, seq, heads, key_dim = k.shape
    value_dim = v.shape[-1]
    if state is None:
        state = torch.zeros(
            batch, heads, value_dim, key_dim, dtype=torch.float32, device=q.device
        )
    else:
        state = state.float()
    out = torch.empty(
        batch, seq, heads, value_dim, dtype=torch.float32, device=q.device
    )
    qf, kf, vf, bf = q.float(), k.float(), v.float(), beta.float()
    for t in range(seq):
        state = state * torch.exp(log_alpha[:, t]).unsqueeze(-2)
        predicted = torch.einsum("bhvk,bhk->bhv", state, kf[:, t])
        error = vf[:, t] - predicted
        state = state + bf[:, t][..., None, None] * torch.einsum(
            "bhv,bhk->bhvk", error, kf[:, t]
        )
        out[:, t] = torch.einsum("bhvk,bhk->bhv", state, qf[:, t])
    return out.to(v.dtype), state


class Raon1MiniRDA(nn.Module):
    """Raon Delta Attention layer: gated delta-rule linear attention with short convolutions and a write gate."""

    def __init__(self, config: Raon1MiniConfig) -> None:
        super().__init__()
        hidden = config.hidden_size
        self.num_heads = config.rda_num_heads
        self.head_dim = config.rda_head_dim
        self.proj_size = self.num_heads * self.head_dim
        self.gate_lower_bound = config.rda_gate_lower_bound
        self.norm_eps = config.rms_norm_eps
        kernel = config.rda_conv_kernel_size
        self.q_proj = nn.Linear(hidden, self.proj_size, bias=False)
        self.k_proj = nn.Linear(hidden, self.proj_size, bias=False)
        self.v_proj = nn.Linear(hidden, self.proj_size, bias=False)
        self.q_conv1d = Raon1MiniShortConv(self.proj_size, kernel)
        self.k_conv1d = Raon1MiniShortConv(self.proj_size, kernel)
        self.v_conv1d = Raon1MiniShortConv(self.proj_size, kernel)
        self.f_a_proj = nn.Linear(hidden, config.rda_decay_rank, bias=False)
        self.f_b_proj = nn.Linear(config.rda_decay_rank, self.proj_size, bias=False)
        self.b_proj = nn.Linear(hidden, self.num_heads, bias=False)
        self.rda_w_proj = nn.Linear(hidden, self.num_heads, bias=False)
        self.g_proj = nn.Linear(hidden, self.proj_size, bias=False)
        self.A_log = nn.Parameter(torch.zeros(self.num_heads, dtype=torch.float32))
        self.dt_bias = nn.Parameter(torch.zeros(self.proj_size, dtype=torch.float32))
        self.o_norm = Raon1MiniRMSNorm(self.head_dim, config.rms_norm_eps)
        self.o_proj = nn.Linear(self.proj_size, hidden, bias=False)

    def forward(
        self,
        hidden: torch.Tensor,
        cache: Raon1MiniDynamicCache | None,
        layer_idx: int,
        token_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        batch, seq, _ = hidden.shape
        dtype = hidden.dtype
        if token_mask is not None:
            # Left padding: padded rows must not enter the conv window or the recurrence.
            hidden = hidden * token_mask[..., None].to(dtype)
        q = self.q_proj(hidden)
        k = self.k_proj(hidden)
        v = self.v_proj(hidden)
        gate = self.g_proj(hidden)
        beta_raw = self.b_proj(hidden)
        rda_delta = self.rda_w_proj(hidden)
        decay_raw = self.f_b_proj(self.f_a_proj(hidden))

        conv_state = cache.conv_states[layer_idx] if cache is not None else None
        q_hist = k_hist = v_hist = None
        if conv_state is not None:
            q_hist, k_hist, v_hist = conv_state.split(self.proj_size, dim=1)
        q, q_hist = self.q_conv1d(q, q_hist)
        k, k_hist = self.k_conv1d(k, k_hist)
        v, v_hist = self.v_conv1d(v, v_hist)
        if cache is not None:
            cache.conv_states[layer_idx] = torch.cat([q_hist, k_hist, v_hist], dim=1)

        q = q.reshape(batch, seq, self.num_heads, self.head_dim)
        k = k.reshape(batch, seq, self.num_heads, self.head_dim)
        v = v.reshape(batch, seq, self.num_heads, self.head_dim)
        # _prepare_after_conv: l2-normalise q/k in fp32 (eps 1e-6), scale q by key_dim**-0.5.
        q = F.normalize(q.float(), dim=-1, eps=1e-6).to(dtype) * (self.head_dim**-0.5)
        k = F.normalize(k.float(), dim=-1, eps=1e-6).to(dtype)
        # Log decay (per-head scale, lower bound): fp32.
        decay = decay_raw.float().reshape(
            batch, seq, self.num_heads, self.head_dim
        ) + self.dt_bias.float().reshape(self.num_heads, self.head_dim)
        scaled = torch.exp(self.A_log.float())[None, None, :, None] * decay
        log_alpha = (
            -F.softplus(scaled)
            if self.gate_lower_bound is None
            else self.gate_lower_bound * torch.sigmoid(scaled)
        )
        beta = torch.sigmoid(beta_raw.float()).to(dtype)
        # RDA: write gate sigmoid(beta+delta), erase gate sigmoid(beta); realised as a fp32 value scale.
        ratio = torch.exp(
            F.logsigmoid(beta_raw.float() + rda_delta.float())
            - F.logsigmoid(beta_raw.float())
        )
        v = (v.float() * ratio.unsqueeze(-1)).to(dtype)

        recurrent_state = (
            cache.recurrent_states[layer_idx] if cache is not None else None
        )
        out, final_state = self._core(
            q, k, v, log_alpha.contiguous(), beta, recurrent_state
        )
        if cache is not None:
            cache.recurrent_states[layer_idx] = final_state
        # Per-head RMSNorm then sigmoid output gate, both in fp32, one rounding.
        normed = self.o_norm(out.reshape(-1, self.head_dim), keep_fp32=True)
        gated = (normed * torch.sigmoid(gate.float().reshape(-1, self.head_dim))).to(
            dtype
        )
        return self.o_proj(gated.reshape(batch, seq, self.proj_size))

    def _core(self, q, k, v, log_alpha, beta, state):
        return torch_recurrent_rda(q, k, v, log_alpha, beta, state)


# ---------------------------------------------------------------------------------------------
# NoPE gated MLA
# ---------------------------------------------------------------------------------------------
class Raon1MiniMLA(nn.Module):
    """Multi-head latent attention without positional encoding and with a sigmoid output gate.

    The 64 "rope" channels are never rotated: they are extra key/query channels shared across heads
    (``k_pe`` is expanded over heads), so scores are 192-wide dot products scaled by ``1/sqrt(192)``.
    """

    def __init__(self, config: Raon1MiniConfig) -> None:
        super().__init__()
        hidden = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.qk_nope = config.qk_nope_head_dim
        self.qk_pe = config.qk_rope_head_dim
        self.v_dim = config.v_head_dim
        self.q_head_dim = self.qk_nope + self.qk_pe
        self.kv_lora_rank = config.kv_lora_rank
        self.softmax_scale = self.q_head_dim**-0.5
        self.use_output_gate = config.mla_use_output_gate
        self.q_a_proj = nn.Linear(hidden, config.q_lora_rank, bias=False)
        self.q_a_layernorm = Raon1MiniRMSNorm(
            config.q_lora_rank, config.mla_latent_norm_eps
        )
        self.q_b_proj = nn.Linear(
            config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
        )
        self.kv_a_proj_with_mqa = nn.Linear(
            hidden, config.kv_lora_rank + self.qk_pe, bias=False
        )
        self.kv_a_layernorm = Raon1MiniRMSNorm(
            config.kv_lora_rank, config.mla_latent_norm_eps
        )
        self.kv_b_proj = nn.Linear(
            config.kv_lora_rank,
            self.num_heads * (self.qk_nope + self.v_dim),
            bias=False,
        )
        if self.use_output_gate:
            self.g_proj = nn.Linear(hidden, self.num_heads * self.v_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.v_dim, hidden, bias=False)

    def forward(
        self,
        hidden: torch.Tensor,
        cache: Raon1MiniDynamicCache | None,
        layer_idx: int,
        attention_mask: torch.Tensor | None,
    ) -> torch.Tensor:
        batch, seq, _ = hidden.shape
        dtype = hidden.dtype
        q_lat = self.q_a_layernorm(self.q_a_proj(hidden))
        q = self.q_b_proj(q_lat).reshape(batch, seq, self.num_heads, self.q_head_dim)
        kv_full = self.kv_a_proj_with_mqa(hidden)
        kv_lat, k_pe = torch.split(kv_full, [self.kv_lora_rank, self.qk_pe], dim=-1)
        kv = self.kv_b_proj(self.kv_a_layernorm(kv_lat)).reshape(
            batch, seq, self.num_heads, self.qk_nope + self.v_dim
        )
        k_nope, value = torch.split(kv, [self.qk_nope, self.v_dim], dim=-1)
        key = torch.cat(
            [k_nope, k_pe.unsqueeze(2).expand(batch, seq, self.num_heads, self.qk_pe)],
            dim=-1,
        )

        query = q.transpose(1, 2)  # [B, H, T, 192]
        key = key.transpose(1, 2)
        value = value.transpose(1, 2)
        if cache is not None:
            key, value = cache.update(key, value, layer_idx)
        kv_len = key.shape[2]
        if attention_mask is None and seq > 1 and kv_len != seq:
            # Chunked prefill after cached tokens: explicit causal mask with the cache offset.
            offset = kv_len - seq
            attention_mask = torch.ones(
                seq, kv_len, dtype=torch.bool, device=hidden.device
            ).tril(offset)[None, None]
        context = F.scaled_dot_product_attention(
            query,
            key,
            value,
            attn_mask=attention_mask,
            is_causal=attention_mask is None and seq > 1,
            scale=self.softmax_scale,
        )
        context = context.transpose(1, 2).reshape(
            batch, seq, self.num_heads * self.v_dim
        )
        if self.use_output_gate:
            # The training NoPE path evaluates the gate in the activation dtype.
            context = context * torch.sigmoid(self.g_proj(hidden))
        return self.o_proj(context.to(dtype))


# ---------------------------------------------------------------------------------------------
# MLPs
# ---------------------------------------------------------------------------------------------
class Raon1MiniMLP(nn.Module):
    """SiTU-GLU MLP: ``down(situ_glu(gate(x), up(x)))``. Used for the dense layer and the shared expert."""

    def __init__(self, config: Raon1MiniConfig, intermediate_size: int) -> None:
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
        self.gate_beta = config.activation_situ_beta
        self.value_beta = config.activation_situ_linear_beta

    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        act = situ_glu(
            self.gate_proj(hidden),
            self.up_proj(hidden),
            self.gate_beta,
            self.value_beta,
        )
        return self.down_proj(act.to(hidden.dtype))


def dequantize_mxfp8(
    weight: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype
) -> torch.Tensor:
    """E4M3 ``[N, K]`` and E8M0 bytes ``[N, K/32]`` -> ``value * 2**(scale - 127)`` in ``dtype``.

    The fp32 product is exact and every released weight is representable in BF16, so the cast is exact.
    """
    rows, cols = weight.shape
    blocks = weight.float().view(rows, cols // 32, 32)
    factor = torch.exp2(_scale_bytes(scale).float() - 127.0).unsqueeze(-1)
    return (blocks * factor).view(rows, cols).to(dtype)


def _scale_bytes(scale: torch.Tensor) -> torch.Tensor:
    """E8M0 scales as their raw exponent bytes, whether stored as uint8 or ``float8_e8m0fnu``."""
    return scale if scale.dtype == torch.uint8 else scale.view(torch.uint8)


def quantize_mxfp8(hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """``[M, K]`` -> E4M3 ``[M, K]`` and E8M0 bytes ``[M, K/32]`` with the production (TE) rule.

    The block exponent is ``amax / 448`` rounded up to a power of two on its fp32 bits, and the
    values are ``x * 2**(127 - e)`` cast round-to-nearest with saturation, as vLLM serving does.
    """
    rows, cols = hidden.shape
    blocks = hidden.float().view(rows, cols // 32, 32)
    bits = (blocks.abs().amax(-1) * (1.0 / 448.0)).view(torch.int32)
    exp = (bits >> 23) & 0xFF
    man = bits & 0x7FFFFF
    round_up = (man > 0) & (exp != 0xFE) & ~((exp == 0) & (man <= 0x400000))
    e = (exp + round_up.int()).clamp(max=254)
    # 2**(127 - e) built from its bits; e == 254 needs the subnormal 2**-127.
    reciprocal = ((254 - e) << 23).view(torch.float32)
    reciprocal = torch.where(e == 254, torch.full_like(reciprocal, 2.0**-127), reciprocal)
    values = (blocks * reciprocal.unsqueeze(-1)).clamp(-448.0, 448.0)
    return values.to(torch.float8_e4m3fn).view(rows, cols), e.to(torch.uint8)


def to_blocked_scales(scale: torch.Tensor) -> torch.Tensor:
    """``[rows, cols]`` E8M0 bytes -> cuBLAS 128x4 swizzled layout (rows padded to 128, cols to 4)."""
    rows, cols = scale.shape
    row_tiles, col_tiles = -(-rows // 128), -(-cols // 4)
    padded = scale.new_zeros(row_tiles * 128, col_tiles * 4)
    padded[:rows, :cols] = scale
    tiles = padded.view(row_tiles, 128, col_tiles, 4).permute(0, 2, 1, 3)
    return tiles.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16).flatten()


def _grouped_activation_scales(
    scale: torch.Tensor, tokens_per_expert: torch.Tensor, cols: int
) -> torch.Tensor:
    """Expert-sorted activation scales -> each expert's rows padded to 128, then swizzled.

    ``scaled_grouped_mm`` reads group ``g``'s scales from a 128-row aligned tile range, so the
    rows are scattered to their padded positions without a host sync.
    """
    groups = tokens_per_expert.shape[0]
    rows = scale.shape[0]
    padded = (tokens_per_expert + 127) // 128 * 128
    padded_start = torch.cumsum(padded, 0) - padded
    start = torch.cumsum(tokens_per_expert, 0) - tokens_per_expert
    expert = torch.repeat_interleave(
        torch.arange(groups, device=scale.device), tokens_per_expert, output_size=rows
    )
    dest = padded_start[expert] + torch.arange(rows, device=scale.device) - start[expert]
    buffer = scale.new_zeros(rows + groups * 128, cols)
    buffer[dest, : scale.shape[1]] = scale
    return to_blocked_scales(buffer).view(-1, cols)


def mxfp8_grouped_mm(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    tokens_per_expert: torch.Tensor,
    offsets: torch.Tensor,
) -> torch.Tensor:
    """Expert-sorted ``[M, K]`` x E4M3 ``[E, N, K]`` (swizzled scales ``[E, *]``) -> ``[M, N]``.

    Activations are quantized to MXFP8 per call and both operands run through the native
    block-scaled tensor cores (W8A8, fp32 accumulation, output in the activation dtype).
    """
    values, scale = quantize_mxfp8(hidden)
    cols = -(-scale.shape[1] // 4) * 4
    act_scale = _grouped_activation_scales(scale, tokens_per_expert, cols)
    return F.scaled_grouped_mm(
        values,
        weight.transpose(-2, -1),
        act_scale.view(torch.float8_e8m0fnu),
        F.ScalingType.BlockWise1x32,
        weight_scale.view(torch.float8_e8m0fnu),
        F.ScalingType.BlockWise1x32,
        swizzle_a=F.SwizzleType.SWIZZLE_32_4_4,
        swizzle_b=F.SwizzleType.SWIZZLE_32_4_4,
        offs=offsets,
        output_dtype=hidden.dtype,
    )


def native_mxfp8_available(device: torch.device) -> bool:
    """Block-scaled grouped GEMM needs datacenter Blackwell (sm_100/sm_103) and torch's MX API."""
    return (
        device.type == "cuda"
        and torch.cuda.get_device_capability(device)[0] == 10
        and hasattr(F, "scaled_grouped_mm")
        and hasattr(getattr(F, "ScalingType", None), "BlockWise1x32")
        and hasattr(getattr(F, "SwizzleType", None), "SWIZZLE_32_4_4")
    )


def routed_linear(linear: nn.Linear, hidden: torch.Tensor) -> torch.Tensor:
    """Routed-expert projection in the activation dtype, also when MXFP8 storage was kept."""
    weight = linear.weight
    if weight.dtype == torch.float8_e4m3fn:
        weight = dequantize_mxfp8(weight, linear.weight_scale_inv, hidden.dtype)
    return F.linear(hidden, weight)


class Raon1MiniRoutedLinear(nn.Linear):
    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        return routed_linear(self, hidden)


class Raon1MiniMXFP8Linear(FP8Linear):
    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        return routed_linear(self, hidden)


def _routed_projection(config: Raon1MiniConfig, in_features: int, out_features: int) -> nn.Linear:
    quantization = getattr(config, "quantization_config", None)
    if isinstance(quantization, dict):
        quantization = FineGrainedFP8Config.from_dict(quantization)
    if getattr(quantization, "quant_method", None) == "mxfp8":
        if not quantization.dequantize:
            return Raon1MiniMXFP8Linear(
                in_features, out_features,
                block_size=quantization.weight_block_size,
                activation_scheme=quantization.activation_scheme,
                scale_fmt="ue8m0",
                has_bias=False,
            )
    return Raon1MiniRoutedLinear(in_features, out_features, bias=False)


class Raon1MiniExpert(nn.Module):
    """One routed SiTU-GLU expert: ``down_proj(situ_glu(gate_proj(x), up_proj(x)))``."""

    def __init__(self, config: Raon1MiniConfig) -> None:
        super().__init__()
        self.gate_proj = _routed_projection(config, config.hidden_size, config.moe_intermediate_size)
        self.up_proj = _routed_projection(config, config.hidden_size, config.moe_intermediate_size)
        self.down_proj = _routed_projection(config, config.moe_intermediate_size, config.hidden_size)
        self.gate_beta = config.activation_situ_beta
        self.value_beta = config.activation_situ_linear_beta

    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        act = situ_glu(
            self.gate_proj(hidden),
            self.up_proj(hidden),
            self.gate_beta,
            self.value_beta,
        )
        return self.down_proj(act.to(hidden.dtype))


class Raon1MiniRouter(nn.Module):
    """Sigmoid router with the aux-loss-free expert bias (``e_score_correction_bias``) used for selection only."""

    def __init__(self, config: Raon1MiniConfig) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(config.num_experts, config.hidden_size))
        self.e_score_correction_bias = nn.Parameter(
            torch.zeros(config.num_experts, dtype=torch.float32)
        )
        self.top_k = config.num_experts_per_token
        self.renormalize = config.moe_renormalize
        self.scaling = config.routed_scaling_factor
        if config.moe_router_activation_func != "sigmoid":
            raise NotImplementedError("raon-1-mini uses the sigmoid router.")

    def forward(self, hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """``[N, C]`` -> (top-k expert indices ``[N, k]``, fp32 routing weights ``[N, k]``)."""
        logits = F.linear(hidden.float(), self.weight.float())
        scores = torch.sigmoid(logits)
        _, top = torch.topk(
            scores + self.e_score_correction_bias.float(),
            self.top_k,
            dim=-1,
            sorted=True,
        )
        chosen = torch.gather(scores, 1, top)
        if self.renormalize:
            chosen = chosen / (chosen.sum(-1, keepdim=True) + 1e-20)
        return top, chosen * self.scaling


class Raon1MiniMoE(nn.Module):
    """Routed experts plus the always-on shared expert.

    With MXFP8 storage on sm_100/sm_103 the routed experts run as two native grouped MXFP8 GEMMs;
    everywhere else (or after ``set_mxfp8_moe_backend("dequantize")``) each used expert is expanded to
    the activation dtype per call, so FP8 storage is kept on every device.
    """

    def __init__(self, config: Raon1MiniConfig) -> None:
        super().__init__()
        self.num_experts = config.num_experts
        self.top_k = config.num_experts_per_token
        self.intermediate_size = config.moe_intermediate_size
        self.gate_beta = config.activation_situ_beta
        self.value_beta = config.activation_situ_linear_beta
        self.gate = Raon1MiniRouter(config)
        self.experts = nn.ModuleList(
            [Raon1MiniExpert(config) for _ in range(config.num_experts)]
        )
        shared_size = config.moe_intermediate_size * config.num_shared_experts
        self.shared_experts = (
            Raon1MiniMLP(config, shared_size) if config.num_shared_experts else None
        )
        # Toggled model-wide by ``Raon1MiniPreTrainedModel.set_mxfp8_moe_backend``.
        self.use_native_mxfp8 = True
        self._packed: tuple[torch.Tensor, ...] | None = None

    def _native_weights(self, hidden: torch.Tensor) -> tuple[torch.Tensor, ...] | None:
        """Stacked E4M3 ``gate_up [E, 2I, H]`` / ``down [E, H, I]`` plus swizzled scales, or None.

        The per-expert parameters are rebound to views of the stacks on first use, so the layer
        keeps one copy of its weights and the checkpoint names are unchanged.
        """
        if not self.use_native_mxfp8:
            return None
        if self._packed is not None:
            return self._packed
        first = self.experts[0].gate_proj.weight
        if (
            first.dtype != torch.float8_e4m3fn
            or hidden.dtype != torch.bfloat16
            or first.device != hidden.device
            or not native_mxfp8_available(hidden.device)
            or any(p.device != first.device for p in self.experts.parameters())
        ):
            return None
        size = self.intermediate_size
        with torch.no_grad():
            gate_up = torch.stack(
                [torch.cat([e.gate_proj.weight, e.up_proj.weight]) for e in self.experts]
            )
            gate_up_scale = torch.stack(
                [torch.cat([_scale_bytes(e.gate_proj.weight_scale_inv), _scale_bytes(e.up_proj.weight_scale_inv)])
                 for e in self.experts]
            )
            down = torch.stack([e.down_proj.weight for e in self.experts])
            down_scale = torch.stack([_scale_bytes(e.down_proj.weight_scale_inv) for e in self.experts])
            for idx, expert in enumerate(self.experts):
                expert.gate_proj.weight.data = gate_up[idx, :size]
                expert.up_proj.weight.data = gate_up[idx, size:]
                expert.down_proj.weight.data = down[idx]
                expert.gate_proj.weight_scale_inv.data = gate_up_scale[idx, :size]
                expert.up_proj.weight_scale_inv.data = gate_up_scale[idx, size:]
                expert.down_proj.weight_scale_inv.data = down_scale[idx]
            gate_up_blocked = torch.stack([to_blocked_scales(s) for s in gate_up_scale])
            down_blocked = torch.stack([to_blocked_scales(s) for s in down_scale])
        self._packed = (gate_up, gate_up_blocked, down, down_blocked)
        return self._packed

    def _native_routed(
        self,
        flat: torch.Tensor,
        top: torch.Tensor,
        probs: torch.Tensor,
        packed: tuple[torch.Tensor, ...],
    ) -> torch.Tensor:
        """Sort the token-expert pairs by expert, run up/gate and down as grouped MXFP8 GEMMs."""
        gate_up, gate_up_blocked, down, down_blocked = packed
        experts = top.flatten()
        order = torch.argsort(experts, stable=True)
        token = order // self.top_k
        tokens_per_expert = torch.bincount(experts, minlength=self.num_experts).int()
        offsets = torch.cumsum(tokens_per_expert, 0, dtype=torch.int32)
        hidden = mxfp8_grouped_mm(
            flat.index_select(0, token), gate_up, gate_up_blocked, tokens_per_expert, offsets
        )
        act = situ_glu(
            hidden[:, : self.intermediate_size],
            hidden[:, self.intermediate_size :],
            self.gate_beta,
            self.value_beta,
        ).to(flat.dtype)
        expert_out = mxfp8_grouped_mm(act, down, down_blocked, tokens_per_expert, offsets)
        # Back to (token, slot) order and a fixed-order sum over the k slots, so the result is
        # deterministic (an index_add_ with k duplicate rows per token would race on atomics).
        weighted = torch.empty(experts.shape[0], flat.shape[1], dtype=torch.float32, device=flat.device)
        weighted[order] = expert_out.float() * probs.flatten().index_select(0, order).unsqueeze(-1)
        return weighted.view(flat.shape[0], self.top_k, flat.shape[1]).sum(1)

    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        batch, seq, dim = hidden.shape
        flat = hidden.reshape(-1, dim)
        top, probs = self.gate(flat)
        packed = self._native_weights(flat) if flat.shape[0] else None
        if packed is not None:
            out = self._native_routed(flat, top, probs, packed).to(hidden.dtype)
            if self.shared_experts is not None:
                out = out + self.shared_experts(flat)
            return out.reshape(batch, seq, dim)
        out = torch.zeros(flat.shape[0], dim, dtype=torch.float32, device=flat.device)
        hit = torch.bincount(top.flatten(), minlength=self.num_experts)
        for expert_idx in hit.nonzero().flatten().tolist():
            rows, slot = torch.where(top == expert_idx)
            expert_out = self.experts[expert_idx](flat.index_select(0, rows))
            out.index_add_(
                0, rows, expert_out.float() * probs[rows, slot].unsqueeze(-1)
            )
        out = out.to(hidden.dtype)
        if self.shared_experts is not None:
            out = out + self.shared_experts(flat)
        return out.reshape(batch, seq, dim)


# ---------------------------------------------------------------------------------------------
# decoder layer
# ---------------------------------------------------------------------------------------------
class Raon1MiniDecoderLayer(nn.Module):
    """One Raon_mHC-wired layer: attention site then MLP site, each with its own selector."""

    def __init__(self, config: Raon1MiniConfig, layer_idx: int) -> None:
        super().__init__()
        self.layer_idx = layer_idx
        self.is_rda = config.is_rda_layer(layer_idx)
        self.is_moe = config.is_moe_layer(layer_idx)
        self.layer_type = "linear_attention" if self.is_rda else "full_attention"
        self.input_layernorm = Raon1MiniRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_attention_layernorm = Raon1MiniRMSNorm(
            config.hidden_size, config.rms_norm_eps
        )
        self.self_attn = Raon1MiniRDA(config) if self.is_rda else Raon1MiniMLA(config)
        self.self_attn_raon_mhc = RaonMHC(config)
        self.mlp_raon_mhc = RaonMHC(config)
        if self.is_moe:
            self.moe = Raon1MiniMoE(config)
        else:
            self.mlp = Raon1MiniMLP(config, config.intermediate_size)

    def forward(
        self,
        stream: torch.Tensor,
        *,
        attention_mask: torch.Tensor | None,
        token_mask: torch.Tensor | None,
        cache: Raon1MiniDynamicCache | None,
    ) -> torch.Tensor:
        batch, seq, n, dim = stream.shape
        dtype = stream.dtype
        rows = stream.reshape(batch * seq, n, dim)

        h_pre, h_post, h_res = self.self_attn_raon_mhc.route(stream)
        sf = rows.float()
        branch_in = self.input_layernorm(
            RaonMHC.aggregate(sf, h_pre, dtype)
        ).reshape(batch, seq, dim)
        if self.is_rda:
            branch_out = self.self_attn(branch_in, cache, self.layer_idx, token_mask)
        else:
            branch_out = self.self_attn(
                branch_in, cache, self.layer_idx, attention_mask
            )
        rows = RaonMHC.mix(
            sf, branch_out.reshape(batch * seq, dim), h_post, h_res, dtype
        )

        h_pre, h_post, h_res = self.mlp_raon_mhc.route(rows.reshape(batch, seq, n, dim))
        sf = rows.float()
        branch_in = self.post_attention_layernorm(
            RaonMHC.aggregate(sf, h_pre, dtype)
        ).reshape(batch, seq, dim)
        branch_out = (
            self.moe(branch_in) if self.is_moe else self.mlp(branch_in)
        )
        rows = RaonMHC.mix(
            sf, branch_out.reshape(batch * seq, dim), h_post, h_res, dtype
        )
        return rows.reshape(batch, seq, n, dim)


# ---------------------------------------------------------------------------------------------
# cache
# ---------------------------------------------------------------------------------------------
class Raon1MiniDynamicCache:
    """Hybrid cache: RDA conv windows + fp32 recurrent states, MLA key/value tensors.

    Per layer: ``conv_states[i]`` is ``[B, 3*proj, K-1]`` of RAW (pre-conv) q|k|v rows,
    ``recurrent_states[i]`` is ``[B, H, V, K]`` fp32; ``key_cache[i]`` / ``value_cache[i]`` are
    ``[B, H, T, 192]`` / ``[B, H, T, 128]``. Implements the hybrid-cache interface ``generate`` uses.
    """

    is_compileable = False

    def __init__(self, config: Raon1MiniConfig) -> None:
        self.layer_types = config.layer_types
        self.transformer_layers = [
            i for i, t in enumerate(self.layer_types) if t == "full_attention"
        ]
        self.last_linear_layer = (
            len(self.layer_types) - 1 - self.layer_types[::-1].index("linear_attention")
        )
        num = config.num_hidden_layers
        self.conv_states: list[torch.Tensor | None] = [None] * num
        self.recurrent_states: list[torch.Tensor | None] = [None] * num
        self.key_cache: list[torch.Tensor | None] = [None] * num
        self.value_cache: list[torch.Tensor | None] = [None] * num

    def __len__(self) -> int:
        return len(self.layer_types)

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs: dict | None = None,
    ):
        if self.key_cache[layer_idx] is None:
            self.key_cache[layer_idx] = key_states
            self.value_cache[layer_idx] = value_states
        else:
            self.key_cache[layer_idx] = torch.cat(
                [self.key_cache[layer_idx], key_states], dim=2
            )
            self.value_cache[layer_idx] = torch.cat(
                [self.value_cache[layer_idx], value_states], dim=2
            )
        return self.key_cache[layer_idx], self.value_cache[layer_idx]

    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
        for i in range(len(self.layer_types)):
            for store in (
                self.key_cache,
                self.value_cache,
                self.conv_states,
                self.recurrent_states,
            ):
                if store[i] is not None:
                    store[i] = store[i].index_select(0, beam_idx.to(store[i].device))

    def get_seq_length(self, layer_idx: int | None = 0) -> int:
        layer_idx = (
            self.transformer_layers[0]
            if layer_idx not in self.transformer_layers
            else layer_idx
        )
        if self.key_cache[layer_idx] is None:
            return 0
        return self.key_cache[layer_idx].shape[-2]

    def get_mask_sizes(
        self, query_length: int, layer_idx: int
    ) -> tuple[int, int]:
        return query_length + self.get_seq_length(layer_idx), 0

    def get_query_offset(self, layer_idx: int = 0) -> int:
        return self.get_seq_length(layer_idx)

    @property
    def has_previous_state(self) -> bool:
        return self.conv_states[self.last_linear_layer] is not None


# ---------------------------------------------------------------------------------------------
# models
# ---------------------------------------------------------------------------------------------
@contextlib.contextmanager
def _keep_mxfp8_storage():
    """Keep MXFP8 storage on devices where Transformers would force a load-time dequantize.

    The fine-grained FP8 quantizer switches to ``dequantize=True`` without a GPU or below
    compute capability 8.9 because its own kernels need FP8 tensor cores. Raon never calls those
    kernels (``Raon1MiniMoE`` reads the E4M3 weights and scales itself and expands them per call
    where native MXFP8 is missing), so for a pre-quantized ``mxfp8`` load the override is only
    active while a Raon model loads; ``dequantize=True`` passed by the caller is still honoured.
    """
    try:
        from transformers.quantizers.quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer
    except ImportError:
        yield
        return
    original = FineGrainedFP8HfQuantizer.validate_environment

    def validate_environment(self, *args: Any, **kwargs: Any) -> None:
        config = self.quantization_config
        if self.pre_quantized and config.quant_method == "mxfp8" and not config.dequantize:
            from transformers.utils import is_accelerate_available

            if not is_accelerate_available():
                raise ImportError("Loading an MXFP8 model requires accelerate (`pip install accelerate`)")
            return
        original(self, *args, **kwargs)

    FineGrainedFP8HfQuantizer.validate_environment = validate_environment
    try:
        yield
    finally:
        FineGrainedFP8HfQuantizer.validate_environment = original


class Raon1MiniPreTrainedModel(PreTrainedModel):
    config: Raon1MiniConfig
    config_class = Raon1MiniConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules: ClassVar[list[str]] = ["Raon1MiniDecoderLayer"]
    _skip_keys_device_placement = "past_key_values"
    _supports_sdpa = True
    _supports_flash_attn = False
    _is_stateful = True
    # The checkpoint stores these in fp32 and the math reads them in fp32; never downcast them on load.
    _keep_in_fp32_modules: ClassVar[list[str]] = [
        "A_log",
        "dt_bias",
        "e_score_correction_bias",
        "eta",
    ]
    _keep_in_fp32_modules_strict: ClassVar[list[str]] = [
        "A_log",
        "dt_bias",
        "e_score_correction_bias",
        "eta",
    ]

    def __init__(self, config: Raon1MiniConfig, *inputs: Any, **kwargs: Any) -> None:
        super().__init__(config, *inputs, **kwargs)
        quantization = getattr(config, "quantization_config", None)
        if isinstance(quantization, dict) and quantization.get("quant_method") == "mxfp8":
            from transformers.quantizers.auto import AUTO_QUANTIZER_MAPPING

            # Without this loader the E4M3 weights would be cast without their scales.
            if "mxfp8" not in AUTO_QUANTIZER_MAPPING:
                raise ImportError(
                    "This checkpoint stores routed experts as MXFP8; use a Transformers "
                    "release with the mxfp8 quantizer (validated with 5.16.1)."
                )
        self._mxfp8_storage_checked = False

    @classmethod
    def from_pretrained(cls, *args: Any, **kwargs: Any):
        with _keep_mxfp8_storage():
            return super().from_pretrained(*args, **kwargs)

    def set_mxfp8_moe_backend(self, backend: str) -> None:
        """``"auto"``: native grouped MXFP8 GEMM where supported; ``"dequantize"``: per-call BF16."""
        if backend not in ("auto", "dequantize"):
            raise ValueError(f"backend must be 'auto' or 'dequantize', got {backend!r}.")
        for module in self.modules():
            if isinstance(module, Raon1MiniMoE):
                module.use_native_mxfp8 = backend == "auto"

    def _check_mxfp8_storage(self) -> None:
        """Reject MXFP8 storage outside the routed experts (a ``dequantize=False`` skip-list error)."""
        if self._mxfp8_storage_checked:
            return
        for name, module in self.named_modules():
            weight = getattr(module, "weight", None)
            if (
                isinstance(weight, torch.Tensor)
                and weight.dtype == torch.float8_e4m3fn
                and ".moe.experts." not in f".{name}"
            ):
                raise ValueError(
                    f"{name} was converted to FP8, but only routed experts are MXFP8 in this "
                    "checkpoint. Pass the repository's quantization_config (its ignored_layers) "
                    "together with dequantize=False."
                )
        self._mxfp8_storage_checked = True

    @torch.no_grad()
    def _init_weights(self, module: nn.Module) -> None:
        # Go through transformers.initialization: those initialisers are no-ops for parameters that
        # from_pretrained already loaded, so a plain in-place normal_() here would clobber checkpoints.
        super()._init_weights(module)
        std = self.config.initializer_range
        if isinstance(module, Raon1MiniRMSNorm):
            init.ones_(module.weight)
        elif isinstance(module, Raon1MiniShortConv):
            init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, Raon1MiniRouter):
            init.normal_(module.weight, mean=0.0, std=std)
            init.zeros_(module.e_score_correction_bias)
        elif isinstance(module, Raon1MiniRDA):
            init.zeros_(module.A_log)
            init.zeros_(module.dt_bias)
        elif isinstance(module, RaonMHCStreamMix):
            init.normal_(module.phi_z, mean=0.0, std=std)
            init.ones_(module.alpha_z)
            init.zeros_(module.b_z)
            init.zeros_(module.eta)
        elif isinstance(module, RaonMHC):
            init.normal_(module.phi_pre, mean=0.0, std=std)
            init.normal_(module.phi_post, mean=0.0, std=std)
            init.ones_(module.alpha_pre)
            init.ones_(module.alpha_post)
            init.zeros_(module.b_pre)
            init.zeros_(module.b_post)


class Raon1MiniModel(Raon1MiniPreTrainedModel):
    """Embedding -> 4-stream expansion -> decoder layers -> mean reduction -> final norm."""

    def __init__(self, config: Raon1MiniConfig) -> None:
        super().__init__(config)
        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size, config.pad_token_id
        )
        self.layers = nn.ModuleList(
            [Raon1MiniDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
        )
        self.norm = Raon1MiniRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_init()

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

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

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Raon1MiniDynamicCache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        use_cache: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **kwargs: Any,
    ) -> BaseModelOutputWithPast | tuple:
        use_cache = self.config.use_cache if use_cache is None else use_cache
        output_hidden_states = (
            self.config.output_hidden_states
            if output_hidden_states is None else output_hidden_states
        )
        return_dict = self.config.return_dict if return_dict is None else return_dict
        self._check_mxfp8_storage()
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Pass exactly one of input_ids or inputs_embeds.")
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        if use_cache and past_key_values is None:
            past_key_values = Raon1MiniDynamicCache(self.config)
        causal_mask = create_causal_mask(
            config=self.config,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            position_ids=position_ids,
        )
        token_mask = None
        if (
            attention_mask is not None
            and attention_mask.dim() == 2
            and not bool(attention_mask.all())
        ):
            token_mask = attention_mask[:, -inputs_embeds.shape[1]:]

        n = self.config.raon_mhc_streams
        stream = inputs_embeds.unsqueeze(2).expand(-1, -1, n, -1).contiguous()
        all_hidden_states = () if output_hidden_states else None
        for layer in self.layers:
            if output_hidden_states:
                # Expose [batch, tokens, hidden] representations by reducing the Raon_mHC streams,
                # as at the model output. Intermediate states are not final-normalized.
                all_hidden_states += (stream.float().mean(2).to(inputs_embeds.dtype),)
            stream = layer(
                stream,
                attention_mask=causal_mask,
                token_mask=token_mask,
                cache=past_key_values,
            )
        hidden = self.norm(stream.float().mean(2).to(inputs_embeds.dtype))
        if output_hidden_states:
            all_hidden_states += (hidden,)
        output = BaseModelOutputWithPast(
            last_hidden_state=hidden,
            past_key_values=past_key_values if use_cache else None,
            hidden_states=all_hidden_states,
        )
        return output if return_dict else output.to_tuple()


class Raon1MiniPresencePenaltyLogitsProcessor:
    """Subtract ``penalty`` from the logits of every token generated so far.

    Matches vLLM's presence penalty: prompt tokens are not penalized and a token
    is penalized once however often it appeared. ``prompt_length`` is the padded
    prompt width, the same for every row of a batch.
    """

    def __init__(self, penalty: float, prompt_length: int) -> None:
        self.penalty = penalty
        self.prompt_length = prompt_length

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        generated = input_ids[:, self.prompt_length:]
        if generated.shape[1] == 0:
            return scores
        seen = torch.zeros_like(scores, dtype=torch.bool).scatter_(1, generated, True)
        return scores - self.penalty * seen.to(scores.dtype)


class Raon1MiniForCausalLM(Raon1MiniPreTrainedModel, GenerationMixin):
    """raon-1-mini causal language model with the untied ``lm_head``."""

    _tied_weights_keys = None

    @classmethod
    def _supports_default_dynamic_cache(cls) -> bool:
        # RDA conv windows and recurrent state require the Raon hybrid cache.
        return False

    def generate(self, *args, presence_penalty: float | None = None, **kwargs):
        # GenerationConfig rejects keys it does not define, so carry this one in ourselves.
        if presence_penalty is not None:
            config = copy.deepcopy(kwargs.get("generation_config") or self.generation_config)
            config.presence_penalty = presence_penalty
            kwargs["generation_config"] = config
        return super().generate(*args, **kwargs)

    def _get_logits_processor(self, generation_config, input_ids_seq_length=None, *args, **kwargs):
        # Transformers has no presence penalty; generation_config.json sets one (1.5)
        # as the serving default, so apply it ahead of temperature/top-k/top-p as vLLM does.
        processors = super()._get_logits_processor(generation_config, input_ids_seq_length, *args, **kwargs)
        penalty = getattr(generation_config, "presence_penalty", None)
        if penalty:
            processors.insert(0, Raon1MiniPresencePenaltyLogitsProcessor(float(penalty), input_ids_seq_length))
        return processors

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

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

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

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

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

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Raon1MiniDynamicCache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        labels: torch.LongTensor | None = None,
        use_cache: bool | None = None,
        logits_to_keep: int | torch.Tensor = 0,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **kwargs: Any,
    ) -> CausalLMOutputWithPast | tuple:
        return_dict = self.config.return_dict if return_dict is None else return_dict
        self._check_mxfp8_storage()
        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_hidden_states=output_hidden_states,
            return_dict=True,
            **kwargs,
        )
        hidden = outputs.last_hidden_state
        slice_indices = (
            slice(-logits_to_keep, None)
            if isinstance(logits_to_keep, int)
            else logits_to_keep
        )
        logits = self.lm_head(hidden[:, slice_indices, :])
        loss = None
        if labels is not None:
            loss = self.loss_function(
                logits=logits,
                labels=labels,
                vocab_size=self.config.vocab_size,
                **kwargs,
            )
        output = CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
        )
        return output if return_dict else output.to_tuple()


__all__ = [
    "Raon1MiniDecoderLayer",
    "Raon1MiniDynamicCache",
    "Raon1MiniForCausalLM",
    "Raon1MiniRDA",
    "RaonMHC",
    "Raon1MiniMLA",
    "Raon1MiniMoE",
    "Raon1MiniModel",
    "Raon1MiniPreTrainedModel",
    "torch_recurrent_rda",
]