# Copyright 2026 The Agnes Model Team. All rights reserved.
#
# This modeling code is an independent implementation of the Agnes model
# that work's Apache-2.0 licensing and attribution accordingly.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from collections.abc import Callable
from typing import Optional

import torch
import torch.nn.functional as F
from torch import nn

from transformers import initialization as init
from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache, DynamicSlidingWindowLayer
from transformers.generation import GenerationMixin
from transformers.integrations import use_experts_implementation
from transformers.masking_utils import create_sliding_window_causal_mask
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
from transformers.processing_utils import Unpack
from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
from transformers.utils.generic import maybe_autocast, merge_with_config_defaults
from transformers.utils.output_capturing import OutputRecorder, capture_outputs
from transformers.conversion_mapping import register_checkpoint_conversion_mapping
from transformers.core_model_loading import (
    Concatenate,
    MergeModulelist,
    WeightConverter,
    WeightRenaming,
)

from .configuration_agnes import AgnesConfig


# ==========================================================================
# Normalization primitives
# ==========================================================================

# ---------------------------------------------------------------------------
# Root-mean-square normalisation.
#
# Agnes fuses the RMS statistic, the reciprocal-sqrt rescale and (optionally)
# the learned gain into a single explicit autograd node rather than letting the
# framework trace the gradient through the individual pointwise ops. The
# statistic is always accumulated in float32 for stability; the closed-form
# backward below is the analytic Jacobian of that statistic, so the saved
# tensors are just the fp32 input, the per-row reciprocal std and (when
# present) the gain — no intermediate activation graph is retained.
#
# For a row x (length n), with r = 1 / sqrt(mean(x^2) + eps) and x_hat = x * r:
#     y            = gain * x_hat
#     dx           = r * (g - x_hat * mean(g * x_hat))     where g = dy * gain
#     d(gain)      = sum_over_rows(dy * x_hat)
# ---------------------------------------------------------------------------


class _FusedRMSNorm(torch.autograd.Function):
    """Weighted RMS norm with a hand-derived backward (see module comment)."""

    @staticmethod
    def forward(ctx, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
        out_dtype = x.dtype
        xf = x.float()
        rstd = torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps)
        x_hat = (xf * rstd).to(out_dtype)
        ctx.save_for_backward(xf, rstd, weight)
        ctx.out_dtype = out_dtype
        return weight * x_hat

    @staticmethod
    def backward(ctx, grad_out: torch.Tensor):
        xf, rstd, weight = ctx.saved_tensors
        n = xf.shape[-1]
        x_hat = xf * rstd
        g = grad_out.float() * weight.float()
        # projection of g onto the manifold orthogonal to x_hat, then rescale by r
        row_corr = (g * x_hat).sum(-1, keepdim=True) / n
        grad_x = (rstd * (g - x_hat * row_corr)).to(ctx.out_dtype)
        reduce_axes = tuple(range(grad_out.dim() - 1))
        grad_weight = (grad_out.float() * x_hat).sum(reduce_axes).to(weight.dtype)
        return grad_x, grad_weight, None


class _FusedUnweightedRMSNorm(torch.autograd.Function):
    """Gain-free RMS norm; backward is the weighted case with gain fixed to 1."""

    @staticmethod
    def forward(ctx, x: torch.Tensor, eps: float) -> torch.Tensor:
        rstd = torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps).to(x.dtype)
        ctx.save_for_backward(x, rstd)
        return x * rstd

    @staticmethod
    def backward(ctx, grad_out: torch.Tensor):
        x, rstd = ctx.saved_tensors
        rf = rstd.float()
        x_hat = x.float() * rf
        n = x.shape[-1]
        g = grad_out.float()
        row_corr = (g * x_hat).sum(-1, keepdim=True) / n
        grad_x = (rf * (g - x_hat * row_corr)).to(grad_out.dtype)
        return grad_x, None


class AgnesRMSNorm(nn.Module):
    """RMS norm with a learned per-channel gain."""

    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return _FusedRMSNorm.apply(hidden_states, self.weight, self.variance_epsilon)

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


class AgnesUnweightedRMSNorm(nn.Module):
    """RMS norm without a learned gain (used inside the attention / mHC blocks)."""

    def __init__(self, eps: float = 1.0e-6):
        super().__init__()
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _FusedUnweightedRMSNorm.apply(x, self.eps)



# ==========================================================================
# Feed-forward network and mixture-of-experts routing
# ==========================================================================

class AgnesMLP(nn.Module):
    def __init__(self, config: AgnesConfig, intermediate_size: int | None = None):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
        self.act_fn = ACT2FN[config.hidden_act]
        self.limit = config.swiglu_limit

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SwiGLU with a symmetric pre-activation clip: the gate branch is capped
        # from above, the value branch on both sides, before they are multiplied
        # and projected back down.
        activated = self.act_fn(self.gate_proj(x).clamp(max=self.limit))
        value = self.up_proj(x).clamp(min=-self.limit, max=self.limit)
        return self.down_proj(activated * value)


@use_experts_implementation
class AgnesExperts(nn.Module):
    """Routed experts held as stacked 3-D weight tensors (one slab per expert)."""

    def __init__(self, config: AgnesConfig):
        super().__init__()
        self.num_experts = config.num_local_experts
        self.hidden_dim = config.hidden_size
        self.intermediate_dim = config.intermediate_size
        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
        self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
        self.act_fn = ACT2FN[config.hidden_act]
        self.limit = config.swiglu_limit

    def forward(
        self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor
    ) -> torch.Tensor:
        # Reference (unfused) dispatch: walk only the experts that at least one
        # token selected, gather that expert's tokens, run its SwiGLU, and
        # scatter-add the weighted result back. The heavy fused kernels are
        # swapped in by `@use_experts_implementation`; this body is the fallback.
        out = torch.zeros_like(hidden_states)
        with torch.no_grad():
            # [experts, slot, token] boolean routing tensor
            routing = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0)
            used = (routing.sum(dim=(-1, -2)) > 0).nonzero().flatten()
        for e in used.tolist():
            if e == self.num_experts:
                continue
            slot, tok = torch.where(routing[e])
            hidden = self._apply_gate(F.linear(hidden_states[tok], self.gate_up_proj[e]))
            hidden = F.linear(hidden, self.down_proj[e]) * top_k_weights[tok, slot, None]
            out.index_add_(0, tok, hidden.to(out.dtype))
        return out

    def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor:
        # Method (not inlined) so the fused grouped_mm / batched_mm expert backends
        # from `@use_experts_implementation` apply the identical clip + SiLU on top
        # of their packed gate/up output instead of bypassing it.
        gate, value = gate_up.chunk(2, dim=-1)
        gated = self.act_fn(gate.clamp(max=self.limit))
        return gated * value.clamp(min=-self.limit, max=self.limit)


class AgnesTopKRouter(nn.Module):
    def __init__(self, config: AgnesConfig):
        super().__init__()
        self.top_k = config.num_experts_per_tok
        self.num_experts = config.num_local_experts
        self.hidden_dim = config.hidden_size
        self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim))
        self.score_fn = ACT2FN[config.scoring_func]
        self.routed_scaling_factor = config.routed_scaling_factor
        self.register_buffer("e_score_correction_bias", torch.zeros(self.num_experts), persistent=True)

    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        tokens = hidden_states.reshape(-1, self.hidden_dim)
        router_logits = F.linear(tokens, self.weight)
        affinity = self.score_fn(router_logits)
        # bias only steers the top-k *selection*; the gathered gate values are the
        # un-biased affinities, renormalised to sum to one over the chosen experts.
        chosen = torch.topk(affinity + self.e_score_correction_bias, self.top_k, dim=-1, sorted=False).indices
        gate = affinity.gather(1, chosen)
        gate = gate / (gate.sum(dim=-1, keepdim=True) + 1e-20)
        return router_logits, gate * self.routed_scaling_factor, chosen


class AgnesHashRouter(nn.Module):
    r"""
    Static hash routing used by the leading `agnes_hash_moe` MoE layers. Which
    experts a token visits is fixed up front by a frozen `tid2eid[input_ids]`
    lookup table (token id -> expert ids) rather than by a learned arg-top-k. The
    learned `weight` is still evaluated to produce the affinities that *weight*
    those experts' outputs; only the selection itself is frozen.
    """

    def __init__(self, config: AgnesConfig):
        super().__init__()
        self.top_k = config.num_experts_per_tok
        self.num_experts = config.num_local_experts
        self.hidden_dim = config.hidden_size
        self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim))
        self.score_fn = ACT2FN[config.scoring_func]
        self.routed_scaling_factor = config.routed_scaling_factor
        self.register_buffer("tid2eid", torch.zeros(config.vocab_size, self.top_k, dtype=torch.long), persistent=True)

    def forward(
        self, hidden_states: torch.Tensor, input_ids: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        tokens = hidden_states.reshape(-1, self.hidden_dim)
        router_logits = F.linear(tokens, self.weight)
        affinity = self.score_fn(router_logits)
        chosen = self.tid2eid[input_ids.reshape(-1)].long()  # frozen per-token expert ids
        gate = affinity.gather(1, chosen)
        gate = gate / (gate.sum(dim=-1, keepdim=True) + 1e-20)
        return router_logits, gate * self.routed_scaling_factor, chosen


class AgnesSparseMoeBlock(nn.Module):
    def __init__(self, config: AgnesConfig, layer_idx: int):
        super().__init__()
        self.is_hash = config.mlp_layer_types[layer_idx] == "agnes_hash_moe"
        self.gate = AgnesHashRouter(config) if self.is_hash else AgnesTopKRouter(config)
        self.experts = AgnesExperts(config)
        self.shared_experts = AgnesMLP(config)
        pffn_size = getattr(config, "parallel_ffn_intermediate_size", 0) or 0
        # parallel dense FFN branch; down_proj is zero-initialized at export so the
        # block is exactly identity-preserving until trained
        self.parallel_ffn = (
            AgnesMLP(config, intermediate_size=pffn_size) if pffn_size and not self.is_hash else None
        )

    def forward(self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None) -> torch.Tensor:
        b, s, d = hidden_states.shape
        tokens = hidden_states.reshape(-1, d)
        # hash layers additionally consume the raw token ids for the frozen lookup
        if self.is_hash:
            _, gate_w, gate_idx = self.gate(hidden_states, input_ids)
        else:
            _, gate_w, gate_idx = self.gate(hidden_states)
        result = self.experts(tokens, gate_idx, gate_w).view(b, s, d)
        result = result + self.shared_experts(hidden_states)
        if self.parallel_ffn is not None:  # zero-init identity branch until trained
            result = result + self.parallel_ffn(hidden_states)
        return result



# ==========================================================================
# Residual hyper-connections (mHC)
# ==========================================================================

class AgnesHyperConnection(nn.Module):
    r"""
    Manifold-constrained hyper-connection (mHC). Where a plain block does a
    single scalar residual add, Agnes keeps `hc_mult` parallel residual streams
    and learns how to (i) collapse them into the sublayer input, (ii) place the
    sublayer output back across them, and (iii) re-mix the streams themselves.

    A single small projection (`fn`, biased by `base`, per-branch `scale`) maps
    the `hc_mult` streams — after an unweighted RMS norm and a flatten — into
    three groups of logits:

    * `pre`  (length `hc_mult`)         : sigmoid gate that collapses the streams
                                          into one sequence for the sublayer.
    * `post` (length `hc_mult`)         : `2·sigmoid` gate in [0, 2] that spreads
                                          the sublayer output back over streams.
    * `comb` (`hc_mult × hc_mult`)      : softmaxed then Sinkhorn-projected onto
                                          the doubly-stochastic manifold; mixes
                                          the streams across the residual.

    The decoder layer holds two of these, one at the attention site and one at
    the MLP site.
    """

    def __init__(self, config: AgnesConfig):
        super().__init__()
        self.hc_mult = config.hc_mult
        self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
        self.hc_eps = config.hc_eps
        self.input_norm = AgnesUnweightedRMSNorm(eps=config.rms_norm_eps)
        n_logits = (2 + self.hc_mult) * self.hc_mult
        self.fn = nn.Parameter(torch.empty(n_logits, self.hc_mult * config.hidden_size))
        self.base = nn.Parameter(torch.empty(n_logits))
        # one learned scale each for the pre / post / comb logit groups
        self.scale = nn.Parameter(torch.empty(3))

    def forward(self, hidden_streams: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        r"""Return `(post, comb, collapsed)`: the output-placement gate, the
        doubly-stochastic stream-mixing matrix, and the collapsed sublayer input.
        `comb` is driven onto the doubly-stochastic manifold by Sinkhorn-Knopp —
        an initial column normalisation followed by `hc_sinkhorn_iters - 1`
        alternating row/column passes."""
        n = self.hc_mult
        normed = self.input_norm(hidden_streams.flatten(start_dim=2).float())
        pre_raw, post_raw, comb_raw = F.linear(normed, self.fn.float()).split([n, n, n * n], dim=-1)
        pre_bias, post_bias, comb_bias = self.base.split([n, n, n * n])
        pre_scale, post_scale, comb_scale = self.scale.unbind(0)

        pre = torch.sigmoid(pre_raw * pre_scale + pre_bias) + self.hc_eps
        post = 2 * torch.sigmoid(post_raw * post_scale + post_bias)

        comb = comb_raw.view(*comb_raw.shape[:-1], n, n) * comb_scale + comb_bias.view(n, n)
        comb = torch.softmax(comb, dim=-1) + self.hc_eps
        comb = comb / (comb.sum(dim=-2, keepdim=True) + self.hc_eps)
        for _ in range(self.hc_sinkhorn_iters - 1):
            comb = comb / (comb.sum(dim=-1, keepdim=True) + self.hc_eps)
            comb = comb / (comb.sum(dim=-2, keepdim=True) + self.hc_eps)

        # weighted sum over the stream axis -> single sequence fed to the sublayer
        collapsed = (pre.unsqueeze(-1) * hidden_streams).sum(dim=2).to(hidden_streams.dtype)
        return post, comb, collapsed


class AgnesHyperHead(nn.Module):
    """Collapses the `hc_mult` residual streams back into one tensor at the top
    of the stack, just before the final shared RMS norm in `AgnesModel`."""

    def __init__(self, config: AgnesConfig):
        super().__init__()
        self.hc_mult = config.hc_mult
        self.input_norm = AgnesUnweightedRMSNorm(eps=config.rms_norm_eps)
        self.eps = config.hc_eps
        self.hc_fn = nn.Parameter(torch.empty(self.hc_mult, self.hc_mult * config.hidden_size))
        self.hc_base = nn.Parameter(torch.empty(self.hc_mult))
        self.hc_scale = nn.Parameter(torch.empty(1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # same gate-and-sum as the `pre` branch of AgnesHyperConnection, but with
        # no sublayer to feed — this is the terminal stream collapse.
        gate_logits = F.linear(self.input_norm(x.flatten(2).float()), self.hc_fn.float())
        gate = torch.sigmoid(gate_logits * self.hc_scale.float() + self.hc_base.float()) + self.eps
        return (gate.unsqueeze(-1) * x).sum(dim=2).to(x.dtype)



# ==========================================================================
# Rotary position embedding
# ==========================================================================

class AgnesRotaryEmbedding(nn.Module):
    """
    Rotary embedding that maintains one inverse-frequency buffer per *rope
    label* and applies interleaved (rather than half-split) rotation with
    partial coverage of each head.

    Two things are specific to Agnes here:

    * Interleaved layout — a single `θ_i` is shared by the two channels of
      each consecutive pair, so only `rope_head_dim // 2` frequencies are
      stored. The doubling back up to the full rope width happens next to the
      rotation itself (in `apply_agnes_rope`) instead of being baked into the
      cached cos/sin, which keeps the pairing explicit at the point of use.

    * Rope labels are independent of the attention schedule. The architectural
      `layer_types` (`agnes_local_attention` / `agnes_sparse_attention` /
      `agnes_pooled_attention`) are decoupled from the rope labels
      (`main` / `compress`), which live as keys in `config.rope_parameters`
      and differ only in their `rope_theta` base. Buffers are therefore built
      by iterating `rope_parameters` rather than the attention schedule.
    """

    inv_freq: torch.Tensor  # fix linting for `register_buffer`

    def __init__(self, config: AgnesConfig):
        super().__init__()
        self.max_seq_len_cached = config.max_position_embeddings
        self.original_max_seq_len = config.max_position_embeddings
        self.config = config
        # Only the nested per-rope-type sub-dicts are real layer types — the top-level
        # `rope_type` key that ``convert_rope_params_to_dict`` may leave on
        # ``config.rope_parameters`` is a flat-shape leftover, not a layer.
        self.layer_types = [k for k, v in config.rope_parameters.items() if isinstance(v, dict)]
        self.rope_type = {}
        for layer_type in self.layer_types:
            rope_params = config.rope_parameters[layer_type]
            self.rope_type[layer_type] = rope_params["rope_type"]
            rope_init_fn = self.compute_default_rope_parameters
            if self.rope_type[layer_type] != "default":
                rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type[layer_type]]
            inv_freq, attention_scaling = rope_init_fn(config, layer_type=layer_type)
            self.register_buffer(f"{layer_type}_inv_freq", inv_freq, persistent=False)
            self.register_buffer(f"{layer_type}_original_inv_freq", inv_freq.clone(), persistent=False)
            setattr(self, f"{layer_type}_attention_scaling", attention_scaling)

    @staticmethod
    def compute_default_rope_parameters(
        config: AgnesConfig | None = None,
        device: Optional["torch.device"] = None,
        seq_len: int | None = None,
        layer_type: str | None = None,
    ) -> tuple["torch.Tensor", float]:
        """Plain (non-scaled) inverse frequencies for one rope label.

        Returns `(inv_freq, attention_factor)`. `attention_factor` is always 1.0
        for this default variant (the scaled variants such as YaRN come from
        `ROPE_INIT_FUNCTIONS` instead). Only the leading `partial_rotary_factor`
        fraction of each head receives rotation, so the frequency table has
        `rope_dim // 2` entries. `seq_len` is accepted for signature parity and
        unused here.
        """
        rope_cfg = config.rope_parameters[layer_type]
        base = rope_cfg["rope_theta"]
        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
        rope_dim = int(head_dim * rope_cfg.get("partial_rotary_factor", 1.0))

        exponents = torch.arange(0, rope_dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / rope_dim
        inv_freq = 1.0 / (base**exponents)
        return inv_freq, 1.0

    @torch.no_grad()
    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
    def forward(self, x, position_ids, layer_type=None):
        # cos/sin carry one entry per interleaved pair; the widening to full rope
        # width is left to `apply_agnes_rope` so the pairing stays where it is used.
        inv_freq = getattr(self, f"{layer_type}_inv_freq")
        scaling = getattr(self, f"{layer_type}_attention_scaling")
        # outer product positions (B, S) x frequencies (rope_dim/2) -> (B, S, rope_dim/2)
        freq_col = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
        pos_row = position_ids[:, None, :].float()
        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
        with maybe_autocast(device_type=device_type, enabled=False):
            angles = (freq_col.float() @ pos_row.float()).transpose(1, 2)
            cos = angles.cos() * scaling
            sin = angles.sin() * scaling
        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)


def _swap_pairs(x: torch.Tensor) -> torch.Tensor:
    """Interleaved-pair rotation partner: within every adjacent `(a, b)` pair
    return `(-b, a)`. Equivalent to the `[-x_odd, x_even]` interleave used by
    the rotation formula below."""
    pairs = x.unflatten(-1, (-1, 2))
    a, b = pairs[..., 0], pairs[..., 1]
    return torch.stack((-b, a), dim=-1).flatten(-2)


def apply_agnes_rope(
    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, unsqueeze_dim: int = 1
) -> torch.Tensor:
    """Apply interleaved rotary position embedding to the trailing rope slice.

    The cached `cos` / `sin` carry one entry per interleaved pair. They are
    widened to the full rope dimension right here (`repeat_interleave(2)`), the
    leading "no-position" channels are passed through untouched, and only the
    final `2 * cos.shape[-1]` channels are rotated via
    `x * cos + swap(x) * sin`. The rotation math is done in float32 and cast
    back to `x`'s dtype. Each head is laid out as `[nope | rope]`.
    """
    cos = cos.repeat_interleave(2, dim=-1).unsqueeze(unsqueeze_dim)
    sin = sin.repeat_interleave(2, dim=-1).unsqueeze(unsqueeze_dim)
    rope_dim = cos.shape[-1]
    split_at = x.shape[-1] - rope_dim
    nope, rope = x[..., :split_at], x[..., split_at:]
    rotated = (rope.float() * cos + _swap_pairs(rope).float() * sin).to(x.dtype)
    return torch.cat([nope, rotated], dim=-1)



# ==========================================================================
# Attention KV caches
# ==========================================================================

class AgnesHCACache(DynamicSlidingWindowLayer):
    r"""Cache layer for HCA blocks. Holds the long-range compressor's
    buffer / running compressed entries / count on top of the sliding-window K=V
    branch. HCA uses *non-overlapping* windows, so there is *no* overlap state,
    and HCA has *no* indexer either.

    State is dict-keyed by entry name — HCA only uses `"compressor"`, but
    :class:`AgnesCSACache` adds `"indexer"` to the same dicts so a single
    set of methods (`store_compression_weights` / `update_compressor_states`)
    serves both:

      * `compressed_kv[name]` — the running list of compressed KV entries
        emitted so far (one every `compress_rate` source tokens; the long-range
        KVs the attention concatenates onto its sliding-window keys / values).
      * `buffer_kv[name]` / `buffer_gate[name]` — source tokens that arrived
        between two full windows; once the buffer hits `compress_rate` tokens
        the compressor closes a window, emits one entry, and drains the buffer.
      * `entry_count[name]` — number of compressed entries emitted so far, so
        `entry_count[name] * compress_rate` is the absolute position of the
        *next* window's first source token. Tracked separately from
        `position_ids` so prefill -> decode -> prefill stays consistent.
    """

    layer_type = "agnes_pooled_attention"

    def __init__(self, config: "AgnesConfig"):
        super().__init__(config)
        self.compress_rate = config.compress_rates["agnes_pooled_attention"]
        self.buffer_kv: dict[str, torch.Tensor | None] = {"compressor": None}
        self.buffer_gate: dict[str, torch.Tensor | None] = {"compressor": None}
        self.compressed_kv: dict[str, torch.Tensor | None] = {"compressor": None}
        self.entry_count: dict[str, int] = {"compressor": 0}

    def update(self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs):
        """Sliding-window key/value update. Because Agnes is shared-KV MQA, keys
        and values are one and the same buffer, so both returns are the full
        (pre-trim) concatenation while the retained cache keeps only the last
        `sliding_window - 1` positions."""
        if not self.is_initialized:
            self.lazy_initialization(key_states, value_states)
            self.values = self.keys
        self.cumulative_length += key_states.shape[-2]
        combined = torch.cat([self.keys, key_states], dim=-2)
        self.keys = combined[:, :, -self.sliding_window + 1 :, :]
        self.values = self.keys
        return combined, combined

    def store_compression_weights(
        self, name: str, kv: torch.Tensor, gate: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, int]:
        r"""Prepend the leftover buffered `(kv, gate)` for entry `name`, split
        off the longest window-aligned prefix (the part that can be compressed
        now), stash the remainder back in the buffer, and hand back
        `(chunk_kv, chunk_gate, first_window_position)`. The compressor then
        softmax-pools each `compress_rate`-token window of that chunk (with its
        `position_bias`) into one compressed entry.
        """
        first_window_position = self.entry_count[name] * self.compress_rate
        held_kv, held_gate = self.buffer_kv[name], self.buffer_gate[name]
        if held_kv is not None and held_kv.shape[1]:
            kv = torch.cat([held_kv, kv], dim=1)
            gate = torch.cat([held_gate, gate], dim=1)
        # split at the largest multiple of compress_rate; carry the tail forward
        cut = (kv.shape[1] // self.compress_rate) * self.compress_rate
        self.buffer_kv[name], self.buffer_gate[name] = kv[:, cut:], gate[:, cut:]
        return kv[:, :cut], gate[:, :cut], first_window_position

    def update_compressor_states(self, name: str, compressed: torch.Tensor) -> torch.Tensor:
        r"""Append the newly emitted compressed entries to the running
        `compressed_kv[name]`, advance `entry_count[name]`, and return the
        accumulated tensor."""
        running = self.compressed_kv[name]
        if running is None:
            self.compressed_kv[name] = compressed
        elif compressed.shape[1] > 0:
            self.compressed_kv[name] = torch.cat([running, compressed], dim=1)
        self.entry_count[name] += compressed.shape[1]
        return self.compressed_kv[name]


class AgnesCSACache(AgnesHCACache):
    r"""CSA cache. On top of :class:`AgnesHCACache` it registers a second entry
    name, `"indexer"`, in each of the inherited state dicts, and it keeps a small
    per-name *overlap* buffer required by the two-series windowing.

    Why the overlap buffer exists: for CSA, `kv_proj` / `gate_proj` emit
    `2 * head_dim` features per token — two interleaved series, Ca in
    `[..., :head_dim]` and Cb in `[..., head_dim:]`. The pooled entry for window
    `w` blends window `w-1`'s Ca with window `w`'s Cb (softmax-gated), giving an
    effective receptive width of `2 * compress_rate_csa` at stride
    `compress_rate_csa`. The only cross-window dependency is therefore the
    previous window's Ca slice, so at a forward boundary we persist exactly
    `chunk[:, -1, :, :head_dim]` (Ca of the last full window) in
    `overlap_kv[name]` / `overlap_gate[name]`; Cb is never revisited.
    """

    layer_type = "agnes_sparse_attention"

    def __init__(self, config: "AgnesConfig"):
        super().__init__(config)
        self.compress_rate = config.compress_rates["agnes_sparse_attention"]
        self.buffer_kv["indexer"] = None
        self.buffer_gate["indexer"] = None
        self.compressed_kv["indexer"] = None
        self.entry_count["indexer"] = 0
        self.overlap_kv: dict[str, torch.Tensor | None] = {"compressor": None, "indexer": None}
        self.overlap_gate: dict[str, torch.Tensor | None] = {"compressor": None, "indexer": None}

    def update_overlap_state(
        self, name: str, chunk_kv: torch.Tensor, chunk_gate: torch.Tensor, head_dim: int
    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
        r"""Swap the overlap buffer: return the Ca slice saved on the previous
        call (or `None` on the first) and store this call's last-window Ca slice
        for the next one. Only `[..., :head_dim]` (Ca) is kept, since Cb is
        already baked into an emitted entry and never read again.
        """
        carried_kv, carried_gate = self.overlap_kv[name], self.overlap_gate[name]
        self.overlap_kv[name] = chunk_kv[:, -1, :, :head_dim].clone()
        self.overlap_gate[name] = chunk_gate[:, -1, :, :head_dim].clone()
        return carried_kv, carried_gate



# ==========================================================================
# Long-range compressors and lightning indexer
# ==========================================================================

class AgnesHCACompressor(nn.Module):
    """
    Heavily Compressed Attention compressor. Compresses
    every `compress_rate_hca` (m'=128) source tokens into a single compressed KV
    entry.

    Each closed window of m' tokens produces one compressed entry:
    `C^{Comp}_i = Σ_{j∈window} softmax(Z_j + B)_j ⊙ C_j`. RoPE on the trailing
    `rope_head_dim` slice is applied at the deterministic absolute position
    `i * compress_rate_hca + first_window_position` so cross-call concatenation
    stays causality-correct. Returns the running list of *all* compressed
    entries emitted so far (shape `[B, 1, T, head_dim]` with
    `T = entry_count["compressor"]`), so the attention can attend over the
    full long-range history.

    When `past_key_values is None` runs in stateless single-shot mode: compress
    every complete window from `hidden_states` and discard the remainder
    (instead of caching it).
    """

    rope_layer_type = "compress"

    def __init__(self, config: AgnesConfig):
        super().__init__()
        self.compress_rate = config.compress_rates["agnes_pooled_attention"]
        self.head_dim = config.head_dim
        self.kv_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False)
        self.gate_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False)
        self.position_bias = nn.Parameter(torch.empty(self.compress_rate, self.head_dim))
        self.kv_norm = AgnesRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.rotary_emb = AgnesRotaryEmbedding(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        q_residual: torch.Tensor,
        position_ids: torch.Tensor,
        past_key_values: Cache | None,
        layer_idx: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        batch = hidden_states.shape[0]
        cache_layer: AgnesHCACache = past_key_values.layers[layer_idx] if past_key_values is not None else None
        kv = self.kv_proj(hidden_states)
        gate = self.gate_proj(hidden_states)
        # stateless mode drops the ragged tail; cached mode buffers it for later
        if cache_layer is None:
            aligned = (kv.shape[1] // self.compress_rate) * self.compress_rate
            chunk_kv, chunk_gate, first_window_position = kv[:, :aligned], gate[:, :aligned], 0
        else:
            chunk_kv, chunk_gate, first_window_position = cache_layer.store_compression_weights("compressor", kv, gate)

        if chunk_kv.shape[1] > 0:  # at least one full window is ready
            n_windows = chunk_kv.shape[1] // self.compress_rate
            chunk_kv = chunk_kv.view(batch, n_windows, self.compress_rate, -1)
            biased_gate = chunk_gate.view(batch, n_windows, self.compress_rate, -1) + self.position_bias
            # softmax pooling within each window (fp32 for stability), then norm
            pool_w = biased_gate.softmax(dim=2, dtype=torch.float32).to(chunk_kv.dtype)
            compressed = self.kv_norm((chunk_kv * pool_w).sum(dim=2))
            abs_pos = torch.arange(n_windows, device=compressed.device) * self.compress_rate + first_window_position
            abs_pos = abs_pos.unsqueeze(0).expand(batch, -1)
            cos, sin = self.rotary_emb(compressed, position_ids=abs_pos, layer_type=self.rope_layer_type)
            compressed = apply_agnes_rope(compressed.unsqueeze(1), cos, sin).squeeze(1)
        else:
            compressed = chunk_kv.new_zeros((batch, 0, self.head_dim))

        if cache_layer is not None:
            compressed = cache_layer.update_compressor_states("compressor", compressed)
        compressed_kv = compressed.unsqueeze(1)

        n_entries = compressed_kv.shape[2]
        seq_len = position_ids.shape[1]
        if seq_len == 1 or n_entries == 0:
            return compressed_kv, None

        # a query at position t may only see entry w once t has passed window w,
        # i.e. w < (t + 1) // compress_rate; everything else is masked to -inf.
        entry_pos = torch.arange(n_entries, device=compressed_kv.device)
        ready = (position_ids + 1) // self.compress_rate  # [B, S]
        block_bias = compressed_kv.new_zeros((batch, 1, seq_len, n_entries))
        block_bias = block_bias.masked_fill(
            entry_pos.view(1, 1, 1, -1) >= ready.unsqueeze(1).unsqueeze(-1),
            float("-inf"),
        )
        return compressed_kv, block_bias


class AgnesIndexerScorer(nn.Module):
    r"""Lightning-indexer scoring head: `∑_h w_{t,h} · ReLU(q_{t,h} · K^IComp_s)`."""

    def __init__(self, config: AgnesConfig):
        super().__init__()
        self.softmax_scale = config.index_head_dim**-0.5
        self.weights_scaling = config.index_n_heads**-0.5
        self.weights_proj = nn.Linear(config.hidden_size, config.index_n_heads, bias=False)

    def forward(self, q: torch.Tensor, compressed_kv: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
        # per-head query·key, ReLU-rectified and scaled -> [B, S, H, T]
        qk = torch.matmul(q.float(), compressed_kv.transpose(-1, -2).float().unsqueeze(1))
        qk = F.relu(qk) * self.softmax_scale
        head_w = self.weights_proj(hidden_states).float() * self.weights_scaling  # [B, S, H]
        # collapse the head axis with the learned per-head weights -> [B, S, T]
        return (qk * head_w.unsqueeze(-1)).sum(dim=2)


class AgnesIndexer(nn.Module):
    r"""Lightning indexer for CSA. For each query it keeps only the top
    `config.index_topk` of the compressed entries, cutting the effective KV set
    from `seq_len / compress_rate_csa` down to `index_topk`.

    It builds a miniature compressor of its own at `index_head_dim` over the
    same windows as the outer CSA compressor, scores each query against those
    compressed keys as `sum_h w_h * ReLU(q_h . k)`, and returns the winning
    indices.

    A dedicated rotary lives here because RoPE has to be applied twice, to two
    tensors whose positions differ: the compressed keys sit at the fixed window
    positions `i * compress_rate + first_window_position`, whereas the queries
    sit at the (per-call) `position_ids`. Both use the compressor's theta
    (`compress_rope_theta`) so that only relative offsets survive the dot
    product; since the query positions change every call, cos/sin can't be
    precomputed, so the indexer calls its rotary twice per forward (label
    always `"compress"`).
    """

    rope_layer_type = "compress"

    def __init__(self, config: AgnesConfig):
        super().__init__()
        self.compress_rate = config.compress_rates["agnes_sparse_attention"]
        self.num_heads = config.index_n_heads
        self.head_dim = config.index_head_dim
        self.index_topk = config.index_topk
        self.kv_proj = nn.Linear(config.hidden_size, 2 * self.head_dim, bias=False)
        self.gate_proj = nn.Linear(config.hidden_size, 2 * self.head_dim, bias=False)
        self.position_bias = nn.Parameter(torch.empty(self.compress_rate, 2 * self.head_dim))
        self.kv_norm = AgnesRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.head_dim, bias=False)
        self.rotary_emb = AgnesRotaryEmbedding(config)
        self.scorer = AgnesIndexerScorer(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        q_residual: torch.Tensor,
        position_ids: torch.Tensor,
        past_key_values: Cache | None,
        layer_idx: int,
    ) -> torch.LongTensor:
        batch, seq_len, _ = hidden_states.shape
        cache_layer: AgnesCSACache = past_key_values.layers[layer_idx] if past_key_values is not None else None
        kv = self.kv_proj(hidden_states)
        gate = self.gate_proj(hidden_states)

        if cache_layer is None:
            aligned = (kv.shape[1] // self.compress_rate) * self.compress_rate
            chunk_kv, chunk_gate, first_window_position = kv[:, :aligned], gate[:, :aligned], 0
        else:
            chunk_kv, chunk_gate, first_window_position = cache_layer.store_compression_weights("indexer", kv, gate)

        if chunk_kv.shape[1] > 0:
            n_windows = chunk_kv.shape[1] // self.compress_rate
            m = self.compress_rate
            chunk_kv = chunk_kv.view(batch, n_windows, m, -1)
            chunk_gate = chunk_gate.view(batch, n_windows, m, -1) + self.position_bias

            # widen each window to 2m slots: current-window Cb in the top half,
            # previous-window Ca in the bottom half (same scheme as outer CSA).
            win_kv = chunk_kv.new_zeros((batch, n_windows, 2 * m, self.head_dim))
            win_gate = chunk_gate.new_full((batch, n_windows, 2 * m, self.head_dim), float("-inf"))
            win_kv[:, :, m:] = chunk_kv[..., self.head_dim :]
            win_gate[:, :, m:] = chunk_gate[..., self.head_dim :]
            if n_windows > 1:
                win_kv[:, 1:, :m] = chunk_kv[:, :-1, :, : self.head_dim]
                win_gate[:, 1:, :m] = chunk_gate[:, :-1, :, : self.head_dim]
            if cache_layer is not None:
                carried_kv, carried_gate = cache_layer.update_overlap_state("indexer", chunk_kv, chunk_gate, self.head_dim)
                if carried_kv is not None:
                    win_kv[:, 0, :m] = carried_kv.to(win_kv.dtype)
                    win_gate[:, 0, :m] = carried_gate.to(win_gate.dtype)

            pool_w = win_gate.softmax(dim=2, dtype=torch.float32).to(win_kv.dtype)
            compressed = self.kv_norm((win_kv * pool_w).sum(dim=2))
            abs_pos = torch.arange(n_windows, device=compressed.device) * self.compress_rate + first_window_position
            abs_pos = abs_pos.unsqueeze(0).expand(batch, -1)
            cos, sin = self.rotary_emb(compressed, position_ids=abs_pos, layer_type=self.rope_layer_type)
            compressed = apply_agnes_rope(compressed.unsqueeze(1), cos, sin).squeeze(1)
        else:
            compressed = chunk_kv.new_zeros((batch, 0, self.head_dim))

        compressed_kv = (
            compressed if cache_layer is None else cache_layer.update_compressor_states("indexer", compressed)
        )

        cos_q, sin_q = self.rotary_emb(hidden_states, position_ids=position_ids, layer_type=self.rope_layer_type)
        q = self.q_b_proj(q_residual).view(batch, seq_len, -1, self.head_dim).transpose(1, 2)
        q = apply_agnes_rope(q, cos_q, sin_q).transpose(1, 2)

        scores = self.scorer(q, compressed_kv, hidden_states)  # [B, S, T]
        n_entries = compressed_kv.shape[1]
        k = min(self.index_topk, n_entries)

        # a query can only see entries whose window has fully closed before it; the
        # `topk` is taken after masking future entries to -inf, and any pick that
        # still points past the ready frontier (too few blocks yet) is flagged -1.
        if n_entries > 0:
            ready = (position_ids + 1) // self.compress_rate  # [B, S]
            entry_pos = torch.arange(n_entries, device=scores.device)
            future = entry_pos.view(1, 1, -1) >= ready.unsqueeze(-1)  # [B, S, T]
            scores = scores.masked_fill(future, float("-inf"))
            picks = scores.topk(k, dim=-1).indices  # [B, S, k]
            stale = picks >= ready.unsqueeze(-1)
            return torch.where(stale, torch.full_like(picks, -1), picks)

        return scores.topk(k, dim=-1).indices


class AgnesCSACompressor(nn.Module):
    """Compressed Sparse Attention compressor. It pools every `compress_rate_csa`
    (m=4) source tokens into one entry and pairs the result with a Lightning
    Indexer that keeps only the top `index_topk` entries per query before core
    attention runs.

    `kv_proj` / `gate_proj` / `position_bias` all emit `2 * head_dim` features,
    holding two series per token: Ca in `[..., :head_dim]` (contributes to the
    *next* window's entry) and Cb in `[..., head_dim:]` (contributes to the
    *current* one). Entry `w` softmax-blends window `w-1`'s Ca with window `w`'s
    Cb across `2 * compress_rate_csa` slots — width `2 * compress_rate_csa`,
    stride `compress_rate_csa`. Window 0 needs the previous forward's last Ca
    slice, which the cache returns via `overlap_kv`; with no cache (or on the
    first call) that half is left as zero-kv / `-inf`-gate so it contributes
    nothing to the softmax.
    """

    rope_layer_type = "compress"

    def __init__(self, config: AgnesConfig):
        super().__init__()
        self.compress_rate = config.compress_rates["agnes_sparse_attention"]
        self.head_dim = config.head_dim
        self.kv_proj = nn.Linear(config.hidden_size, 2 * self.head_dim, bias=False)
        self.gate_proj = nn.Linear(config.hidden_size, 2 * self.head_dim, bias=False)
        self.position_bias = nn.Parameter(torch.empty(self.compress_rate, 2 * self.head_dim))
        self.kv_norm = AgnesRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.rotary_emb = AgnesRotaryEmbedding(config)
        self.indexer = AgnesIndexer(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        q_residual: torch.Tensor,
        position_ids: torch.Tensor,
        past_key_values: Cache | None,
        layer_idx: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        batch, seq_len, _ = hidden_states.shape
        cache_layer: AgnesCSACache = past_key_values.layers[layer_idx] if past_key_values is not None else None
        kv = self.kv_proj(hidden_states)
        gate = self.gate_proj(hidden_states)

        if cache_layer is None:
            aligned = (kv.shape[1] // self.compress_rate) * self.compress_rate
            chunk_kv, chunk_gate, first_window_position = kv[:, :aligned], gate[:, :aligned], 0
        else:
            chunk_kv, chunk_gate, first_window_position = cache_layer.store_compression_weights("compressor", kv, gate)

        if chunk_kv.shape[1] > 0:
            n_windows = chunk_kv.shape[1] // self.compress_rate
            m = self.compress_rate
            chunk_kv = chunk_kv.view(batch, n_windows, m, -1)
            chunk_gate = chunk_gate.view(batch, n_windows, m, -1) + self.position_bias

            # 2m-wide slot layout: top half = current window's Cb, bottom half =
            # previous window's Ca. For window 0 the bottom half is either the Ca
            # slice returned by the cache, or (no cache / first call) left as
            # zero-kv + -inf-gate so it drops out of the softmax.
            win_kv = chunk_kv.new_zeros((batch, n_windows, 2 * m, self.head_dim))
            win_gate = chunk_gate.new_full((batch, n_windows, 2 * m, self.head_dim), float("-inf"))
            win_kv[:, :, m:] = chunk_kv[..., self.head_dim :]
            win_gate[:, :, m:] = chunk_gate[..., self.head_dim :]
            if n_windows > 1:
                win_kv[:, 1:, :m] = chunk_kv[:, :-1, :, : self.head_dim]
                win_gate[:, 1:, :m] = chunk_gate[:, :-1, :, : self.head_dim]
            if cache_layer is not None:
                carried_kv, carried_gate = cache_layer.update_overlap_state(
                    "compressor", chunk_kv, chunk_gate, self.head_dim
                )
                if carried_kv is not None:
                    win_kv[:, 0, :m] = carried_kv.to(win_kv.dtype)
                    win_gate[:, 0, :m] = carried_gate.to(win_gate.dtype)

            # fp32 softmax pooling: in bf16/fp16 close logits over wide windows
            # can collapse together, so accumulate the weights in float.
            pool_w = win_gate.softmax(dim=2, dtype=torch.float32).to(win_kv.dtype)
            compressed = self.kv_norm((win_kv * pool_w).sum(dim=2))
            abs_pos = torch.arange(n_windows, device=compressed.device) * self.compress_rate + first_window_position
            abs_pos = abs_pos.unsqueeze(0).expand(batch, -1)
            cos, sin = self.rotary_emb(compressed, position_ids=abs_pos, layer_type=self.rope_layer_type)
            compressed = apply_agnes_rope(compressed.unsqueeze(1), cos, sin).squeeze(1)
        else:
            compressed = chunk_kv.new_zeros((batch, 0, self.head_dim))

        if cache_layer is not None:
            compressed = cache_layer.update_compressor_states("compressor", compressed)
        compressed_kv = compressed.unsqueeze(1)

        # ask the indexer which entries each query keeps; it returns -1 for picks
        # that are not yet attendable. Clamp those to a scratch column, scatter
        # zeros for the kept ones, then slice the scratch column back off so the
        # bias is -inf everywhere except the valid picks.
        picks = self.indexer(hidden_states, q_residual, position_ids, past_key_values, layer_idx)  # [B, S, k]
        n_entries = compressed_kv.shape[2]
        kept = picks >= 0  # [B, S, k]
        gather_idx = torch.where(kept, picks, torch.full_like(picks, n_entries))
        block_bias = compressed_kv.new_full((batch, 1, seq_len, n_entries + 1), float("-inf"))
        block_bias.scatter_(-1, gather_idx.unsqueeze(1), 0.0)
        return compressed_kv, block_bias[..., :n_entries]



# ==========================================================================
# Core attention
# ==========================================================================

class AgnesGroupedLinear(nn.Linear):
    """Block-diagonal grouped linear used by the grouped output projection
    The core attention's stacked output is `num_attention_heads* head_dim`-dim,
    which is *very* large (Agnes-Flash: 32768; Agnes-Pro: 65536). A direct
    `num_attention_heads*head_dim → hidden_size` projection would dominate the per-token cost.

    Agnes sidesteps that by splitting the heads into `g` groups, projecting
    each `num_attention_heads * head_dim/g`-dim group independently to a `d_g`-dim intermediate output
    (with `d_g < num_attention_heads * head_dim/g`), and then mixing the resulting `g·d_g` vector to
    `hidden_size` through a single follow-up linear (`self_attn.o_b_proj`). This
    module owns the per-group block (`self_attn.o_a_proj`).

    For Agnes-Flash (num_attention_heads=64, head_dim=512, o_groups=8, o_lora_rank=1024,
    hidden_size=4096), g=8 groups of 4096-dim each are projected to 1024-dim, then
    mixed to 4096-dim; for Agnes-Pro (num_attention_heads=128, head_dim=512, o_groups=16,
    o_lora_rank=1024, hidden_size=7168), g=16 groups of 4096-dim each are projected
    to 1024-dim, then mixed to 7168-dim.
    """

    def __init__(self, in_features_per_group: int, out_features: int, n_groups: int, bias: bool = False):
        super().__init__(in_features_per_group, out_features, bias=bias)
        self.n_groups = n_groups

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        lead = x.shape[:-2]
        in_dim = x.shape[-1]
        # per-group matmul: bring the group axis to the front so a single bmm
        # applies each group's block independently.
        blocks = self.weight.view(self.n_groups, -1, in_dim).transpose(1, 2)
        grouped = x.reshape(-1, self.n_groups, in_dim).transpose(0, 1)
        out = torch.bmm(grouped, blocks).transpose(0, 1)
        return out.reshape(*lead, self.n_groups, -1)


def _broadcast_kv_heads(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """Broadcast the single shared KV head out to `n_rep` query heads. Same
    result as `torch.repeat_interleave(x, n_rep, dim=1)`, taking
    `(B, n_kv, S, D)` to `(B, n_kv * n_rep, S, D)`, but expressed as an
    `expand` + `reshape` to avoid materialising the intermediate copy.
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)


def _eager_attention(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    scaling: float,
    dropout: float | int = 0.0,
    **kwargs,
):
    # Shared-KV MQA: broadcast the one stored KV head across all query heads.
    key_states = _broadcast_kv_heads(key, module.num_key_value_groups)
    value_states = _broadcast_kv_heads(value, module.num_key_value_groups)
    logits = torch.matmul(query, key_states.transpose(2, 3)) * scaling
    if attention_mask is not None:
        logits = logits + attention_mask

    # Append one learned per-head sink logit as an extra "always-visible" column,
    # then take the softmax over [keys | sink] and read back only the key
    # columns. The shift by the row max is a plain numerical-stability
    # subtraction (matters in bf16/fp16, invariant of the softmax).
    sink = module.sinks.reshape(1, -1, 1, 1).expand(query.shape[0], -1, query.shape[-2], -1)
    logits = torch.cat([logits, sink], dim=-1)
    logits = logits - logits.amax(dim=-1, keepdim=True)
    probs = F.softmax(logits, dim=-1, dtype=logits.dtype)[..., :-1]
    probs = nn.functional.dropout(probs, p=dropout, training=module.training).to(value_states.dtype)
    attn_output = torch.matmul(probs, value_states).transpose(1, 2).contiguous()
    return attn_output, probs


_COMPRESSOR_BY_LAYER_TYPE = {
    "agnes_local_attention": None,
    "agnes_sparse_attention": AgnesCSACompressor,
    "agnes_pooled_attention": AgnesHCACompressor,
}


class AgnesAttention(nn.Module):
    r"""
    Agnes attention block. It departs from a textbook multi-head block in five
    ways:

    * Shared-KV multi-query attention: `num_key_value_heads = 1`; `kv_proj`
      emits that single KV head and it is read as both key and value.
    * Partial, interleaved RoPE over the leading `rope_head_dim` of each head.
      The conjugate rotation (position `-i`) is re-applied to the output rope
      slice so each entry's contribution depends only on query/key relative
      distance.
    * A learned per-head attention-sink logit that absorbs probability mass
      away from the real keys.
    * A grouped low-rank output projection, to keep the wide stacked-head
      output affordable.
    * Three interchangeable cache regimes: plain sliding window, sliding+CSA,
      and sliding+HCA.
    """

    def __init__(self, config: AgnesConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.layer_type = config.layer_types[layer_idx]
        # Sliding-only layers use the "main" (plain θ=10000) rope; CSA/HCA layers
        # share the same yarn-scaled "compress" rope as their compressor.
        self.rope_layer_type = "main" if self.layer_type == "agnes_local_attention" else "compress"
        self.num_heads = config.num_attention_heads
        self.num_key_value_groups = config.num_attention_heads  # single KV head, broadcast to all
        self.head_dim = config.head_dim
        self.sliding_window = config.sliding_window
        self.attention_dropout = config.attention_dropout
        self.is_causal = True
        self.scaling = self.head_dim**-0.5

        self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=False)
        self.q_a_norm = AgnesRMSNorm(config.q_lora_rank, eps=config.rms_norm_eps)
        self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.head_dim, bias=False)
        self.q_b_norm = AgnesUnweightedRMSNorm(eps=config.rms_norm_eps)
        self.kv_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False)
        self.kv_norm = AgnesRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.o_a_proj = AgnesGroupedLinear(
            self.num_heads * self.head_dim // config.o_groups, config.o_groups * config.o_lora_rank, config.o_groups
        )
        self.o_b_proj = nn.Linear(config.o_groups * config.o_lora_rank, config.hidden_size, bias=False)
        self.sinks = nn.Parameter(torch.empty(self.num_heads))
        self.compressor = (
            _COMPRESSOR_BY_LAYER_TYPE[self.layer_type](config) if self.layer_type != "agnes_local_attention" else None
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: dict[str, tuple[torch.Tensor, torch.Tensor]] | tuple[torch.Tensor, torch.Tensor],
        position_ids: torch.Tensor,
        attention_mask: torch.Tensor | None,
        past_key_values: Cache | None = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        lead = hidden_states.shape[:-1]
        head_shape = (*lead, -1, self.head_dim)
        # the model hands down a {"main", "compress"} cos/sin dict; this layer
        # takes whichever entry matches its rope label (local -> main, else compress).
        cos, sin = position_embeddings[self.rope_layer_type]

        # low-rank query path: down-proj + norm, then up-proj into heads, a second
        # (unweighted) norm, and finally rotary.
        q_lora = self.q_a_norm(self.q_a_proj(hidden_states))
        q = self.q_b_norm(self.q_b_proj(q_lora).view(*head_shape).transpose(1, 2))
        q = apply_agnes_rope(q, cos, sin)

        kv = self.kv_norm(self.kv_proj(hidden_states)).view(*head_shape).transpose(1, 2)
        kv = apply_agnes_rope(kv, cos, sin)

        if past_key_values is not None:  # shared-KV sliding window (K is V)
            kv = past_key_values.update(kv, kv, self.layer_idx)[0]

        block_bias = None
        if self.compressor is not None:  # long-range CSA / HCA entries
            compressed_kv, block_bias = self.compressor(
                hidden_states, q_lora, position_ids, past_key_values, self.layer_idx
            )
            kv = torch.cat([kv, compressed_kv], dim=2)

        # The compressor path concatenates extra entries onto the KV axis after the
        # standard sliding-window cache update, so a tensor `attention_mask` (built
        # for the pre-concat KV length) needs to be extended to cover them. The
        # compressor returns a `block_bias` carrying per-query causality + indexer
        # validity over those new slots — cat it in instead of zero-padding (which
        # would let every query see every compressed slot).
        if isinstance(attention_mask, torch.Tensor) and kv.shape[2] > attention_mask.shape[-1]:
            if block_bias is not None:
                attention_mask = torch.cat([attention_mask, block_bias.to(attention_mask.dtype)], dim=-1)
            else:
                attention_mask = F.pad(attention_mask, (0, kv.shape[2] - attention_mask.shape[-1]), value=0.0)

        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, _eager_attention
        )
        attn_output, attn_weights = attention_interface(
            self,
            q,
            kv,
            kv,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            sliding_window=self.sliding_window,
            s_aux=self.sinks,
            **kwargs,
        )

        # Because V shares K's storage it also carries K's rotary; undo it on the
        # output's rope slice with the conjugate rotation (`-sin`) at the query
        # position, so what reaches the output projection is rope-free. The
        # transpose pair only re-lays [B, H, S, D] as the [B, S, H, D] that
        # apply_agnes_rope expects.
        attn_output = apply_agnes_rope(attn_output.transpose(1, 2), cos, -sin).transpose(1, 2)

        # grouped low-rank output projection: split heads into o_groups, project
        # each group down, concatenate, then mix up to hidden_size.
        heads = attn_output.reshape(*lead, self.config.o_groups, -1)
        mixed = self.o_a_proj(heads).flatten(2)
        return self.o_b_proj(mixed), attn_weights



# ==========================================================================
# Decoder block
# ==========================================================================

class AgnesDecoderLayer(GradientCheckpointingLayer):
    r"""Agnes decoder block. Unlike a textbook residual block, the residual
    here is not a single tensor but a stack of `hc_mult` parallel streams kept
    in shape `[B, S, hc_mult, D]` for the whole block. Two
    :class:`AgnesHyperConnection` modules (one at the attention site, one at the
    MLP site) collapse those streams into the sublayer input and re-expand the
    sublayer output back across them. Their stream-mixing matrix is projected
    onto the doubly-stochastic manifold by Sinkhorn-Knopp iteration, which keeps
    the residual transform non-expansive across a deep stack.
    """

    def __init__(self, config: AgnesConfig, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.self_attn = AgnesAttention(config, layer_idx)
        self.mlp = AgnesSparseMoeBlock(config, layer_idx)
        self.input_layernorm = AgnesRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = AgnesRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.attn_hc = AgnesHyperConnection(config)
        self.ffn_hc = AgnesHyperConnection(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        input_ids: torch.Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> torch.Tensor:
        # The stream stack stays [B, S, hc_mult, hidden] throughout. Each site
        # runs the same recipe: collapse streams -> sublayer -> place the output
        # back across streams (`post`) and carry the old streams through the
        # doubly-stochastic mixer (`comb`). `post`/`comb` arrive in fp32 (Sinkhorn
        # runs in float) and are cast down before mixing. `comb` is used
        # transposed — sum over its FIRST axis, i.e. comb.T @ streams — because
        # the Sinkhorn matrix is doubly-stochastic but not symmetric.
        dtype = hidden_states.dtype

        post, comb, collapsed = self.attn_hc(hidden_states)
        attn_out, _ = self.self_attn(self.input_layernorm(collapsed), **kwargs)
        placed = post.to(dtype).unsqueeze(-1) * attn_out.unsqueeze(-2)
        carried = torch.matmul(comb.to(dtype).transpose(-1, -2), hidden_states)
        hidden_states = placed + carried

        post, comb, collapsed = self.ffn_hc(hidden_states)
        mlp_out = self.mlp(self.post_attention_layernorm(collapsed), input_ids=input_ids)
        placed = post.to(dtype).unsqueeze(-1) * mlp_out.unsqueeze(-2)
        carried = torch.matmul(comb.to(dtype).transpose(-1, -2), hidden_states)
        return placed + carried



# ==========================================================================
# Auxiliary load-balancing loss
# ==========================================================================

def load_balancing_loss_func(
    gate_logits: torch.Tensor | tuple[torch.Tensor] | None,
    num_experts: int | None = None,
    top_k=2,
    attention_mask: torch.Tensor | None = None,
) -> torch.Tensor | int:
    r"""
    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.

    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
    experts is too unbalanced.

    Args:
        gate_logits:
            Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
            shape [batch_size X sequence_length, num_experts].
        num_experts:
            Number of experts
        top_k:
            The number of experts to route per-token, can be also interpreted as the `top-k` routing
            parameter.
        attention_mask (`torch.Tensor`, *optional*):
            The attention_mask used in forward function
            shape [batch_size X sequence_length] if not None.

    Returns:
        The auxiliary loss.
    """
    if gate_logits is None or not isinstance(gate_logits, tuple):
        return 0

    if isinstance(gate_logits, tuple):
        compute_device = gate_logits[0].device
        concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)

    routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)

    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)

    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)

    if attention_mask is None:
        # Compute the percentage of tokens routed to each experts
        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)

        # Compute the average probability of routing to these experts
        router_prob_per_expert = torch.mean(routing_weights, dim=0)
    else:
        batch_size, sequence_length = attention_mask.shape
        num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)

        # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
        expert_attention_mask = (
            attention_mask[None, :, :, None, None]
            .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
            .reshape(-1, top_k, num_experts)
            .to(compute_device)
        )

        # Compute the percentage of tokens routed to each experts
        tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
            expert_attention_mask, dim=0
        )

        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
        router_per_expert_attention_mask = (
            attention_mask[None, :, :, None]
            .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
            .reshape(-1, num_experts)
            .to(compute_device)
        )

        # Compute the average probability of routing to these experts
        router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
            router_per_expert_attention_mask, dim=0
        )

    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
    return overall_loss * num_experts



# ==========================================================================
# Pretrained base and full models
# ==========================================================================

@auto_docstring
class AgnesPreTrainedModel(PreTrainedModel):
    config: AgnesConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["AgnesDecoderLayer"]
    _skip_keys_device_placement = ["past_key_values"]
    # Agnes runs the eager attention path only; the optimized backends are all
    # unavailable for a concrete structural reason:
    #
    #   * FlashAttention 2 / 3 / 4 cap the supported head dim at 256, but Agnes
    #     uses `head_dim=512`, so those kernels raise
    #     `FlashAttention forward only supports head dimension at most 256`.
    #   * SDPA: the fused torch SDPA kernel has no slot for the per-head learned
    #     attention-sink logit, which is part of Agnes's softmax.
    #   * FlexAttention: the compressor concatenates a *variable* number of
    #     entries onto the KV axis inside the block, after the model-level mask
    #     was already built, so the KV length no longer matches the BlockMask's
    #     `kv_len`. BlockMask can't be resized at runtime, and the compressor
    #     already carries its own causal bookkeeping, so wiring that into a
    #     `mask_mod` is not worthwhile.
    _supports_flash_attn = False
    _supports_sdpa = False
    _supports_flex_attn = False
    # The compressor's rolling-window buffer / compressed-entries / overlap state
    # lives on the per-layer cache (:class:`AgnesHCACache` /
    # :class:`AgnesCSACache`) and isn't compatible with :class:`StaticCache`
    # — that path would hand the compressor a :class:`StaticSlidingWindowLayer`
    # with no `store_compression_weights` method. Disabling fullgraph compile
    # keeps generation tests on the dynamic cache build that does dispatch to
    # Agnes's own cache layers.
    _can_compile_fullgraph = False
    _supports_attention_backend = True
    _can_record_outputs = {
        "router_logits": OutputRecorder(AgnesTopKRouter, index=0),
        "hidden_states": AgnesDecoderLayer,
        "attentions": AgnesAttention,
    }
    config_class = AgnesConfig
    _keep_in_fp32_modules_strict = [
        "attn_hc",
        "ffn_hc",
        "hc_head",
        "sinks",
        "position_bias",
        "e_score_correction_bias",
        "q_a_norm",
        "kv_norm",
        "input_layernorm",
        "post_attention_layernorm",
        "norm",
    ]
    # Agnes-Flash checkpoints mix FP8 and BF16 in the attention compressor /
    # indexer branch: these projections ship in BF16 with no companion `scale_inv`.
    # Listed here (non-strict) so the FP8 quantizer's `get_modules_to_not_convert`
    # auto-skips them; non-strict has no dtype effect at BF16, so they stay BF16.
    _keep_in_fp32_modules = [
        "self_attn.compressor.kv_proj",
        "self_attn.compressor.gate_proj",
        "self_attn.compressor.indexer.kv_proj",
        "self_attn.compressor.indexer.gate_proj",
        "self_attn.compressor.indexer.scorer.weights_proj",
    ]
    _keys_to_ignore_on_load_unexpected = [r"(^|\.)mtp\..*"]
    # ``_is_stateful`` opts out of generation modes that need to roll the cache
    # back across drafts (assisted generation, prompt lookup, contrastive search).
    # The compressor's running-window state isn't rewindable, so `generate`
    # raises a clear error early instead of failing deep in the compressor with
    # a missing-method `AttributeError`.
    _is_stateful = True

    @torch.no_grad()
    def _init_weights(self, module):
        super()._init_weights(module)
        std = self.config.initializer_range
        if isinstance(module, (AgnesTopKRouter, AgnesHashRouter)):
            init.normal_(module.weight, mean=0.0, std=std)
            if isinstance(module, AgnesTopKRouter):
                init.zeros_(module.e_score_correction_bias)  # buffer
            if isinstance(module, AgnesHashRouter):
                init.zeros_(module.tid2eid)  # buffer; real values come from the checkpoint
        elif isinstance(module, AgnesExperts):
            init.normal_(module.gate_up_proj, mean=0.0, std=std)
            init.normal_(module.down_proj, mean=0.0, std=std)
        elif isinstance(module, AgnesAttention):
            init.zeros_(module.sinks)
        elif isinstance(module, AgnesHyperConnection):
            init.normal_(module.fn, mean=0.0, std=std)
            init.zeros_(module.base)
            init.ones_(module.scale)
        elif isinstance(module, AgnesHyperHead):
            init.normal_(module.hc_fn, mean=0.0, std=std)
            init.zeros_(module.hc_base)
            init.ones_(module.hc_scale)
        elif isinstance(module, (AgnesHCACompressor, AgnesCSACompressor, AgnesIndexer)):
            init.zeros_(module.position_bias)
        elif isinstance(module, AgnesRotaryEmbedding):
            for layer_type in module.layer_types:
                rope_init_fn = module.compute_default_rope_parameters
                if module.rope_type[layer_type] != "default":
                    rope_init_fn = ROPE_INIT_FUNCTIONS[module.rope_type[layer_type]]
                curr_inv_freq, _ = rope_init_fn(module.config, layer_type=layer_type)
                init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)
                init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)


@auto_docstring
class AgnesModel(AgnesPreTrainedModel):
    def __init__(self, config: AgnesConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList(
            [AgnesDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = AgnesRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = AgnesRotaryEmbedding(config)
        self.gradient_checkpointing = False
        self.hc_head = AgnesHyperHead(config)

        # Initialize weights and apply final processing
        self.post_init()

    @merge_with_config_defaults
    @capture_outputs
    @auto_docstring
    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        use_cache: bool | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> MoeModelOutputWithPast:
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
        # what we hand back is the caller's own cache (or None); a cache we spin
        # up internally is used for this pass only, never returned.
        return_cache = past_key_values if use_cache else None
        if past_key_values is None:
            past_key_values = DynamicCache(config=self.config)
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        if position_ids is None:
            start = past_key_values.get_seq_length()
            position_ids = (torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + start).unsqueeze(0)

        # generate() can pass an already-built per-layer-type mask dict; every
        # Agnes layer shares one sliding-window mask, so reuse it if present.
        if isinstance(attention_mask, dict):
            causal_mask = next(iter(attention_mask.values()))
        else:
            causal_mask = create_sliding_window_causal_mask(
                config=self.config,
                inputs_embeds=inputs_embeds,
                attention_mask=attention_mask,
                past_key_values=past_key_values,
                position_ids=position_ids,
            )

        # fan the embeddings into hc_mult parallel residual streams and precompute
        # both rope flavours once for the whole stack.
        hidden_states = inputs_embeds.unsqueeze(2).expand(-1, -1, self.config.hc_mult, -1).contiguous()
        position_embeddings = {
            "main": self.rotary_emb(inputs_embeds, position_ids=position_ids, layer_type="main"),
            "compress": self.rotary_emb(inputs_embeds, position_ids=position_ids, layer_type="compress"),
        }

        for layer in self.layers:
            hidden_states = layer(
                hidden_states,
                position_embeddings=position_embeddings,
                position_ids=position_ids,
                attention_mask=causal_mask,
                input_ids=input_ids,
                past_key_values=past_key_values,
                **kwargs,
            )

        hidden_states = self.norm(self.hc_head(hidden_states))
        return MoeModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=return_cache)


@auto_docstring
class AgnesForCausalLM(AgnesPreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
    _tp_plan = {"lm_head": "colwise_gather_output"}
    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}

    def __init__(self, config):
        super().__init__(config)
        self.model = AgnesModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.router_aux_loss_coef = config.router_aux_loss_coef
        self.num_experts = config.num_local_experts
        self.num_experts_per_tok = config.num_experts_per_tok

        # Initialize weights and apply final processing
        self.post_init()

    @can_return_tuple
    @auto_docstring
    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        labels: torch.LongTensor | None = None,
        use_cache: bool | None = None,
        output_router_logits: bool | None = None,
        logits_to_keep: int | torch.Tensor = 0,
        **kwargs: Unpack[TransformersKwargs],
    ) -> MoeCausalLMOutputWithPast:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

        Example:

        ```python
        >>> from transformers import AutoTokenizer, AutoModelForCausalLM

        >>> model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, trust_remote_code=True)
        >>> tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)

        >>> inputs = tokenizer("The capital of France is", return_tensors="pt")
        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True)[0]
        ```"""

        output_router_logits = (
            output_router_logits if output_router_logits is not None else self.config.output_router_logits
        )

        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
        outputs: MoeModelOutputWithPast = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_router_logits=output_router_logits,
            **kwargs,
        )

        hidden_states = outputs.last_hidden_state
        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(hidden_states[:, slice_indices, :])

        loss = None
        if labels is not None:
            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)

        aux_loss = None
        if output_router_logits:
            aux_loss = load_balancing_loss_func(
                outputs.router_logits,
                self.num_experts,
                self.num_experts_per_tok,
                attention_mask,
            )
            if labels is not None:
                loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device

        return MoeCausalLMOutputWithPast(
            loss=loss,
            aux_loss=aux_loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            router_logits=outputs.router_logits,
        )


__all__ = ["AgnesPreTrainedModel", "AgnesModel", "AgnesForCausalLM"]



# ==========================================================================
# Checkpoint weight conversion
# ==========================================================================

# Agnes checkpoints ship native-format tensor names (`embed.weight`,
# `layers.N.attn.*`, `layers.N.ffn.*`, flat `hc_*` hyper-connection params,
# per-expert `w1`/`w2`/`w3` weights in the native gate/down/up notation). The mapping below
# renames them onto this implementation's module tree and merges the
# per-expert weights into the fused `gate_up_proj` / `down_proj` tensors.
# Registered for `model_type="agnes"`, it is applied automatically by
# `from_pretrained` and reversed by `save_pretrained`.
_AGNES_CHECKPOINT_CONVERSION_MAPPING = [
    WeightRenaming(source_patterns=r"^embed\.weight$", target_patterns="embed_tokens.weight"),
    WeightRenaming(source_patterns=r"^head\.weight$", target_patterns="lm_head.weight"),
    WeightRenaming(source_patterns=r"^norm\.weight$", target_patterns="norm.weight"),
    WeightRenaming(source_patterns=r"^hc_head_fn$", target_patterns="hc_head.hc_fn"),
    WeightRenaming(source_patterns=r"^hc_head_base$", target_patterns="hc_head.hc_base"),
    WeightRenaming(source_patterns=r"^hc_head_scale$", target_patterns="hc_head.hc_scale"),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.attn_norm\.",
        target_patterns=r"layers.\1.input_layernorm.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.ffn_norm\.",
        target_patterns=r"layers.\1.post_attention_layernorm.",
    ),
    WeightRenaming(source_patterns=r"^layers\.(\d+)\.hc_attn_fn$", target_patterns=r"layers.\1.attn_hc.fn"),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.hc_attn_base$", target_patterns=r"layers.\1.attn_hc.base"
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.hc_attn_scale$", target_patterns=r"layers.\1.attn_hc.scale"
    ),
    WeightRenaming(source_patterns=r"^layers\.(\d+)\.hc_ffn_fn$", target_patterns=r"layers.\1.ffn_hc.fn"),
    WeightRenaming(source_patterns=r"^layers\.(\d+)\.hc_ffn_base$", target_patterns=r"layers.\1.ffn_hc.base"),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.hc_ffn_scale$", target_patterns=r"layers.\1.ffn_hc.scale"
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.attn\.",
        target_patterns=r"layers.\1.self_attn.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.ffn\.",
        target_patterns=r"layers.\1.mlp.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.attn_sink$",
        target_patterns=r"layers.\1.self_attn.sinks",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.indexer\.compressor\.norm\.",
        target_patterns=r"layers.\1.self_attn.compressor.indexer.kv_norm.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.indexer\.compressor\.ape$",
        target_patterns=r"layers.\1.self_attn.compressor.indexer.position_bias",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.indexer\.compressor\.",
        target_patterns=r"layers.\1.self_attn.compressor.indexer.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.indexer\.",
        target_patterns=r"layers.\1.self_attn.compressor.indexer.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.compressor\.indexer\.weights_proj\.",
        target_patterns=r"layers.\1.self_attn.compressor.indexer.scorer.weights_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.compressor\.norm\.",
        target_patterns=r"layers.\1.self_attn.compressor.kv_norm.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.compressor\.ape$",
        target_patterns=r"layers.\1.self_attn.compressor.position_bias",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.(.*?)\.wq_a\.",
        target_patterns=r"layers.\1.self_attn.\2.q_a_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.(.*?)\.wq_b\.",
        target_patterns=r"layers.\1.self_attn.\2.q_b_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.(.*?)\.wkv\.",
        target_patterns=r"layers.\1.self_attn.\2.kv_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.(.*?)\.wgate\.",
        target_patterns=r"layers.\1.self_attn.\2.gate_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.(.*?)\.wo_a\.",
        target_patterns=r"layers.\1.self_attn.\2.o_a_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.(.*?)\.wo_b\.",
        target_patterns=r"layers.\1.self_attn.\2.o_b_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.wq_a\.",
        target_patterns=r"layers.\1.self_attn.q_a_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.wq_b\.",
        target_patterns=r"layers.\1.self_attn.q_b_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.wkv\.",
        target_patterns=r"layers.\1.self_attn.kv_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.wo_a\.",
        target_patterns=r"layers.\1.self_attn.o_a_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.wo_b\.",
        target_patterns=r"layers.\1.self_attn.o_b_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.self_attn\.q_norm\.",
        target_patterns=r"layers.\1.self_attn.q_a_norm.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.mlp\.gate\.bias$",
        target_patterns=r"layers.\1.mlp.gate.e_score_correction_bias",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.mlp\.shared_experts\.w1\.",
        target_patterns=r"layers.\1.mlp.shared_experts.gate_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.mlp\.shared_experts\.w2\.",
        target_patterns=r"layers.\1.mlp.shared_experts.down_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.mlp\.shared_experts\.w3\.",
        target_patterns=r"layers.\1.mlp.shared_experts.up_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.mlp\.parallel_ffn\.w1\.",
        target_patterns=r"layers.\1.mlp.parallel_ffn.gate_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.mlp\.parallel_ffn\.w2\.",
        target_patterns=r"layers.\1.mlp.parallel_ffn.down_proj.",
    ),
    WeightRenaming(
        source_patterns=r"^layers\.(\d+)\.mlp\.parallel_ffn\.w3\.",
        target_patterns=r"layers.\1.mlp.parallel_ffn.up_proj.",
    ),
    WeightConverter(
        source_patterns=[
            "mlp.experts.*.w1.weight",
            "mlp.experts.*.w3.weight",
        ],
        target_patterns="mlp.experts.gate_up_proj",
        operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
    ),
    WeightConverter(
        source_patterns="mlp.experts.*.w2.weight",
        target_patterns="mlp.experts.down_proj",
        operations=[MergeModulelist(dim=0)],
    ),
]


try:
    register_checkpoint_conversion_mapping("agnes", _AGNES_CHECKPOINT_CONVERSION_MAPPING)
except ValueError:
    pass  # already registered (module imported more than once)