"""Senryu-2: a Gated DeltaNet-2 + Parallax hybrid on a 4-branch gated residual stream.

Every sublayer reads its input from the residual branches through Qwen's
Gated Residual (per-branch RMSNorm, elementwise sigmoid gate from a low-rank
bottleneck over all branches) and writes its output back to every branch with
one data-dependent scalar per branch. A hashed n-gram memory, whose tables live
in host memory during training, is injected before block ``ngram_layer_index``
(docs/architecture.md, "Senryu-2").

Training and scoring run packed: the batch is flattened into one row and split
into documents at every BOS token, so GDN-2, Parallax, the n-gram hashes and
the n-gram convolution never look across a document boundary.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F
from fla.layers.gdn2 import GatedDeltaNet2
from fla.models.utils import Cache as FlaCache
from fla.modules.rotary import rotary_embedding
from safetensors.torch import save_model
from torch import nn
from torch.utils.checkpoint import checkpoint

from tiny_gdn.config import TinyGDNConfig
from tiny_gdn.gated_residual_ops import fused_gated_residual_read, fused_gated_residual_write
from tiny_gdn.haiku_layers import SiTUGLU
from tiny_gdn.model import TinyGDNOutput, linear_cross_entropy
from tiny_gdn.ngram import (
    NGRAM_TABLES_EMA_FILE,
    NGRAM_TABLES_FILE,
    NGramLayout,
    document_distance,
    load_table_file,
    ngram_row_indices,
)
from tiny_gdn.nn_common import RMSNorm
from tiny_gdn.parallax_ops import parallax_attention, parallax_decode_cached

GATED_RESIDUAL_CHUNK_ROWS = 4096
KV_CACHE_GRANULE = 256


def mark_muon_row_blocks(parameter: nn.Parameter, blocks: int) -> None:
    """Orthogonalize ``parameter`` as ``blocks`` equal row blocks (heads, fused halves)."""

    if parameter.shape[0] % blocks:
        raise ValueError(f"{blocks} row blocks do not divide {parameter.shape[0]} rows")
    parameter._muon_row_blocks = blocks


def mark_adamw_only(*parameters: nn.Parameter) -> None:
    """Keep 2D weights that are gates or elongated low-rank maps on AdamW."""

    for parameter in parameters:
        parameter._adamw_only = True


def mark_gain(*parameters: nn.Parameter) -> None:
    """2D per-branch norm gains and filters: AdamW without weight decay, like 1D gains."""

    mark_adamw_only(*parameters)
    for parameter in parameters:
        parameter._no_weight_decay = True


# --------------------------------------------------------------------------- layout


@dataclass
class Senryu2Context:
    """Per-forward state shared by every block (passes through activation checkpointing)."""

    batch: int
    length: int
    cu_seqlens: torch.Tensor | None
    distance: torch.Tensor | None
    rotary: tuple[torch.Tensor, torch.Tensor]
    position_offset: int = 0
    cache: Senryu2Cache | None = None


class RotaryTables:
    """Partial-RoPE cos/sin rows ``[length, rotary_dim / 2]``, built in FP32 and cached."""

    def __init__(self, rotary_dim: int, theta: float) -> None:
        self.rotary_dim = rotary_dim
        self.theta = theta
        self._cached: tuple[torch.Tensor, torch.Tensor] | None = None
        self._key: tuple[torch.device, torch.dtype, bool] | None = None

    def get(self, length: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
        key = (device, dtype, torch.is_inference_mode_enabled())
        cached = self._cached
        if cached is None or self._key != key or cached[0].shape[0] < length:
            rows = max(length, 4096 if cached is None else 2 * cached[0].shape[0])
            inverse_frequency = 1.0 / (
                self.theta
                ** (torch.arange(0, self.rotary_dim, 2, device=device, dtype=torch.float32) / self.rotary_dim)
            )
            angles = torch.outer(torch.arange(rows, device=device, dtype=torch.float32), inverse_frequency)
            cached = (angles.cos().to(dtype), angles.sin().to(dtype))
            self._cached = cached
            self._key = key
        return cached


def parallax_reference(
    q: torch.Tensor,
    r: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    *,
    scale: float,
    cu_seqlens: torch.Tensor | None = None,
) -> torch.Tensor:
    """FP32 Parallax for queries at the last ``Tq`` of ``Tk`` key positions.

    ``out = O1 / d1 * (1 + d2 / d1) - O2 / d1`` with ``p1 = exp(scale q·k)``,
    ``p2 = p1 * (r·k)``; ``cu_seqlens`` (batch 1, ``Tq == Tk``) masks across
    packed documents.
    """

    batch, queries, query_heads, head_dim = q.shape
    keys, kv_heads = k.shape[1], k.shape[2]
    groups = query_heads // kv_heads
    qf = q.float().reshape(batch, queries, kv_heads, groups, head_dim)
    rf = r.float().reshape(batch, queries, kv_heads, groups, head_dim)
    kf, vf = k.float(), v.float()
    s1 = torch.einsum("bqhgd,bkhd->bhgqk", qf, kf) * scale
    s2 = torch.einsum("bqhgd,bkhd->bhgqk", rf, kf)
    query_positions = torch.arange(keys - queries, keys, device=q.device)[:, None]
    key_positions = torch.arange(keys, device=q.device)[None, :]
    masked = key_positions > query_positions
    if cu_seqlens is not None:
        segment = torch.zeros(keys, dtype=torch.long, device=q.device)
        segment[cu_seqlens[1:-1].long()] = 1
        segment = segment.cumsum(0)
        masked = masked | (segment[keys - queries :, None] != segment[None, :])
    s1 = s1.masked_fill(masked, float("-inf"))
    p1 = (s1 - s1.amax(dim=-1, keepdim=True)).exp()
    d1 = p1.sum(dim=-1)
    p2 = p1 * s2
    d2 = p2.sum(dim=-1)
    o1 = torch.einsum("bhgqk,bkhd->bqhgd", p1, vf)
    o2 = torch.einsum("bhgqk,bkhd->bqhgd", p2, vf)
    inverse_d1 = (1.0 / d1).permute(0, 3, 1, 2)[..., None]
    correction = (d2 / d1).permute(0, 3, 1, 2)[..., None]
    out = o1 * inverse_d1 * (1.0 + correction) - o2 * inverse_d1
    return out.reshape(batch, queries, query_heads, head_dim).to(q.dtype)


# ------------------------------------------------------------------ gated residual


class _GatedResidualRead(torch.autograd.Function):
    """``x = (1/n) Σ σ(W_u SiLU(W_d vec(R̂) / n))_i ⊙ R̂_i`` and ``s = 2σ(W_w vec(R̂) / n)``.

    Saves only the branches, the per-branch inverse RMS and the bottleneck
    pre-activation; the normalized branches and gates are recomputed in
    chunks during backward, so a 32k-token read keeps ~0.4 GB instead of ~4 GB.
    CUDA runs the fused kernels in ``gated_residual_ops`` instead.
    """

    @staticmethod
    def forward(ctx, residual, gamma, down, up, write, eps):
        rows, branches, hidden_size = residual.shape
        dtype = down.dtype
        gain = 1.0 + gamma.float()
        out = residual.new_empty(rows, hidden_size)
        rstd = residual.new_empty(rows, branches, dtype=torch.float32)
        hidden = residual.new_empty(rows, down.shape[0])
        scales = residual.new_empty(rows, 0 if write is None else branches, dtype=torch.float32)
        for start in range(0, rows, GATED_RESIDUAL_CHUNK_ROWS):
            stop = min(start + GATED_RESIDUAL_CHUNK_ROWS, rows)
            values = residual[start:stop].float()
            chunk_rstd = torch.rsqrt(values.square().mean(dim=-1, keepdim=True) + eps)
            hat = values * chunk_rstd * gain
            flat = hat.to(dtype).flatten(1)
            chunk_hidden = (flat @ down.t()) / branches
            activated = F.silu(chunk_hidden.float()).to(dtype)
            gates = torch.sigmoid((activated @ up.t()).float()).view(stop - start, branches, hidden_size)
            out[start:stop] = ((gates * hat).sum(dim=1) / branches).to(out.dtype)
            rstd[start:stop] = chunk_rstd.squeeze(-1)
            hidden[start:stop] = chunk_hidden
            if write is not None:
                scales[start:stop] = 2.0 * torch.sigmoid((flat @ write.t()).float() / branches)
        ctx.save_for_backward(residual, gamma, down, up, write, rstd, hidden, scales)
        scales_out = scales.to(residual.dtype, copy=True)
        if write is None:
            ctx.mark_non_differentiable(scales_out)
        return out, scales_out

    @staticmethod
    def backward(ctx, grad_out, grad_scales):
        residual, gamma, down, up, write, rstd, hidden, scales = ctx.saved_tensors
        rows, branches, hidden_size = residual.shape
        dtype = down.dtype
        gain = 1.0 + gamma.float()
        grad_residual = torch.empty(residual.shape, dtype=residual.dtype, device=residual.device)
        grad_gamma = torch.zeros(branches, hidden_size, dtype=torch.float32, device=residual.device)
        grad_down = torch.zeros(down.shape, dtype=torch.float32, device=down.device)
        grad_up = torch.zeros(up.shape, dtype=torch.float32, device=up.device)
        grad_write = None if write is None else torch.zeros(write.shape, dtype=torch.float32, device=write.device)
        use_scales = write is not None and grad_scales is not None
        for start in range(0, rows, GATED_RESIDUAL_CHUNK_ROWS):
            stop = min(start + GATED_RESIDUAL_CHUNK_ROWS, rows)
            count = stop - start
            normalized = residual[start:stop].float() * rstd[start:stop, :, None]
            hat = normalized * gain
            flat = hat.to(dtype).flatten(1)
            pre = hidden[start:stop].float()
            sigmoid_pre = torch.sigmoid(pre)
            activated = (pre * sigmoid_pre).to(dtype)
            gates = torch.sigmoid((activated @ up.t()).float()).view(count, branches, hidden_size)
            grad_x = grad_out[start:stop].float().unsqueeze(1) / branches
            grad_hat = gates * grad_x
            grad_gate_logits = (grad_x * hat * gates * (1.0 - gates)).flatten(1).to(dtype)
            grad_up += (grad_gate_logits.t() @ activated).float()
            grad_activated = (grad_gate_logits @ up).float()
            grad_pre = grad_activated * sigmoid_pre * (1.0 + pre * (1.0 - sigmoid_pre))
            grad_pre = (grad_pre / branches).to(dtype)
            grad_down += (grad_pre.t() @ flat).float()
            grad_flat = (grad_pre @ down).float()
            if use_scales:
                scale_values = scales[start:stop]
                grad_scale_logits = grad_scales[start:stop].float() * scale_values * (1.0 - scale_values / 2.0)
                grad_scale_logits = (grad_scale_logits / branches).to(dtype)
                grad_write += (grad_scale_logits.t() @ flat).float()
                grad_flat += (grad_scale_logits @ write).float()
            grad_hat = grad_hat + grad_flat.view(count, branches, hidden_size)
            grad_gamma += (grad_hat * normalized).sum(dim=0)
            grad_normalized = grad_hat * gain
            grad_residual[start:stop] = (
                rstd[start:stop, :, None]
                * (grad_normalized - normalized * (grad_normalized * normalized).mean(dim=-1, keepdim=True))
            ).to(grad_residual.dtype)
        return (
            grad_residual,
            grad_gamma.to(gamma.dtype),
            grad_down.to(down.dtype),
            grad_up.to(up.dtype),
            None if grad_write is None else grad_write.to(write.dtype),
            None,
        )


def gated_residual_read_reference(
    residual: torch.Tensor,
    gamma: torch.Tensor,
    down: torch.Tensor,
    up: torch.Tensor,
    write: torch.Tensor | None,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Eager autograd version of :class:`_GatedResidualRead`, for tests."""

    rows, branches, hidden_size = residual.shape
    gain = 1.0 + gamma.float()
    values = residual.float()
    hat = values * torch.rsqrt(values.square().mean(dim=-1, keepdim=True) + eps) * gain
    flat = hat.to(down.dtype).flatten(1)
    hidden = (flat @ down.t()) / branches
    gates = torch.sigmoid((F.silu(hidden.float()).to(up.dtype) @ up.t()).float()).view(rows, branches, hidden_size)
    out = ((gates * hat).sum(dim=1) / branches).to(residual.dtype)
    if write is None:
        return out, residual.new_empty(rows, 0)
    scales = 2.0 * torch.sigmoid((flat @ write.t()).float() / branches)
    return out, scales.to(residual.dtype)


class GatedResidual(nn.Module):
    """Qwen's Gated Residual read (and write scales) over ``branches`` residual branches."""

    def __init__(self, branches: int, hidden_size: int, rank: int, eps: float, *, write: bool) -> None:
        super().__init__()
        self.branches = branches
        self.eps = eps
        self.weight = nn.Parameter(torch.zeros(branches, hidden_size))
        self.down_proj = nn.Linear(branches * hidden_size, rank, bias=False)
        self.up_proj = nn.Linear(rank, branches * hidden_size, bias=False)
        self.write_proj = nn.Linear(branches * hidden_size, branches, bias=False) if write else None
        mark_gain(self.weight)
        mark_adamw_only(self.down_proj.weight, self.up_proj.weight)
        if self.write_proj is not None:
            mark_adamw_only(self.write_proj.weight)

    def read(self, residual: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """``(x [N, d], scales [N, n])`` from branches ``[N, n, d]``."""

        write = None if self.write_proj is None else self.write_proj.weight
        read = fused_gated_residual_read if residual.is_cuda else _GatedResidualRead.apply
        return read(residual, self.weight, self.down_proj.weight, self.up_proj.weight, write, self.eps)

    @staticmethod
    def write(residual: torch.Tensor, scales: torch.Tensor, update: torch.Tensor) -> torch.Tensor:
        if residual.is_cuda:
            return fused_gated_residual_write(residual, scales, update)
        return residual + scales.unsqueeze(-1) * update.unsqueeze(-2)


# ------------------------------------------------------------------- token mixers


class ParallaxAttention(nn.Module):
    """Parallax attention (softmax plus a learned first-order correction from ``r``).

    Grouped query/key/value heads with per-head zero-centered RMSNorm on q, r
    and k, partial RoPE on q, r and k, and a full-rank sigmoid output gate.
    """

    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
        super().__init__()
        self.layer_index = layer_index
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.attention_head_dim
        self.scale = self.head_dim ** -0.5
        query_width = self.num_heads * self.head_dim
        kv_width = self.num_kv_heads * self.head_dim
        hidden = config.hidden_size
        self.q_proj = nn.Linear(hidden, query_width, bias=False)
        self.r_proj = nn.Linear(hidden, query_width, bias=False)
        self.k_proj = nn.Linear(hidden, kv_width, bias=False)
        self.v_proj = nn.Linear(hidden, kv_width, bias=False)
        self.g_proj = nn.Linear(hidden, query_width, bias=True)
        self.o_proj = nn.Linear(query_width, hidden, bias=False)
        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
        self.r_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
        for projection in (self.q_proj, self.r_proj):
            mark_muon_row_blocks(projection.weight, self.num_heads)
        for projection in (self.k_proj, self.v_proj):
            mark_muon_row_blocks(projection.weight, self.num_kv_heads)
        mark_adamw_only(self.g_proj.weight)

    def forward(self, hidden_states: torch.Tensor, context: Senryu2Context) -> torch.Tensor:
        batch, length, _ = hidden_states.shape
        q = self.q_norm(self.q_proj(hidden_states).view(batch, length, self.num_heads, self.head_dim))
        r = self.r_norm(self.r_proj(hidden_states).view(batch, length, self.num_heads, self.head_dim))
        k = self.k_norm(self.k_proj(hidden_states).view(batch, length, self.num_kv_heads, self.head_dim))
        v = self.v_proj(hidden_states).view(batch, length, self.num_kv_heads, self.head_dim)
        cosine, sine = context.rotary
        q, r, k = (
            rotary_embedding(
                tensor,
                cosine,
                sine,
                seqlen_offsets=context.position_offset,
                cu_seqlens=context.cu_seqlens,
            )
            for tensor in (q, r, k)
        )
        use_kernel = q.is_cuda and q.dtype in (torch.bfloat16, torch.float16)
        cache = context.cache
        if cache is not None:
            keys, values = cache.append_attention(self.layer_index, k, v, context.position_offset)
            if not use_kernel:
                out = parallax_reference(q, r, keys, values, scale=self.scale)
            elif keys.shape[1] == length:
                out = parallax_attention(q, r, keys, values, scale=self.scale)
            else:
                out = parallax_decode_cached(q, r, keys, values, scale=self.scale)
        elif use_kernel:
            out = parallax_attention(q, r, k, v, scale=self.scale, cu_seqlens=context.cu_seqlens)
        else:
            out = parallax_reference(q, r, k, v, scale=self.scale, cu_seqlens=context.cu_seqlens)
        gate = torch.sigmoid(self.g_proj(hidden_states)).view(batch, length, self.num_heads, self.head_dim)
        return self.o_proj((out * gate).reshape(batch, length, self.num_heads * self.head_dim))


def build_gated_deltanet2(config: TinyGDNConfig, layer_index: int) -> GatedDeltaNet2:
    mixer = GatedDeltaNet2(
        hidden_size=config.hidden_size,
        expand_v=config.linear_expand_v,
        head_dim=config.linear_head_dim,
        num_heads=config.linear_num_heads,
        num_v_heads=config.linear_num_value_heads,
        mode="chunk",
        use_short_conv=True,
        allow_neg_eigval=config.allow_negative_eigenvalues,
        conv_size=config.linear_conv_kernel_dim,
        conv_bias=False,
        layer_idx=layer_index,
        norm_eps=config.rms_norm_eps,
    )
    if config.linear_full_rank_gate:
        # The layer calls `self.g_proj(hidden_states)`, so a single full-rank
        # map drops in for fla's low-rank pair.
        mixer.g_proj = nn.Linear(config.hidden_size, mixer.value_dim, bias=True)
        mark_adamw_only(mixer.g_proj.weight)
    else:
        mark_adamw_only(*(layer.weight for layer in mixer.g_proj))
    mark_adamw_only(*(layer.weight for layer in mixer.f_proj))
    for projection in (mixer.q_proj, mixer.k_proj, mixer.b_proj):
        mark_muon_row_blocks(projection.weight, mixer.num_heads)
    for projection in (mixer.v_proj, mixer.w_proj):
        mark_muon_row_blocks(projection.weight, mixer.num_v_heads)
    return mixer


# ------------------------------------------------------------------- n-gram memory


class BranchRMSNorm(nn.Module):
    """Zero-centered RMSNorm over the last dim with one gain vector per branch."""

    def __init__(self, branches: int, hidden_size: int, eps: float) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(branches, hidden_size))
        self.eps = eps
        mark_gain(self.weight)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        values = hidden_states.float()
        values = values * torch.rsqrt(values.square().mean(dim=-1, keepdim=True) + self.eps)
        return values * (1.0 + self.weight.float())


class NGramMemory(nn.Module):
    """Engram-style memory: hashed n-gram rows, gated per residual branch, then a short conv.

    ``e`` concatenates one row per table. Each branch gates the shared value
    ``W_V e`` with ``σ(RMSNorm(R_m) · RMSNorm(W_K,m e) / √d)``; the gated values
    pass a depthwise causal convolution (dilation = largest order, zero init)
    and are added to their branches.
    """

    def __init__(self, config: TinyGDNConfig) -> None:
        super().__init__()
        self.layout = NGramLayout.from_config(config)
        self.branches = config.residual_branches
        self.hidden_size = config.hidden_size
        self.kernel = config.ngram_conv_kernel
        self.dilation = self.layout.max_order
        memory_dim = self.layout.memory_dim
        self.value_proj = nn.Linear(memory_dim, config.hidden_size, bias=False)
        self.key_proj = nn.Linear(memory_dim, self.branches * config.hidden_size, bias=False)
        self.hidden_norm = BranchRMSNorm(self.branches, config.hidden_size, config.rms_norm_eps)
        self.key_norm = BranchRMSNorm(self.branches, config.hidden_size, config.rms_norm_eps)
        self.conv_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.conv_weight = nn.Parameter(torch.zeros(config.hidden_size, self.kernel))
        mark_muon_row_blocks(self.key_proj.weight, self.branches)
        mark_gain(self.conv_weight)
        self.register_buffer("canonical_ids", torch.arange(config.vocab_size, dtype=torch.int64), persistent=False)
        self.register_buffer("tables", None, persistent=False)

    @property
    def conv_history(self) -> int:
        return (self.kernel - 1) * self.dilation

    def set_canonical_ids(self, canonical: torch.Tensor) -> None:
        if canonical.shape != self.canonical_ids.shape:
            raise ValueError("canonical map does not match the vocabulary")
        self.canonical_ids = canonical.to(device=self.canonical_ids.device, dtype=torch.int64)

    def attach_tables(self, tables: torch.Tensor, canonical: torch.Tensor) -> None:
        """Look rows up in ``tables [total_rows, embedding_dim]``.

        The tables may stay in host memory while the model runs on the GPU:
        decoding reads only ``table_count`` rows per token.
        """

        if tables.shape != (self.layout.total_rows, self.layout.embedding_dim):
            raise ValueError(f"tables have shape {tuple(tables.shape)}")
        self.set_canonical_ids(canonical)
        self.tables = tables

    def load_tables(self, path: Path, *, device: torch.device | str, dtype: torch.dtype) -> None:
        tables, canonical = load_table_file(path, self.layout, device=device, dtype=dtype)
        self.attach_tables(tables, canonical)

    def lookup(
        self,
        input_ids: torch.Tensor,
        *,
        distance: torch.Tensor | None = None,
        history: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Concatenated rows ``[..., memory_dim]`` for every position of ``input_ids``."""

        if self.tables is None:
            raise RuntimeError(
                "Senryu-2 needs n-gram embeddings: pass ngram_embeddings or attach tables "
                "with model.ngram.load_tables(...)"
            )
        canonical = self.canonical_ids[input_ids]
        indices = ngram_row_indices(canonical, self.layout, distance=distance, history=history)
        tables = self.tables
        if tables.device == indices.device:
            return F.embedding(indices, tables).flatten(-2)
        rows = F.embedding(indices.to(tables.device), tables).flatten(-2)
        return rows.to(indices.device, non_blocking=True)

    def _convolve(
        self,
        values: torch.Tensor,
        distance: torch.Tensor | None,
        history: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Causal dilated depthwise conv over ``[B, T, n, d]``; returns ``(out, new_history)``."""

        batch, length = values.shape[:2]
        span = self.conv_history
        if history is None:
            history = values.new_zeros(batch, span, self.branches, self.hidden_size)
        extended = torch.cat([history, values], dim=1)
        weight = self.conv_weight.to(values.dtype)
        out = values * weight[:, -1]
        for tap in range(1, self.kernel):
            shift = tap * self.dilation
            shifted = extended[:, span - shift : span - shift + length]
            if distance is not None:
                shifted = shifted * (distance >= shift).to(values.dtype)[..., None, None]
            out = out + shifted * weight[:, -1 - tap]
        return out, extended[:, -span:]

    def forward(
        self,
        residual: torch.Tensor,
        embeddings: torch.Tensor,
        *,
        batch: int,
        distance: torch.Tensor | None,
        conv_state: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Branch updates ``[N, n, d]`` for branches ``[N, n, d]`` and rows ``[N, memory_dim]``."""

        rows = residual.shape[0]
        value = self.value_proj(embeddings)
        keys = self.key_proj(embeddings).view(rows, self.branches, self.hidden_size)
        similarity = (self.hidden_norm(residual) * self.key_norm(keys)).sum(dim=-1)
        alpha = torch.sigmoid(similarity / math.sqrt(self.hidden_size)).to(value.dtype)
        gated = alpha.unsqueeze(-1) * value.unsqueeze(1)
        normalized = self.conv_norm(gated).view(batch, rows // batch, self.branches, self.hidden_size)
        convolved, new_state = self._convolve(
            normalized,
            None if distance is None else distance.view(batch, rows // batch),
            conv_state,
        )
        return F.silu(convolved).view(rows, self.branches, self.hidden_size) + gated, new_state


# ------------------------------------------------------------------------- blocks


class Senryu2Block(nn.Module):
    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
        super().__init__()
        self.layer_index = layer_index
        self.layer_type = config.layer_types[layer_index]
        residual = (config.residual_branches, config.hidden_size, config.residual_gate_rank, config.rms_norm_eps)
        self.mixer_residual = GatedResidual(*residual, write=True)
        self.mlp_residual = GatedResidual(*residual, write=True)
        if self.layer_type == "gdn2":
            self.token_mixer: nn.Module = build_gated_deltanet2(config, layer_index)
        elif self.layer_type == "parallax":
            self.token_mixer = ParallaxAttention(config, layer_index)
        else:
            raise ValueError(f"Senryu-2 has no {self.layer_type} layers")
        self.mlp = SiTUGLU(config)
        mark_muon_row_blocks(self.mlp.gate_up_proj.weight, 2)

    def _mix(self, hidden_states: torch.Tensor, context: Senryu2Context) -> torch.Tensor:
        batched = hidden_states.view(context.batch, context.length, -1)
        if self.layer_type == "parallax":
            mixed = self.token_mixer(batched, context)
        else:
            cache = context.cache
            mixed, _, _ = self.token_mixer(
                batched,
                past_key_values=None if cache is None else cache.linear,
                use_cache=cache is not None,
                cu_seqlens=context.cu_seqlens,
            )
        return mixed.reshape(hidden_states.shape)

    def forward(self, residual: torch.Tensor, context: Senryu2Context) -> torch.Tensor:
        hidden_states, scales = self.mixer_residual.read(residual)
        residual = GatedResidual.write(residual, scales, self._mix(hidden_states, context))
        hidden_states, scales = self.mlp_residual.read(residual)
        return GatedResidual.write(residual, scales, self.mlp(hidden_states))


# -------------------------------------------------------------------------- cache


@dataclass
class Senryu2Cache:
    """Decode state: GDN-2 states, Parallax keys/values, and the n-gram tails."""

    linear: Any = field(default_factory=FlaCache)
    keys: dict[int, torch.Tensor] = field(default_factory=dict)
    values: dict[int, torch.Tensor] = field(default_factory=dict)
    length: int = 0
    ngram_tokens: torch.Tensor | None = None
    ngram_conv: torch.Tensor | None = None

    def append_attention(
        self,
        layer_index: int,
        keys: torch.Tensor,
        values: torch.Tensor,
        start: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Write ``[B, T, H, D]`` keys/values at ``start``; return views of the whole prefix."""

        stop = start + keys.shape[1]
        stored_keys = self.keys.get(layer_index)
        if stored_keys is None or stored_keys.shape[1] < stop:
            capacity = max(stop, 0 if stored_keys is None else 2 * stored_keys.shape[1])
            capacity = -(-capacity // KV_CACHE_GRANULE) * KV_CACHE_GRANULE
            grown_keys = keys.new_empty(keys.shape[0], capacity, *keys.shape[2:])
            grown_values = values.new_empty(values.shape[0], capacity, *values.shape[2:])
            if stored_keys is not None:
                grown_keys[:, :start] = stored_keys[:, :start]
                grown_values[:, :start] = self.values[layer_index][:, :start]
            self.keys[layer_index] = stored_keys = grown_keys
            self.values[layer_index] = grown_values
        stored_values = self.values[layer_index]
        stored_keys[:, start:stop] = keys
        stored_values[:, start:stop] = values
        return stored_keys[:, :stop], stored_values[:, :stop]


# -------------------------------------------------------------------------- model


@dataclass
class _Packing:
    """How a ``[B, T]`` batch maps onto one packed row of ``valid`` tokens."""

    batch: int
    length: int
    index: torch.Tensor | None
    row_starts: torch.Tensor

    def pack(self, tensor: torch.Tensor) -> torch.Tensor:
        flat = tensor.reshape(self.batch * self.length, *tensor.shape[2:])
        return flat if self.index is None else flat.index_select(0, self.index)

    def unpack(self, tensor: torch.Tensor) -> torch.Tensor:
        if self.index is None:
            return tensor.view(self.batch, self.length, *tensor.shape[1:])
        out = tensor.new_zeros(self.batch * self.length, *tensor.shape[1:])
        out.index_copy_(0, self.index, tensor)
        return out.view(self.batch, self.length, *tensor.shape[1:])


class Senryu2ForCausalLM(nn.Module):
    def __init__(self, config: TinyGDNConfig) -> None:
        super().__init__()
        if config.residual_mode != "gated_residual":
            raise ValueError("Senryu2ForCausalLM needs residual_mode 'gated_residual'")
        self.config = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(Senryu2Block(config, index) for index in range(config.num_hidden_layers))
        self.final_residual = GatedResidual(
            config.residual_branches,
            config.hidden_size,
            config.residual_gate_rank,
            config.rms_norm_eps,
            write=False,
        )
        self.ngram = NGramMemory(config) if config.ngram_orders else None
        self.rotary = RotaryTables(config.rotary_dim, config.rope_theta)
        self.checkpointed_layers: frozenset[int] = frozenset()
        self.apply(self._initialize_module)
        self._initialize_residual_projections()

    # ---------------------------------------------------------------- setup
    def _initialize_module(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)

    def _initialize_residual_projections(self) -> None:
        residual_std = self.config.initializer_range / math.sqrt(2 * self.config.num_hidden_layers)
        for layer in self.layers:
            nn.init.normal_(layer.token_mixer.o_proj.weight, mean=0.0, std=residual_std)
            nn.init.normal_(layer.mlp.down_proj.weight, mean=0.0, std=residual_std)

    def enable_gradient_checkpointing(self, enabled: bool = True, *, layers: int | None = None) -> None:
        """Recompute block activations in backward; ``layers`` picks GDN-2 blocks first."""

        if not enabled:
            if layers is not None:
                raise ValueError("layers needs gradient checkpointing enabled")
            self.checkpointed_layers = frozenset()
            return
        count = len(self.layers) if layers is None else layers
        if not 1 <= count <= len(self.layers):
            raise ValueError(f"layers must be between 1 and {len(self.layers)}")
        cheapest_replay_first = sorted(
            range(len(self.layers)),
            key=lambda index: self.layers[index].layer_type != "gdn2",
        )
        self.checkpointed_layers = frozenset(cheapest_replay_first[:count])

    def describe_gradient_checkpointing(self) -> str:
        if not self.checkpointed_layers:
            return "off"
        layer_types = sorted({self.layers[index].layer_type for index in self.checkpointed_layers})
        return f"{len(self.checkpointed_layers)}/{len(self.layers)} blocks ({', '.join(layer_types)})"

    def project_to_vocabulary(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return F.linear(hidden_states, self.embed_tokens.weight)

    # -------------------------------------------------------------- packing
    def _packing(self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None) -> _Packing:
        batch, length = input_ids.shape
        device = input_ids.device
        if attention_mask is None or bool(attention_mask.all()):
            row_starts = torch.arange(0, batch * length, length, device=device)
            return _Packing(batch, length, None, row_starts)
        if attention_mask.shape != input_ids.shape:
            raise ValueError("attention_mask must match input_ids")
        valid = attention_mask.reshape(-1).bool()
        index = valid.nonzero().flatten()
        lengths = attention_mask.bool().sum(dim=1)
        if bool((lengths == 0).any()):
            raise ValueError("every row needs at least one unmasked token")
        row_starts = torch.cumsum(lengths, dim=0) - lengths
        return _Packing(batch, length, index, row_starts)

    def packed_layout(
        self,
        tokens: torch.Tensor,
        row_starts: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """``(cu_seqlens int32, distance int64)`` for a flat token stream split at BOS."""

        distance = document_distance(tokens, self.config.bos_token_id, row_starts)
        starts = (distance == 0).nonzero().flatten()
        cu_seqlens = torch.cat([starts, starts.new_tensor([tokens.numel()])]).to(torch.int32)
        return cu_seqlens, distance

    # -------------------------------------------------------------- forward
    def _run_block(self, index: int, residual: torch.Tensor, context: Senryu2Context) -> torch.Tensor:
        layer = self.layers[index]
        if self.training and index in self.checkpointed_layers:
            return checkpoint(layer, residual, context, use_reentrant=False)
        return layer(residual, context)

    def _inject_ngrams(
        self,
        residual: torch.Tensor,
        embeddings: torch.Tensor,
        context: Senryu2Context,
    ) -> torch.Tensor:
        assert self.ngram is not None
        cache = context.cache
        conv_state = None if cache is None else cache.ngram_conv
        distance = context.distance if cache is None else None
        if self.training and self.checkpointed_layers:
            update, new_state = checkpoint(
                self.ngram,
                residual,
                embeddings,
                batch=context.batch,
                distance=distance,
                conv_state=conv_state,
                use_reentrant=False,
            )
        else:
            update, new_state = self.ngram(
                residual,
                embeddings,
                batch=context.batch,
                distance=distance,
                conv_state=conv_state,
            )
        if cache is not None:
            cache.ngram_conv = new_state
        return residual + update

    def _ngram_rows(
        self,
        input_ids: torch.Tensor,
        packing: _Packing | None,
        distance: torch.Tensor | None,
        ngram_embeddings: torch.Tensor | None,
        cache: Senryu2Cache | None,
    ) -> torch.Tensor:
        assert self.ngram is not None
        layout = self.ngram.layout
        if cache is not None:
            history = cache.ngram_tokens
            if history is None:
                history = input_ids.new_full((input_ids.shape[0], layout.max_order - 1), layout.pad_token_id)
            canonical = self.ngram.canonical_ids[input_ids]
            cache.ngram_tokens = torch.cat([history, canonical], dim=1)[:, -(layout.max_order - 1) :]
            if ngram_embeddings is not None:
                return ngram_embeddings.reshape(-1, layout.memory_dim)
            return self.ngram.lookup(input_ids, history=history).reshape(-1, layout.memory_dim)
        assert packing is not None
        if ngram_embeddings is not None:
            if ngram_embeddings.shape != (*input_ids.shape, layout.memory_dim):
                raise ValueError(
                    f"ngram_embeddings must have shape {(*input_ids.shape, layout.memory_dim)}"
                )
            return packing.pack(ngram_embeddings)
        return self.ngram.lookup(packing.pack(input_ids), distance=distance)

    def forward(
        self,
        input_ids: torch.Tensor,
        labels: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        past_key_values: Senryu2Cache | None = None,
        *,
        use_cache: bool = False,
        return_logits: bool = True,
        return_hidden_states: bool = False,
        labels_are_shifted: bool = False,
        include_mtp_loss: bool = True,
        mtp_loss_weight: float | None = None,
        z_loss_coefficient: float = 0.0,
        logits_to_keep: int | None = None,
        ngram_embeddings: torch.Tensor | None = None,
    ) -> TinyGDNOutput:
        del include_mtp_loss, mtp_loss_weight
        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, sequence]")
        if labels is not None and labels.shape != input_ids.shape:
            raise ValueError("labels must have the same shape as input_ids")
        if z_loss_coefficient < 0.0:
            raise ValueError("z_loss_coefficient cannot be negative")
        if logits_to_keep is not None and logits_to_keep <= 0:
            raise ValueError("logits_to_keep must be positive")
        if use_cache and labels is not None:
            raise ValueError("use_cache is not supported with labels")
        if past_key_values is not None and not isinstance(past_key_values, Senryu2Cache):
            raise TypeError("past_key_values must be a Senryu2Cache or None")
        cached = use_cache or past_key_values is not None
        batch, length = input_ids.shape
        offset = 0 if past_key_values is None else past_key_values.length
        if offset + length > self.config.max_position_embeddings:
            raise ValueError("Input exceeds max_position_embeddings")

        embeddings = self.embed_tokens(input_ids)
        device, dtype = embeddings.device, embeddings.dtype
        packing: _Packing | None = None
        distance: torch.Tensor | None = None
        if cached:
            if attention_mask is not None and not bool(attention_mask.all()):
                raise ValueError("Senryu-2 decodes unpadded rows; drop the padding instead of masking it")
            cache = past_key_values if past_key_values is not None else Senryu2Cache()
            context = Senryu2Context(
                batch=batch,
                length=length,
                cu_seqlens=None,
                distance=None,
                rotary=self.rotary.get(offset + length, device, dtype),
                position_offset=offset,
                cache=cache,
            )
            hidden = embeddings.reshape(batch * length, -1)
        else:
            cache = None
            packing = self._packing(input_ids, attention_mask)
            tokens = packing.pack(input_ids)
            cu_seqlens, distance = self.packed_layout(tokens, packing.row_starts)
            context = Senryu2Context(
                batch=1,
                length=tokens.numel(),
                cu_seqlens=cu_seqlens,
                distance=distance,
                rotary=self.rotary.get(tokens.numel(), device, dtype),
            )
            hidden = packing.pack(embeddings)
        residual = hidden.unsqueeze(1).expand(-1, self.config.residual_branches, -1)

        for index in range(len(self.layers)):
            if self.ngram is not None and index == self.config.ngram_layer_index:
                rows = self._ngram_rows(input_ids, packing, distance, ngram_embeddings, cache)
                residual = self._inject_ngrams(residual, rows.to(dtype), context)
            residual = self._run_block(index, residual, context)
        hidden_states, _ = self.final_residual.read(residual)
        if cache is not None:
            cache.length = offset + length

        main_loss = None
        z_loss = None
        total_loss = None
        if labels is not None:
            assert packing is not None
            if labels_are_shifted:
                targets = labels
            else:
                targets = torch.cat([labels[:, 1:], labels.new_full((batch, 1), -100)], dim=1)
            total_loss, main_loss, z_loss = linear_cross_entropy(
                hidden_states,
                self.embed_tokens.weight,
                packing.pack(targets).reshape(-1),
                z_loss_coefficient=z_loss_coefficient,
            )
            if z_loss_coefficient == 0.0:
                z_loss = None

        batched_states = (
            hidden_states.view(batch, length, -1) if packing is None else packing.unpack(hidden_states)
        )
        output_states = batched_states if logits_to_keep is None else batched_states[:, -logits_to_keep:]
        logits = self.project_to_vocabulary(output_states) if return_logits else None
        return TinyGDNOutput(
            loss=total_loss,
            logits=logits,
            main_loss=main_loss,
            mtp_loss=None,
            z_loss=z_loss,
            hidden_states=batched_states if return_hidden_states else None,
            past_key_values=cache if use_cache else None,
        )

    # ------------------------------------------------------------ reporting
    def parameter_report(self) -> dict[str, int]:
        total = sum(parameter.numel() for parameter in self.parameters())
        embeddings = self.embed_tokens.weight.numel()
        tables = 0 if self.ngram is None else self.ngram.layout.total_rows * self.ngram.layout.embedding_dim
        return {
            "deployable_core": total,
            "training_total": total,
            "embedding": embeddings,
            "mtp_auxiliary": 0,
            "non_embedding_core": total - embeddings,
            "ngram_table": tables,
        }

    def save_checkpoint(self, output_dir: Path) -> None:
        output_dir.mkdir(parents=True, exist_ok=True)
        self.config.save_json(output_dir / "config.json")
        save_model(self, output_dir / "model.safetensors")

    def load_ngram_tables(
        self,
        checkpoint_dir: Path,
        *,
        weights: str,
        device: str | torch.device,
        dtype: torch.dtype,
    ) -> None:
        """Attach the checkpoint's tables that match ``weights`` ("normal" or "ema")."""

        if self.ngram is None:
            return
        table_file = checkpoint_dir / (NGRAM_TABLES_EMA_FILE if weights == "ema" else NGRAM_TABLES_FILE)
        if not table_file.is_file():
            raise FileNotFoundError(f"Senryu-2 checkpoint has no n-gram tables: {table_file}")
        self.ngram.load_tables(table_file, device=device, dtype=dtype)

    def extra_repr(self) -> str:
        report = self.parameter_report()
        return f"core_parameters={report['deployable_core']:,}, ngram_table_parameters={report['ngram_table']:,}"

    def get_architecture_metadata(self) -> dict[str, Any]:
        return {
            "architecture": self.config.architecture,
            "layer_types": list(self.config.layer_types),
            "effective_num_layers": self.config.effective_num_layers,
            "shared_layer_indices": [],
            "parameter_report": self.parameter_report(),
            "features": [
                "3:1 Gated DeltaNet-2 + Parallax attention hybrid",
                "4-branch Gated Residual stream (gated read, per-branch write scales)",
                "hashed n-gram memory (orders 2-3) with per-branch context gating",
                "full-rank sigmoid output gates",
                "partial RoPE, per-head QK/R normalization",
                "SiTU-GLU",
                "document-masked packed training",
                "tied input-output embeddings",
            ],
        }