delta / modeling_tamelm.py
Upload TAMELM
f2c3f04 verified about 1 month ago
74.4 kB
#@title Model — ready
"""TAMELM model definitions: routed attention, depth attention, and dropless MoE.
This module stays self-contained for notebook execution and portable Hugging Face
exports. Definitions are ordered from configuration and low-level operations to
model components, cache/output interfaces, and the complete language model.
Routing uses shared GQA projections and original top-k softmax probabilities.
Tile routing and graph selection are causal; cached decoding retains tile state.
CUDA FlexAttention and Triton grouped GEMM are optional acceleration backends.
"""

from __future__ import annotations

import ast as _ast
import inspect as _inspect
import linecache as _linecache
import math
from dataclasses import dataclass
from typing import Dict, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    GenerationMixin,
    PretrainedConfig,
    PreTrainedModel,
)
from transformers.cache_utils import Cache
from transformers.utils import ModelOutput

try:
    from torch.nn.attention.flex_attention import (
        BlockMask,
        create_block_mask,
        flex_attention,
    )

    HAS_FLEX = True
except ImportError:  # pragma: no cover
    HAS_FLEX = False

try:
    import triton  # type: ignore
    import triton.language as tl  # type: ignore

    HAS_TRITON = True
except ImportError:
    HAS_TRITON = False


# Notebook source capture and semantic fingerprints.
# Capture source independently of execution counts, paths, and timestamps.
_NOTEBOOK_SOURCES = globals().get("_NOTEBOOK_SOURCES", {})


# =============================================================================
# Source identity helpers
# =============================================================================


def _capture_cell_source(component):
    caller = _inspect.currentframe().f_back
    source = "".join(_linecache.getlines(caller.f_code.co_filename))
    if not source:
        try:
            source = get_ipython().history_manager.input_hist_raw[-1]
        except (NameError, AttributeError, IndexError):
            source = ""
    if not source.strip():
        raise RuntimeError("Cannot capture executed source; run the definition cells in Jupyter/Colab.")
    _NOTEBOOK_SOURCES[component] = source


class _WithoutDocstrings(_ast.NodeTransformer):
    def _strip(self, node):
        self.generic_visit(node)
        if (
            node.body
            and isinstance(node.body[0], _ast.Expr)
            and isinstance(node.body[0].value, _ast.Constant)
            and isinstance(node.body[0].value.value, str)
        ):
            node.body = node.body[1:] or [_ast.Pass()]
        return node

    visit_Module = _strip
    visit_ClassDef = _strip
    visit_FunctionDef = _strip
    visit_AsyncFunctionDef = _strip


def _semantic_source(source):
    return _ast.dump(
        _WithoutDocstrings().visit(_ast.parse(source)), include_attributes=False
    )


# Routing mechanism IDs are part of the attention/cache contract.
GLOBAL, LOCAL, STRIDED, GRAPH = 0, 1, 2, 3
MECH_NAMES = ("global", "local", "strided", "graph")


# =============================================================================
# Configuration
# =============================================================================


@dataclass
class RoutedAttentionConfig:
    d_model: int
    n_q_heads: int = 16
    n_kv_heads: int = 4
    head_dim: Optional[int] = None
    top_k: int = 2
    tile: int = 128  # routing granularity == FlexAttention block size
    local_window: int = 256  # tokens
    strided_recent: int = 128  # tokens of dense recent context for 'strided'
    strided_every: int = 4  # plus every S-th tile back (tile-relative)
    graph_top_r: int = 0  # past tiles picked by 'graph'; 0 = legacy 3-way router
    graph_recent: int = 128  # tokens of dense recent context for 'graph'
    max_seq_len: int = 2048
    router_init_std: float = 1e-2
    balance_coef: float = 1e-2
    backend: str = "flex"  # "flex" | "dense"; decode/cached paths always dense





# =============================================================================
# Shared routing precision
# =============================================================================


def router_logits_fp32(router, x):
    """Keep routing decisions out of AMP; casting rounded BF16 logits is too late."""
    with torch.autocast(device_type=x.device.type, enabled=False):
        return F.linear(
            x.float(),
            router.weight.float(),
            None if router.bias is None else router.bias.float(),
        )


# =============================================================================
# Position transforms
# =============================================================================


class RotaryPos(nn.Module):
    """Standard RoPE with interleaved pairs (same pairing convention as BlackHoleRoPE)."""

    def __init__(self, head_dim: int, base: float = 10_000.0):
        super().__init__()
        half = head_dim // 2
        self.register_buffer(
            "inv", 1.0 / base ** (torch.arange(half).float() / half), persistent=False
        )

    def forward(self, x, start_pos: int = 0, **_):  # x: [B, H, T, dh]
        T = x.shape[2]
        t = torch.arange(start_pos, start_pos + T, device=x.device, dtype=torch.float32)
        ang = t[:, None] * self.inv[None]
        cos, sin = ang.cos(), ang.sin()
        xf = x.float()
        x1, x2 = xf[..., 0::2], xf[..., 1::2]
        return (
            torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], -1)
            .flatten(-2)
            .to(x.dtype)
        )


class WarpedRoPE(nn.Module):
    """RoPE on a content-defined clock, one clock per KV group.
        tau_t = softplus(w . x_t + b) > 0        clock rate of token t
        T_t   = sum_{s<t} tau_s                  token time (exclusive prefix sum)
        angle = omega_f * T_t
    q_m . k_n then depends on T_m - T_n = sum_{n<=s<m} tau_s: only the tokens between
    them, so attention is relative in warped time and independent of where a document
    sits in a packed window. Causal, cumsum-parallel, O(1) decode state (the running
    time). At init (w = 0, softplus(b) = 1) it is exactly standard RoPE.
    """

    def __init__(
        self,
        d_model: int,
        head_dim: int,
        n_q_heads: int,
        n_kv_heads: int,
        base: float = 10_000.0,
    ):
        super().__init__()
        if head_dim % 2 or n_q_heads % n_kv_heads:
            raise ValueError(
                "WarpedRoPE needs an even head_dim and grouped query heads"
            )
        self.half, self.base = head_dim // 2, float(base)
        self.group = n_q_heads // n_kv_heads
        self.rate = nn.Linear(d_model, n_kv_heads)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.zeros_(self.rate.weight)
        nn.init.constant_(self.rate.bias, math.log(math.e - 1.0))  # softplus(b) = 1

    def clock(self, x, t0=None):
        """x [B,T,d] (normed). t0 [B,Hkv] time already elapsed, or None.
        Returns token times [B,T,Hkv] and inclusive times [B,T,Hkv] (cache state)."""
        tau = F.softplus(router_logits_fp32(self.rate, x))  # fp32
        inclusive = tau.cumsum(1)
        if t0 is not None:
            inclusive = inclusive + t0[:, None]
        return inclusive - tau, inclusive

    def rotate(self, x, times):  # x [B,h,T,dh]; times [B,h,T]
        # Frequencies are recomputed rather than buffered: nothing here can be left
        # uninitialized by meta-device / no_init_weights loading.
        inv = 1.0 / self.base ** (
            torch.arange(self.half, device=x.device, dtype=torch.float32) / self.half
        )
        ang = times[..., None] * inv
        cos, sin = ang.cos(), ang.sin()
        xf = x.float()
        x1, x2 = xf[..., 0::2], xf[..., 1::2]
        return (
            torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], -1)
            .flatten(-2)
            .to(x.dtype)
        )


class BlackHoleRoPE(nn.Module):
    def __init__(
        self,
        dim,
        max_seq_len=2048,
        base_vortex_strength=1.0,
        energy_amplification=0.1,
        theta_base=10_000.0,
        adaptive_vortex=True,
        discrepancy_modulation=True,
    ):
        super().__init__()
        if dim <= 0 or dim % 2 or max_seq_len < 1:
            raise ValueError(
                "BlackHoleRoPE requires an even dimension and positive max_seq_len"
            )
        self.dim = int(dim)
        self.half_dim = dim // 2
        self.max_seq_len = int(max_seq_len)
        self.energy_amplification = float(energy_amplification)
        self.discrepancy_modulation = bool(discrepancy_modulation)
        freqs = 1.0 / (
            theta_base
            ** (torch.arange(0, self.half_dim, dtype=torch.float32) / self.half_dim)
        )
        self.register_buffer("base_freqs", freqs)
        self.vortex_strength = nn.Parameter(
            torch.full((max_seq_len,), base_vortex_strength),
            requires_grad=adaptive_vortex,
        )
        self.phase_shift = nn.Parameter(torch.randn(max_seq_len, self.half_dim) * 0.01)
        self.radial_warp = nn.Parameter(torch.ones(max_seq_len, self.half_dim))
        if discrepancy_modulation:
            self.discrepancy_projector = nn.Sequential(
                nn.Linear(dim, self.half_dim),
                nn.Tanh(),
                nn.Linear(self.half_dim, self.half_dim),
            )
        self.register_buffer("energy_clamp_min", torch.tensor(0.1))
        self.register_buffer("energy_clamp_max", torch.tensor(10.0))

    def _compute_positional_vortex(self, seq_len, device, dtype, start_pos=0):
        positions = torch.arange(
            start_pos, start_pos + seq_len, dtype=torch.float32, device=device
        )
        position_freqs = positions.unsqueeze(1) * self.base_freqs.unsqueeze(0)
        base_phase = position_freqs + self.phase_shift[start_pos : start_pos + seq_len]
        vortex_phase = self.vortex_strength[start_pos : start_pos + seq_len].unsqueeze(
            1
        ) * torch.sin(base_phase + position_freqs * 0.1)
        final_phase = base_phase + vortex_phase
        rotation = torch.complex(torch.cos(final_phase), torch.sin(final_phase))
        radial_factor = torch.sqrt(position_freqs + 1e-8)
        energy = 1.0 + self.energy_amplification * torch.sin(
            self.radial_warp[start_pos : start_pos + seq_len] * radial_factor
            + vortex_phase.abs()
        )
        energy = energy.clamp(self.energy_clamp_min, self.energy_clamp_max)
        return rotation, energy.to(dtype=dtype)

    def apply_discrepancy_modulation(self, x, symplectic_state):
        if not self.discrepancy_modulation or not symplectic_state:
            return None
        bsz, seq_len, _ = x.shape
        total = torch.zeros(bsz, seq_len, 1, device=x.device, dtype=x.dtype)
        count = 0
        for theta, p in symplectic_state.values():
            if theta is None or p is None:
                continue
            ev = theta.float().var(dim=-1, keepdim=True, correction=0) + p.float().var(
                dim=-1, keepdim=True, correction=0
            )
            ev = (
                ev.unsqueeze(1).expand(-1, seq_len, -1)
                if theta.dim() == 2
                else ev[:, :seq_len]
            )
            total = total + ev.to(dtype=x.dtype)
            count += 1
        if count == 0:
            return None
        lifted = F.pad(total / count, (0, self.dim - 1))
        return self.discrepancy_projector(lifted)

    def forward(self, x, symplectic_state=None, start_pos=0):
        bsz, seq_len, dim = x.shape
        if dim != self.dim or start_pos < 0 or start_pos + seq_len > self.max_seq_len:
            raise ValueError(
                "BlackHoleRoPE dimension/position exceeds configured bounds"
            )
        input_dtype = x.dtype
        xp = x.reshape(bsz, seq_len, self.half_dim, 2)
        xc = torch.complex(xp[..., 0].float(), xp[..., 1].float())
        rotation, energy = self._compute_positional_vortex(
            seq_len, x.device, torch.float32, start_pos
        )
        rotation = rotation.unsqueeze(0)
        energy = energy.unsqueeze(0)
        disc = self.apply_discrepancy_modulation(x, symplectic_state)
        if disc is not None:
            phase_scale = 1.0 + 0.1 * torch.tanh(disc.float())
            base_phase = torch.angle(rotation)
            rotation = torch.polar(
                torch.ones_like(base_phase), base_phase * phase_scale
            )
        amp = xc * rotation * energy.float()
        return (
            torch.stack([amp.real, amp.imag], dim=-1)
            .reshape(bsz, seq_len, dim)
            .to(input_dtype)
        )


class BlackHolePos(nn.Module):
    """Adapter: folds heads into batch for BlackHoleRoPE and repeats the symplectic
    state per head so discrepancy modulation keeps the right batch alignment."""

    def __init__(self, bh: nn.Module):
        super().__init__()
        self.bh = bh

    def forward(self, x, start_pos: int = 0, symplectic_state=None):
        B, H, T, dh = x.shape
        state = None
        if symplectic_state:
            state = {
                n: (th.repeat_interleave(H, 0), p.repeat_interleave(H, 0))
                for n, (th, p) in symplectic_state.items()
                if th is not None and p is not None
            }
        y = self.bh(
            x.reshape(B * H, T, dh), symplectic_state=state, start_pos=start_pos
        )
        return y.view(B, H, T, dh)


# FlexAttention compilation is lazy and shared by CUDA attention calls.
_COMPILED_FLEX = None


# =============================================================================
# Attention backends and visibility geometry
# =============================================================================


def _run_flex(q, k, v, block_mask):
    global _COMPILED_FLEX
    if q.is_cuda:
        if min(q.shape[-1], k.shape[-1], v.shape[-1]) < 16:
            raise ValueError(
                "CUDA FlexAttention requires head_dim >= 16; use a supported GPU model configuration"
            )
        if _COMPILED_FLEX is None:
            _COMPILED_FLEX = torch.compile(flex_attention, dynamic=False)
        return _COMPILED_FLEX(q, k, v, block_mask=block_mask, enable_gqa=True)
    return flex_attention(q, k, v, block_mask=block_mask, enable_gqa=True)


def geometry_allows(mech, q_idx, kv_idx, cfg: RoutedAttentionConfig, picked=None):
    """Boolean key-visibility for a given mechanism id. Pure tensor algebra so the
    same function serves FlexAttention's mask_mod and the dense reference mask.
    `picked` is True where kv_idx lies in a graph-selected tile of q_idx's tile."""
    d = q_idx - kv_idx
    causal = d >= 0
    local = d < cfg.local_window
    tile_gap = torch.div(q_idx, cfg.tile, rounding_mode="floor") - torch.div(
        kv_idx, cfg.tile, rounding_mode="floor"
    )
    strided = (d < cfg.strided_recent) | (
        torch.remainder(tile_gap, cfg.strided_every) == 0
    )
    allowed = (
        (mech == GLOBAL) | ((mech == LOCAL) & local) | ((mech == STRIDED) & strided)
    )
    if cfg.graph_top_r:
        graph = d < cfg.graph_recent
        if picked is not None:
            graph = graph | picked
        allowed = allowed | ((mech == GRAPH) & graph)
    return causal & allowed


def geometry_block_mask(mech, length, cfg, mask_mod, picked=None):
    """Exact block classification in O(B * H * (T/tile)^2), never O(B*H*T^2).
    A partial block uses the original elementwise predicate. Fully allowed blocks
    skip that predicate. No attention probabilities or causal rules are approximated.
    `picked` is the tile-level graph selection [B, H, n, n] (query tile, kv tile).
    """
    c = cfg.tile
    n = mech.shape[-1]
    starts = torch.arange(n, device=mech.device) * c
    ends = (starts + c - 1).clamp(max=length - 1)
    dmin = starts[:, None] - ends[None, :]
    dmax = ends[:, None] - starts[None, :]
    gap = (
        torch.arange(n, device=mech.device)[:, None]
        - torch.arange(n, device=mech.device)[None, :]
    )
    periodic = gap.remainder(cfg.strided_every) == 0
    mechanism = mech[..., None]
    any_geometry = (
        (mechanism == GLOBAL)
        | ((mechanism == LOCAL) & (dmin < cfg.local_window))
        | ((mechanism == STRIDED) & ((dmin < cfg.strided_recent) | periodic))
    )
    full_geometry = (
        (mechanism == GLOBAL)
        | ((mechanism == LOCAL) & (dmax < cfg.local_window))
        | ((mechanism == STRIDED) & ((dmax < cfg.strided_recent) | periodic))
    )
    if cfg.graph_top_r:
        g_any, g_full = dmin < cfg.graph_recent, dmax < cfg.graph_recent
        if picked is not None:
            # A picked tile is a strictly earlier, complete tile: every pair in the
            # block is allowed, so it is safe for the full-block path.
            g_any, g_full = g_any | picked, g_full | picked
        any_geometry = any_geometry | ((mechanism == GRAPH) & g_any)
        full_geometry = full_geometry | ((mechanism == GRAPH) & g_full)
    any_block = (dmax >= 0) & any_geometry
    # Out-of-range positions in tail tiles are masked. Such blocks must not
    # enter the full-block path, which skips mask_mod in the CUDA kernel.
    complete_tiles = starts + c <= length
    full = (
        (dmin >= 0) & full_geometry & complete_tiles[:, None] & complete_tiles[None, :]
    )
    partial = any_block & ~full

    def encode(mask):
        counts = mask.sum(-1).to(
            dtype=torch.int32, memory_format=torch.contiguous_format
        )
        # Router permutes tile/head axes; argsort preserves that nonstandard
        # layout. CUDA Flex flattens B/H using the H stride and requires
        # contiguous metadata (the CPU eager path does not enforce this).
        indices = (
            mask.to(torch.int32)
            .argsort(dim=-1, descending=True, stable=True)
            .to(dtype=torch.int32, memory_format=torch.contiguous_format)
        )
        return counts, indices

    pnum, pidx = encode(partial)
    fnum, fidx = encode(full)
    return BlockMask.from_kv_blocks(
        pnum,
        pidx,
        fnum,
        fidx,
        BLOCK_SIZE=c,
        mask_mod=mask_mod,
        seq_lengths=(length, length),
    )


# TAMELM routed sequence attention, v2.
#
# One shared GQA projection/cache; a tiny tile-level router picks, per (tile, q-head),
# top-k attention geometries from {global, local, strided, graph}; all geometries run
# through ONE attention call (FlexAttention block mask, or an equivalent dense mask).
#
# Invariants:
#   * K/V are projected, position-encoded and cached BEFORE routing.
#   * Tile routing is causal: tile c is routed, and its graph neighbours are chosen,
#     from its first token only, so training == chunked prefill == token-by-token decode.
#   * Top-k outputs are weighted by the ORIGINAL softmax probabilities (no
#     renormalization), so top-1 still trains the router.
#   * Graph geometry: a query tile sees its recent window plus the top-r COMPLETED
#     past tiles ranked by (first-token query) . (tile-mean key). Ranking uses
#     un-rotated q/k, so it is content similarity rather than position. The tile choice
#     is hard (no gradient); the router weight on 'graph' is what trains. This is the
#     causal, tile-level repair of the learned k-NN adjacency from the graph router.
#
# Positional transform is pluggable: RotaryPos (standard RoPE), WarpedRoPE (RoPE on a
# content-defined clock; relative in warped time), or BlackHolePos (legacy adapter;
# note it depends on absolute position through per-position phase/vortex/energy).


# =============================================================================
# Routed sequence attention
# =============================================================================


class RoutedSequenceAttention(nn.Module):

    def __init__(self, cfg: RoutedAttentionConfig, pos: Optional[nn.Module] = None):
        super().__init__()
        H, Hkv, K = cfg.n_q_heads, cfg.n_kv_heads, cfg.top_k
        dh = cfg.head_dim or cfg.d_model // H
        if cfg.graph_top_r < 0:
            raise ValueError("graph_top_r must be nonnegative")
        self.M = 4 if cfg.graph_top_r else 3  # number of mechanisms
        if (
            min(
                H,
                Hkv,
                dh,
                cfg.tile,
                cfg.local_window,
                cfg.strided_recent,
                cfg.strided_every,
                cfg.graph_recent,
                cfg.max_seq_len,
            )
            <= 0
        ):
            raise ValueError("Attention dimensions/windows must be positive")
        if H % Hkv or dh % 2 or not 1 <= K <= self.M:
            raise ValueError("Invalid GQA heads, odd rotary dimension, or top_k")
        if cfg.backend not in {"flex", "dense"}:
            raise ValueError("backend must be flex or dense")
        self.cfg, self.H, self.Hkv, self.K, self.dh = cfg, H, Hkv, K, dh
        self.graph = bool(cfg.graph_top_r)
        self.n_max = -(-cfg.max_seq_len // cfg.tile)  # tiles in a full window
        self._force_flex = False  # test hook: Flex on CPU (eager)

        self.q_proj = nn.Linear(cfg.d_model, H * dh, bias=False)
        self.k_proj = nn.Linear(cfg.d_model, Hkv * dh, bias=False)
        self.v_proj = nn.Linear(cfg.d_model, Hkv * dh, bias=False)
        self.o_proj = nn.Linear(H * dh, cfg.d_model, bias=False)

        self.router = nn.Linear(cfg.d_model, H * self.M, bias=False)
        nn.init.normal_(self.router.weight, std=cfg.router_init_std)
        # [H, M] mask: this revision has no mandatory anchor heads.
        # The run contract records actual constraints; anchored variants are separate experiments.
        self.register_buffer("allowed", torch.ones(H, self.M, dtype=torch.bool))

        self.pos = pos if pos is not None else RotaryPos(dh)

    # --------------------------------------------------------------------- routing
    def _route(self, x_starts):  # [B, n, d]
        B, n, _ = x_starts.shape
        logits = router_logits_fp32(self.router, x_starts).view(B, n, self.H, self.M)
        logits = logits.masked_fill(~self.allowed, float("-inf"))
        P = logits.softmax(-1)
        w, idx = P.topk(self.K, dim=-1)  # un-renormalized weights
        return P, w, idx

    def _balance(self, P, idx):
        M = self.M
        f = F.one_hot(idx, M).float().sum(-2).reshape(-1, M).mean(0) / self.K
        Pbar = P.reshape(-1, M).mean(0)
        loss = self.cfg.balance_coef * M * (f * Pbar).sum()
        ent = -(P.clamp_min(1e-9).log() * P).sum(-1).mean()
        return loss, f.detach(), ent.detach()

    @torch.no_grad()
    def _select_tiles(self, q_first, k_raw, tiles):
        """q_first [B,H,n_new,dh]: un-rotated query of each new tile's first token.
        k_raw [B,Hkv,S,dh]: every un-rotated key known so far. tiles [n_new].
        Returns [B, n_new, H, n_max] bool: past complete tiles picked per query tile."""
        C, r = self.cfg.tile, self.cfg.graph_top_r
        B, H, n_new, dh = q_first.shape
        sel = torch.zeros(
            B, n_new, H, self.n_max, dtype=torch.bool, device=q_first.device
        )
        n_full = min(k_raw.shape[2] // C, self.n_max)
        if n_full == 0 or n_new == 0:
            return sel
        with torch.autocast(device_type=q_first.device.type, enabled=False):
            kbar = k_raw[:, :, : n_full * C].float().unflatten(2, (n_full, C)).mean(3)
            kbar = kbar.repeat_interleave(H // self.Hkv, dim=1)  # [B,H,n_full,dh]
            scores = torch.einsum("bhnd,bhtd->bnht", q_first.float(), kbar)
        # Tile t is a candidate for query tile c iff t < c (strictly earlier, hence
        # complete by the time tile c's first token exists).
        past = torch.arange(n_full, device=sel.device)[None, :] < tiles[:, None]
        scores = scores.masked_fill(~past[None, :, None, :], float("-inf"))
        vals, idx = scores.topk(min(r, n_full), dim=-1)
        part = torch.zeros(B, n_new, H, n_full, dtype=torch.bool, device=sel.device)
        part.scatter_(-1, idx, torch.isfinite(vals))
        sel[..., :n_full] = part
        return sel

    # --------------------------------------------------------------------- forward
    def forward(
        self, x, cache: Optional[Dict] = None, start_pos: int = 0, symplectic_state=None
    ):
        """x: [B, T, d]. With `cache` (a dict, mutated in place) supports chunked
        prefill and decode; `start_pos` is the absolute position of x[:, 0]."""
        cfg, H, Hkv, K, dh, C = (
            self.cfg,
            self.H,
            self.Hkv,
            self.K,
            self.dh,
            self.cfg.tile,
        )
        B, T, _ = x.shape
        if T < 1 or start_pos < 0:
            raise ValueError("Nonempty input and nonnegative start_pos required")
        if cache is not None and torch.is_grad_enabled():
            raise ValueError("KV caching is inference-only; use torch.no_grad()")
        if cache is None and start_pos != 0:
            raise ValueError("A nonzero start_pos requires the prior KV cache")
        if cache is not None and "k" not in cache and start_pos != 0:
            raise ValueError("An empty cache must begin at position zero")

        q_raw = self.q_proj(x).view(B, T, H, dh).transpose(1, 2)
        k_raw = self.k_proj(x).view(B, T, Hkv, dh).transpose(1, 2)
        v = self.v_proj(x).view(B, T, Hkv, dh).transpose(1, 2)
        if isinstance(self.pos, WarpedRoPE):
            t0 = (
                cache["clock"][:, -1]
                if (cache is not None and "clock" in cache)
                else None
            )
            times, clock = self.pos.clock(x, t0)  # [B, T, Hkv]
            if cache is not None:
                cache["clock"] = (
                    torch.cat([cache["clock"], clock], 1) if "clock" in cache else clock
                )
            times = times.transpose(1, 2)  # [B, Hkv, T]
            q = self.pos.rotate(q_raw, times.repeat_interleave(H // Hkv, dim=1))
            k = self.pos.rotate(k_raw, times)
        else:
            # Compute the positional transform once for the shared Q/K position range.
            qk = self.pos(
                torch.cat((q_raw, k_raw), dim=1),
                start_pos=start_pos,
                symplectic_state=symplectic_state,
            )
            q, k = qk.split((H, Hkv), dim=1)

        # 1) write sequence memory before any routing decision
        k_raw_all = k_raw
        if cache is not None:
            if "k" in cache:
                assert (
                    cache["k"].shape[2] == start_pos
                ), "cache length must equal start_pos"
                k = torch.cat([cache["k"], k], 2)
                v = torch.cat([cache["v"], v], 2)
                if self.graph:
                    k_raw_all = torch.cat([cache["k_raw"], k_raw], 2)
            cache["k"], cache["v"] = k, v
            if self.graph:
                cache["k_raw"] = k_raw_all

        # 2) causal tile routing: route every tile whose FIRST token lies in this chunk
        first_tile, last_tile = start_pos // C, (start_pos + T - 1) // C
        n_known = cache["w"].shape[1] if (cache is not None and "w" in cache) else 0
        new_tiles = torch.arange(n_known, last_tile + 1, device=x.device)
        offsets = new_tiles * C - start_pos
        if n_known <= last_tile and n_known * C < start_pos:
            raise ValueError("A tile first token was skipped; routes out of sync")
        aux = {}
        sel_new = None
        if len(new_tiles):
            P_new, w_new, idx_new = self._route(x[:, offsets])
            loss, f, ent = self._balance(P_new, idx_new)
            aux = {"balance_loss": loss, "mech_freq": f, "router_entropy": ent}
            with torch.no_grad():
                selected = F.one_hot(idx_new, self.M).float()
                selected_weight = torch.zeros_like(P_new).scatter_add_(
                    -1, idx_new, w_new
                )
                aux.update(
                    query_tile_indices=new_tiles.detach(),
                    mech_frequency_by_query_tile=selected.mean(dim=(0, 2, 3)),
                    mech_probability_by_query_tile=P_new.detach().mean(dim=(0, 2)),
                    mech_selected_weight_by_query_tile=selected_weight.mean(dim=(0, 2)),
                )
            if self.graph:
                sel_new = self._select_tiles(q_raw[:, :, offsets], k_raw_all, new_tiles)
        else:
            w_new = x.new_zeros(B, 0, H, K, dtype=torch.float32)
            idx_new = torch.zeros(B, 0, H, K, dtype=torch.long, device=x.device)
        if self.graph and sel_new is None:
            sel_new = torch.zeros(
                B, 0, H, self.n_max, dtype=torch.bool, device=x.device
            )
        if cache is not None and n_known:
            w_all = torch.cat([cache["w"], w_new], 1)
            idx_all = torch.cat([cache["idx"], idx_new], 1)
            sel_all = torch.cat([cache["sel"], sel_new], 1) if self.graph else None
        else:
            w_all, idx_all, sel_all = w_new, idx_new, sel_new
        if cache is not None:
            cache["w"], cache["idx"] = w_all.detach(), idx_all
            if self.graph:
                cache["sel"] = sel_all

        # slot s = h*K + j  (maps to kv head s // (H*K/Hkv) == h // (H/Hkv))
        w_t = w_all[:, first_tile : last_tile + 1]  # [B, n, H, K]
        i_t = idx_all[:, first_tile : last_tile + 1]
        n = w_t.shape[1]
        mech = i_t.permute(0, 2, 3, 1).reshape(B, H * K, n)  # [B, HK, n]
        wslot = w_t.permute(0, 2, 3, 1).reshape(B, H * K, n)
        picked_slots = None
        if self.graph:  # [B, HK, n, n_max]
            picked_slots = (
                sel_all[:, first_tile : last_tile + 1]
                .permute(0, 2, 1, 3)
                .repeat_interleave(K, dim=1)
            )

        q_slots = q.repeat_interleave(K, dim=1)  # [B, HK, T, dh]
        q_pos = torch.arange(start_pos, start_pos + T, device=x.device)
        q_tile = q_pos // C - first_tile  # [T]

        if cfg.backend == "flex" and not HAS_FLEX:
            raise RuntimeError(
                "FlexAttention unavailable; install compatible torch or explicitly select dense"
            )
        use_flex = (
            cfg.backend == "flex"
            and (x.is_cuda or self._force_flex)
            and cache is None
            and start_pos == 0
        )
        if use_flex:
            n_max = self.n_max
            if picked_slots is not None:

                def mask_mod(b, h, q_idx, kv_idx):
                    qt = (q_idx // C).clamp(max=n - 1)
                    picked = picked_slots[b, h, qt, (kv_idx // C).clamp(max=n_max - 1)]
                    return (
                        (q_idx < T)
                        & (kv_idx < T)
                        & geometry_allows(mech[b, h, qt], q_idx, kv_idx, cfg, picked)
                    )

            else:

                def mask_mod(b, h, q_idx, kv_idx):
                    return (
                        (q_idx < T)
                        & (kv_idx < T)
                        & geometry_allows(
                            mech[b, h, (q_idx // C).clamp(max=n - 1)],
                            q_idx,
                            kv_idx,
                            cfg,
                        )
                    )

            if C == 128:
                bm = geometry_block_mask(
                    mech,
                    T,
                    cfg,
                    mask_mod,
                    None if picked_slots is None else picked_slots[..., :n],
                )
            else:
                bm = create_block_mask(
                    mask_mod,
                    B,
                    H * K,
                    T,
                    T,
                    device=x.device,
                    BLOCK_SIZE=C,
                    _compile=x.is_cuda,
                )
            out = _run_flex(q_slots, k, v, bm)
        else:
            S = k.shape[2]
            kv_pos = torch.arange(S, device=x.device)
            m = mech[:, :, q_tile]  # [B, HK, T]
            picked = None
            if picked_slots is not None:
                picked = picked_slots[:, :, q_tile][..., kv_pos // C]  # [B, HK, T, S]
            allow = geometry_allows(
                m[..., None],
                q_pos[None, None, :, None],
                kv_pos[None, None, None, :],
                cfg,
                picked,
            )  # [B, HK, T, S]
            out = F.scaled_dot_product_attention(
                q_slots, k, v, attn_mask=allow, enable_gqa=True
            )

        # 3) merge: o_h = sum_j P_{h,m_j} * A_{m_j}
        out = out * wslot[:, :, q_tile, None].to(out.dtype)
        out = out.view(B, H, K, T, dh).sum(2).transpose(1, 2).reshape(B, T, H * dh)
        return self.o_proj(out), aux


# =============================================================================
# Grouped GEMM acceleration
# =============================================================================

# Ampere-compatible dropless grouped GEMM with explicit autograd.
#
# No Hopper TMA instructions, token dropping, padding capacity, or CPU count reads.
# A is [sum(tokens), K], W is [experts,K,N], offsets are cumulative token counts.
# Persistent forward/input-gradient kernels use tensor-core dot products; each
# weight-gradient tile has a unique writer and accumulates in FP32.

if HAS_TRITON:

    @triton.jit
    def _grouped_forward(
        A,
        W,
        O,
        Y,
        K: tl.constexpr,
        N: tl.constexpr,
        AS0: tl.constexpr,
        AS1: tl.constexpr,
        WS0: tl.constexpr,
        WS1: tl.constexpr,
        WS2: tl.constexpr,
        E: tl.constexpr,
        GRID: tl.constexpr,
        BM: tl.constexpr = 32,
        BN: tl.constexpr = 64,
        BK: tl.constexpr = 32,
    ):
        tile_id = tl.program_id(0)
        prior_tiles = 0
        start = 0
        for expert in range(E):
            end = tl.load(O + expert)
            count = end - start
            nt = tl.cdiv(N, BN)
            tiles = tl.cdiv(count, BM) * nt
            while tile_id >= prior_tiles and tile_id < prior_tiles + tiles:
                local = tile_id - prior_tiles
                mi = (local // nt) * BM + tl.arange(0, BM)
                ni = (local % nt) * BN + tl.arange(0, BN)
                ki = tl.arange(0, BK)
                acc = tl.zeros((BM, BN), tl.float32)
                for kk in range(tl.cdiv(K, BK)):
                    kval = kk * BK + ki
                    av = tl.load(
                        A + (start + mi[:, None]) * AS0 + kval[None, :] * AS1,
                        (mi[:, None] < count) & (kval[None, :] < K),
                        other=0,
                    )
                    wv = tl.load(
                        W + expert * WS0 + kval[:, None] * WS1 + ni[None, :] * WS2,
                        (kval[:, None] < K) & (ni[None, :] < N),
                        other=0,
                    )
                    acc += tl.dot(av, wv)
                tl.store(
                    Y + (start + mi[:, None]) * N + ni[None, :],
                    acc,
                    (mi[:, None] < count) & (ni[None, :] < N),
                )
                tile_id += GRID
            prior_tiles += tiles
            start = end

    @triton.jit
    def _grouped_weight_grad(
        A,
        DY,
        O,
        DW,
        K: tl.constexpr,
        N: tl.constexpr,
        AS0: tl.constexpr,
        AS1: tl.constexpr,
        DS0: tl.constexpr,
        DS1: tl.constexpr,
        BM: tl.constexpr = 32,
        BN: tl.constexpr = 64,
        BT: tl.constexpr = 32,
    ):
        expert = tl.program_id(0)
        ki = tl.program_id(1) * BM + tl.arange(0, BM)
        ni = tl.program_id(2) * BN + tl.arange(0, BN)
        end = tl.load(O + expert)
        start = 0
        if expert > 0:
            start = tl.load(O + expert - 1)
        ti = tl.arange(0, BT)
        acc = tl.zeros((BM, BN), tl.float32)
        for block in range(tl.cdiv(end - start, BT)):
            token = start + block * BT + ti
            av = tl.load(
                A + token[None, :] * AS0 + ki[:, None] * AS1,
                (token[None, :] < end) & (ki[:, None] < K),
                other=0,
            )
            dy = tl.load(
                DY + token[:, None] * DS0 + ni[None, :] * DS1,
                (token[:, None] < end) & (ni[None, :] < N),
                other=0,
            )
            acc += tl.dot(av, dy)
        tl.store(
            DW + expert * K * N + ki[:, None] * N + ni[None, :],
            acc,
            (ki[:, None] < K) & (ni[None, :] < N),
        )


def _forward(a, w, offsets):
    if not HAS_TRITON or not a.is_cuda:
        raise RuntimeError("Ampere grouped GEMM requires CUDA and Triton")
    if a.dtype not in (torch.bfloat16, torch.float16) or w.dtype != a.dtype:
        raise ValueError("Grouped GEMM expects matching BF16/FP16 operands")
    if (
        a.ndim != 2
        or w.ndim != 3
        or a.shape[1] != w.shape[1]
        or offsets.numel() != w.shape[0]
    ):
        raise ValueError("Grouped GEMM shape mismatch")
    out = torch.empty((a.shape[0], w.shape[2]), device=a.device, dtype=a.dtype)
    grid = torch.cuda.get_device_properties(a.device).multi_processor_count
    with torch.cuda.device(a.device):
        _grouped_forward[(grid,)](
            a,
            w,
            offsets,
            out,
            a.shape[1],
            w.shape[2],
            *a.stride(),
            *w.stride(),
            w.shape[0],
            grid,
            num_warps=4,
            num_stages=3,
        )
    return out


class _GroupedGEMM(torch.autograd.Function):
    @staticmethod
    def forward(ctx, a, w, offsets):
        ctx.save_for_backward(a, w, offsets)
        return _forward(a, w, offsets)

    @staticmethod
    def backward(ctx, grad):
        a, w, offsets = ctx.saved_tensors
        da = (
            _forward(grad, w.transpose(-1, -2), offsets)
            if ctx.needs_input_grad[0]
            else None
        )
        dw = None
        if ctx.needs_input_grad[1]:
            dw = torch.empty(w.shape, device=w.device, dtype=w.dtype)
            grid = (
                w.shape[0],
                triton.cdiv(w.shape[1], 32),
                triton.cdiv(w.shape[2], 64),
            )
            with torch.cuda.device(a.device):
                _grouped_weight_grad[grid](
                    a,
                    grad,
                    offsets,
                    dw,
                    w.shape[1],
                    w.shape[2],
                    *a.stride(),
                    *grad.stride(),
                    num_warps=4,
                    num_stages=3,
                )
        return da, dw, None


def ampere_grouped_mm(a, w, offsets):
    return _GroupedGEMM.apply(a, w, offsets)


def grouped_mm(a, b, offsets):
    return ampere_grouped_mm(a, b, offsets)


# =============================================================================
# Model building blocks
# =============================================================================


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return (
            x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + self.eps)
        ).to(x.dtype) * self.weight


class DepthAttention(nn.Module):
    """Bank entries [B,T,H,d] preserve token index exactly; never reduce over T.
    Keep the input embedding plus the most recent layer states. Keys and values
    are projected once when written, with autograd retained through the bank.
    """

    def __init__(self, cfg):
        super().__init__()
        self.heads = cfg.depth_heads
        self.dim = cfg.depth_dim
        self.q_proj = nn.Linear(cfg.embed_dim, cfg.depth_dim, bias=False)
        self.out_proj = nn.Linear(cfg.depth_dim, cfg.embed_dim, bias=False)
        self.gate = nn.Parameter(torch.tensor(cfg.depth_gate_init))

    def forward(self, x, bank):
        b, t, _ = x.shape
        q = self.q_proj(x).view(b, t, self.heads, self.dim // self.heads)
        keys = torch.stack([kv[0] for kv in bank], dim=-2)
        vals = torch.stack([kv[1] for kv in bank], dim=-2)
        scores = (q.unsqueeze(-2).float() * keys.float()).sum(-1) / math.sqrt(
            q.shape[-1]
        )
        weights = scores.softmax(-1).to(vals.dtype)
        y = (weights.unsqueeze(-1) * vals).sum(-2).reshape(b, t, self.dim)
        return self.gate.tanh() * self.out_proj(y)


class DroplessMoE(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.router = nn.Linear(cfg.embed_dim, cfg.num_experts, bias=False)
        self.up = nn.Parameter(
            torch.empty(cfg.num_experts, cfg.embed_dim, cfg.expert_hidden_dim * 2)
        )
        self.down = nn.Parameter(
            torch.empty(cfg.num_experts, cfg.expert_hidden_dim, cfg.embed_dim)
        )
        nn.init.normal_(self.up, std=0.02)
        nn.init.normal_(self.down, std=0.02 / math.sqrt(2 * cfg.num_layers))

    def forward(self, x):
        cfg = self.cfg
        flat = x.reshape(-1, x.shape[-1])
        prob = router_logits_fp32(self.router, flat).softmax(-1)
        weights, experts = prob.topk(cfg.expert_top_k, -1)
        weights = weights / weights.sum(-1, keepdim=True)
        order = experts.flatten().argsort(stable=True)
        token_ids = order // cfg.expert_top_k
        dispatched = flat[token_ids]
        counts = torch.zeros(cfg.num_experts, device=x.device, dtype=torch.int32)
        counts.scatter_add_(
            0, experts.flatten(), torch.ones_like(experts.flatten(), dtype=torch.int32)
        )
        if cfg.expert_backend == "grouped" and x.is_cuda:
            # Explicit casts are differentiable and avoid ambiguous autocast support.
            offsets = counts.cumsum(0, dtype=torch.int32)
            h = grouped_mm(
                dispatched.to(torch.bfloat16), self.up.to(torch.bfloat16), offsets
            )
            gate, value = h.chunk(2, -1)
            h = (F.silu(gate) * value).contiguous()
            result = grouped_mm(h, self.down.to(torch.bfloat16), offsets)
        else:
            # One synchronization for all counts, instead of nonzero()+sync per expert.
            sizes = counts.cpu().tolist()
            chunks = dispatched.split(sizes, 0)
            result = torch.cat(
                [self._expert(chunk, i) for i, chunk in enumerate(chunks)], 0
            )
        # Autocast produces BF16 expert outputs while the residual may be FP32.
        # Both backends must mix in the residual dtype; index_add does not promote.
        result = result.to(flat.dtype)
        weighted = result * weights.flatten()[order, None].to(result.dtype)
        mixed = torch.zeros_like(flat).index_add(0, token_ids, weighted)
        frequency = counts.float() / experts.numel()
        balance = cfg.balance_coef * cfg.num_experts * (frequency * prob.mean(0)).sum()
        return mixed.view_as(x), balance, frequency.detach()

    def _expert(self, x, i):
        gate, value = (x @ self.up[i]).chunk(2, -1)
        return (F.silu(gate) * value) @ self.down[i]


class AttentionBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.hnorm, self.dnorm, self.mnorm = (RMSNorm(cfg.embed_dim, cfg.rms_norm_eps) for _ in range(3))
        ac = RoutedAttentionConfig(
            d_model=cfg.embed_dim,
            n_q_heads=cfg.n_q_heads,
            n_kv_heads=cfg.n_kv_heads,
            top_k=cfg.attention_top_k,
            tile=cfg.tile,
            local_window=cfg.local_window,
            strided_recent=cfg.strided_recent,
            strided_every=cfg.strided_every,
            graph_top_r=cfg.graph_top_r,
            graph_recent=cfg.graph_recent,
            max_seq_len=cfg.max_seq_len,
            router_init_std=cfg.router_init_std,
            balance_coef=cfg.balance_coef,
            backend=cfg.attention_backend,
        )
        head_dim = cfg.embed_dim // cfg.n_q_heads
        if cfg.rope_type == "warped":
            pos = WarpedRoPE(cfg.embed_dim, head_dim, cfg.n_q_heads, cfg.n_kv_heads)
        elif cfg.rope_type == "blackhole":
            pos = BlackHolePos(
                BlackHoleRoPE(head_dim, cfg.max_seq_len, discrepancy_modulation=False)
            )
        else:
            pos = None
        self.horizontal = RoutedSequenceAttention(ac, pos)
        self.vertical = DepthAttention(cfg)
        self.moe = DroplessMoE(cfg)
        self.depth_mode = cfg.depth_mode

    def forward(self, x, bank, cache=None, start_pos=0):
        h, aux = self.horizontal(self.hnorm(x), cache=cache, start_pos=start_pos)
        x = x + h
        if self.depth_mode == "attention":
            x = x + self.vertical(self.dnorm(x), bank)
        m, mloss, freq = self.moe(self.mnorm(x))
        aux["expert_frequency"] = freq
        return x + m, aux.get("balance_loss", x.new_zeros(())) + mloss, aux


# =============================================================================
# Hugging Face cache and output interfaces
# =============================================================================


class TAMECache(Cache):
    """Per-sequence KV and tile-route state with HF beam and rollback operations.
    Physical token count includes padding; each row's rotary position counts only
    real tokens. Depth attention is same-token across layers, so its bank is
    reconstructed for new tokens and needs no cross-token cache.
    """

    def __init__(self, num_layers, tile):
        super().__init__(layers=[])
        self.num_layers, self.tile = num_layers, tile
        self.rows = []
        self.mask = None

    def get_seq_length(self, layer_idx=0):
        return 0 if self.mask is None else self.mask.shape[1]

    def get_max_cache_shape(self, layer_idx=0):
        return -1

    def get_mask_sizes(self, query_length, layer_idx=0):
        return self.get_seq_length() + query_length, 0

    @property
    def is_compileable(self):
        return False

    @property
    def is_initialized(self):
        return bool(self.rows)

    @property
    def is_sliding(self):
        return [False] * self.num_layers

    @property
    def max_batch_size(self):
        return len(self.rows)

    def reset(self):
        self.rows = []
        self.mask = None

    def row_length(self, index):
        row = self.rows[index]
        return row[0]["k"].shape[2] if row and "k" in row[0] else 0

    def reorder_cache(self, beam_idx):
        indices = beam_idx.tolist()
        self.rows = [[dict(layer) for layer in self.rows[i]] for i in indices]
        if self.mask is not None:
            self.mask = self.mask.index_select(0, beam_idx.to(self.mask.device))

    def batch_select_indices(self, indices):
        self.reorder_cache(indices)

    def batch_repeat_interleave(self, repeats):
        if self.mask is not None:
            indices = torch.arange(
                len(self.rows), device=self.mask.device
            ).repeat_interleave(repeats)
            self.reorder_cache(indices)

    def crop(self, max_length):
        seen = self.get_seq_length()
        end = max(0, seen + max_length) if max_length < 0 else min(seen, max_length)
        if self.mask is None:
            return
        self.mask = self.mask[:, :end]
        for i, row in enumerate(self.rows):
            length = int(self.mask[i].sum())
            tiles = (length + self.tile - 1) // self.tile
            for layer in row:
                for key in ("k", "v", "k_raw"):  # [B, Hkv, S, dh]
                    if key in layer:
                        layer[key] = layer[key][:, :, :length]
                if "clock" in layer:  # [B, S, Hkv] warped time
                    layer["clock"] = layer["clock"][:, :length]
                for key in ("w", "idx", "sel"):  # [B, tiles, ...]
                    if key in layer:
                        layer[key] = layer[key][:, :tiles]


@dataclass
class TAMEOutput(ModelOutput):
    loss: torch.Tensor | None = None
    logits: torch.Tensor | None = None
    past_key_values: TAMECache | list | None = None
    aux_losses: dict | None = None
    attention_stats: list | None = None


# =============================================================================
# Complete language model configuration
# =============================================================================

class TAMEConfig(PretrainedConfig):
    model_type = "tamelm_two_axis"
    attribute_map = {
        "hidden_size": "embed_dim",
        "num_hidden_layers": "num_layers",
        "num_attention_heads": "n_q_heads",
        "num_key_value_heads": "n_kv_heads",
        "max_position_embeddings": "max_seq_len",
    }
    # Keys introduced after the alpha checkpoint, with the value that reproduces
    # the older architecture when a saved config predates them.
    _LEGACY_DEFAULTS = {"graph_top_r": 0, "init_scheme": "pilot"}

    def __init__(
        # Configuration instance; supplied automatically by Python.
        self,
        # Tokenizer vocabulary size; number of token IDs to embed/predict.
        vocab_size=44469,
        # Main token representation width.
        embed_dim=1024,
        # Number of stacked attention/depth/expert blocks.
        num_layers=14,
        # Available feed-forward experts per layer.
        num_experts=8,
        # Experts selected per token; their outputs are weighted together.
        expert_top_k=2,
        # Width of each expert gate/value branch before projecting back.
        expert_hidden_dim=4096,
        # Query heads; head width is embed_dim / n_q_heads (64 here).
        n_q_heads=16,
        # Shared key/value heads; eight query heads share each KV head here.
        n_kv_heads=2,
        # Attention geometries selected per query head and routing tile.
        attention_top_k=2,
        # Tokens per routing tile; first token sets routes. Also mask block size.
        tile=128,
        # Local visibility: current token plus up to 255 prior tokens here.
        local_window=256,
        # Dense recent window for strided attention, including current token.
        strided_recent=128,
        # Also expose past tiles at tile-distance multiples of this value.
        strided_every=4,
        # Content-selected completed past tiles; 0 disables graph geometry.
        graph_top_r=2,
        # Dense recent window added to graph-selected tiles.
        graph_recent=128,
        # "flex" or "dense"; cached decoding uses the dense path.
        attention_backend="flex",
        # "grouped" uses CUDA/Triton; "reference" is PyTorch. CPU uses reference.
        expert_backend="grouped",
        # Maximum real tokens per sequence, including generated continuation.
        max_seq_len=2048,
        # "attention" reads same-token earlier-layer states; "baseline" skips it.
        depth_mode="attention",
        # Depth bank capacity; here embedding entry plus recent layer entries.
        depth_slots=4,
        # Width of projected depth-attention representations.
        depth_dim=384,
        # Depth-attention heads; each has depth_dim / depth_heads dimensions.
        depth_heads=4,
        # Initial learned gate; tanh(gate) scales the depth contribution.
        depth_gate_init=0.01,
        # Auxiliary loss strength for balancing attention routes and experts.
        balance_coef=0.01,
        # "warped" learns token clock rates; alternatives: "standard", "blackhole".
        rope_type="warped",
        # Attention-router weight std under "v1"; not expert-router scale.
        router_init_std=0.01,
        # "v1" restores router/output scales; "pilot" retains older initialization.
        init_scheme="v1",
        # Token positions per vocabulary projection/CE chunk; saves memory.
        loss_chunk_size=256,
        # Tokenizer ID for padding; use attention_mask to mark padded positions.
        pad_token_id=0,
        # Tokenizer ID for beginning of sequence.
        bos_token_id=1,
        # Tokenizer ID for end of sequence; can stop generation.
        eos_token_id=2,
        # Generation cache default; direct forward() still defaults to False.
        use_cache=True,
        # RMS normalization stability epsilon; applied to every RMSNorm.
        rms_norm_eps=1e-6,
        # Additional Hugging Face configuration fields.
        **kwargs,
    ):
        kwargs.setdefault("tie_word_embeddings", True)
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            **kwargs,
        )
        values = locals().copy()
        for key, value in values.items():
            if key not in {"self", "kwargs", "__class__", "values"}:
                setattr(self, key, value)
        if not math.isfinite(rms_norm_eps) or rms_norm_eps <= 0:
            raise ValueError("rms_norm_eps must be finite and positive")
        if embed_dim % n_q_heads or (embed_dim // n_q_heads) % 2:
            raise ValueError(
                "Width must divide query heads with even rotary head dimension"
            )
        if depth_dim % depth_heads or depth_mode not in {"attention", "baseline"}:
            raise ValueError("Invalid depth dimensions/mode")
        if (
            not 1 <= expert_top_k <= num_experts
            or min(num_layers, depth_slots, loss_chunk_size) < 1
        ):
            raise ValueError("Invalid expert/depth/loss dimensions")
        if expert_backend not in {"reference", "grouped"} or rope_type not in {
            "blackhole",
            "standard",
            "warped",
        }:
            raise ValueError("Unknown expert backend or position transform")
        if graph_top_r < 0 or init_scheme not in {"v1", "pilot"} or router_init_std < 0:
            raise ValueError("Invalid graph_top_r, init_scheme or router_init_std")

    @classmethod
    def from_dict(cls, config_dict, **kwargs):
        config_dict = dict(config_dict)
        for key, legacy in cls._LEGACY_DEFAULTS.items():
            config_dict.setdefault(key, legacy)
        return super().from_dict(config_dict, **kwargs)


# =============================================================================
# Complete language model
# =============================================================================


class TAMELM(PreTrainedModel, GenerationMixin):
    config_class = TAMEConfig
    base_model_prefix = "tamelm"
    _tied_weights_keys = {"lm_head.weight": "token_embedding.weight"}

    # Construction and checkpoint compatibility

    def _init_weights(self, module):
        # Initialization below is the existing TAMELM scheme. HF post_init
        # records model metadata but must not draw extra random numbers or
        # change the seeded backbone used by the pilot and matched controls.
        return

    def __init__(self, config):
        super().__init__(config)
        self.token_embedding = nn.Embedding(config.vocab_size, config.embed_dim)
        self.blocks = nn.ModuleList(
            [AttentionBlock(config) for _ in range(config.num_layers)]
        )
        self.final_norm = RMSNorm(config.embed_dim, config.rms_norm_eps)
        self.depth_norm = RMSNorm(config.embed_dim, config.rms_norm_eps)
        self.depth_kv = nn.Linear(config.embed_dim, 2 * config.depth_dim, bias=False)
        self.lm_head = nn.Linear(config.embed_dim, config.vocab_size, bias=False)
        self.lm_head.weight = self.token_embedding.weight
        self.post_init()
        nn.init.normal_(self.token_embedding.weight, std=0.02)
        for module in self.modules():
            if isinstance(module, nn.Linear) and module is not self.lm_head:
                nn.init.normal_(module.weight, std=0.02)
        # HF loads under no_init_weights(); torch.ones there may leave this
        # nonpersistent routing mask uninitialized. It is fixed config state.
        for block in self.blocks:
            block.horizontal.allowed.fill_(True)
        if config.init_scheme == "v1":
            # The generic Linear loop above overwrites module-specific schemes.
            # Restore them so the run config describes the network that trains.
            out_std = 0.02 / math.sqrt(2 * config.num_layers)
            for block in self.blocks:
                nn.init.normal_(
                    block.horizontal.router.weight, std=config.router_init_std
                )
                nn.init.normal_(block.horizontal.o_proj.weight, std=out_std)
                nn.init.normal_(block.vertical.out_proj.weight, std=out_std)
        for block in self.blocks:
            if isinstance(block.horizontal.pos, WarpedRoPE):
                block.horizontal.pos.reset_parameters()  # exactly standard RoPE at step 0
        self.tokenizer = None

    # Embedding and tokenizer interfaces

    def get_input_embeddings(self):
        return self.token_embedding

    def load_state_dict(self, state_dict, strict=True, assign=False):
        # Pilot checkpoints predate serialization of this fixed routing mask.
        # Fill only missing masks; leave every learned tensor untouched.
        missing = [
            f"blocks.{i}.horizontal.allowed"
            for i in range(len(self.blocks))
            if f"blocks.{i}.horizontal.allowed" not in state_dict
        ]
        if missing:
            state_dict = dict(state_dict)
            for key in missing:
                state_dict[key] = torch.ones_like(self.get_buffer(key))
        return super().load_state_dict(state_dict, strict=strict, assign=assign)

    def set_input_embeddings(self, value):
        self.token_embedding = value
        self.lm_head.weight = value.weight

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, value):
        self.lm_head = value
        self.token_embedding.weight = value.weight

    def attach_tokenizer(self, tokenizer):
        import hashlib
        import json

        if len(tokenizer) != self.config.vocab_size:
            raise ValueError("Tokenizer vocabulary does not match model")
        for name in ("pad", "bos", "eos"):
            if getattr(tokenizer, name + "_token_id") != getattr(
                self.config, name + "_token_id"
            ):
                raise ValueError(f"Tokenizer {name} ID does not match model")
        backend = json.loads(tokenizer.backend_tokenizer.to_str())
        backend.pop("truncation", None)
        backend.pop("padding", None)
        identity = {
            "backend": backend,
            "special_ids": {
                name: getattr(tokenizer, name + "_token_id")
                for name in ("pad", "bos", "eos")
            },
        }
        digest = hashlib.sha256(
            json.dumps(identity, sort_keys=True, separators=(",", ":")).encode()
        ).hexdigest()
        if getattr(self.config, "tokenizer_sha256", digest) != digest:
            raise ValueError(
                "Tokenizer identity differs from the checkpoint vocabulary/policy"
            )
        self.config.tokenizer_sha256 = digest
        self.tokenizer = tokenizer

    # Depth state and packed forward path

    def _write_depth(self, x):
        b, t, _ = x.shape
        k, v = self.depth_kv(self.depth_norm(x)).chunk(2, -1)
        shape = (
            b,
            t,
            self.config.depth_heads,
            self.config.depth_dim // self.config.depth_heads,
        )
        return k.reshape(shape), v.reshape(shape)

    def _forward_unpadded(
        self,
        input_ids,
        labels=None,
        attention_mask=None,
        use_cache=False,
        past_key_values=None,
        start_pos=0,
        return_logits=None,
        **kwargs,
    ):
        if (
            input_ids.ndim != 2
            or not 0 < input_ids.shape[1] <= self.config.max_seq_len - start_pos
        ):
            raise ValueError("Input must be [batch, tokens] within max_seq_len")
        # Packed training has no padding. Do not accept a mask and silently ignore it.
        if attention_mask is not None:
            raise ValueError(
                "Use packed unpadded input_ids; padding masks are not supported"
            )
        if (use_cache or past_key_values is not None) and (
            labels is not None or torch.is_grad_enabled()
        ):
            raise ValueError(
                "KV cache requires inference without labels under torch.no_grad()"
            )
        caches = (
            past_key_values
            if past_key_values is not None
            else ([{} for _ in self.blocks] if use_cache else None)
        )
        if caches is not None and len(caches) != len(self.blocks):
            raise ValueError("One cache per layer required")
        x = self.token_embedding(input_ids)
        bank = [self._write_depth(x)] if self.config.depth_mode == "attention" else []
        losses, stats = [], []
        for i, block in enumerate(self.blocks):
            x, loss, info = block(
                x, bank, None if caches is None else caches[i], start_pos
            )
            losses.append(loss)
            stats.append({k: v.detach() for k, v in info.items()})
            if self.config.depth_mode == "attention" and i + 1 < len(self.blocks):
                bank.append(self._write_depth(x))
                if len(bank) > self.config.depth_slots:
                    bank = (
                        [bank[0]] + bank[-(self.config.depth_slots - 1) :]
                        if self.config.depth_slots > 1
                        else bank[-1:]
                    )
        x = self.final_norm(x)
        aux = torch.stack(losses).mean()
        primary, total = None, None
        if labels is not None:
            # Checkpoint the vocabulary projection/CE chunks: no retained [B,T,V]
            # logits. Recompute only this head in backward, not the full model.
            from torch.utils.checkpoint import checkpoint

            states, targets = x[:, :-1].reshape(-1, x.shape[-1]), labels[:, 1:].reshape(
                -1
            )
            if targets.numel() == 0:
                raise ValueError("Language-model loss requires at least two tokens")
            parts = []

            def chunk_loss(h, y):
                return F.cross_entropy(
                    F.linear(h, self.lm_head.weight).float(),
                    y,
                    reduction="sum",
                    ignore_index=-100,
                )

            for offset in range(0, targets.numel(), self.config.loss_chunk_size):
                h, y = (
                    states[offset : offset + self.config.loss_chunk_size],
                    targets[offset : offset + self.config.loss_chunk_size],
                )
                parts.append(
                    checkpoint(chunk_loss, h, y, use_reentrant=False)
                    if torch.is_grad_enabled()
                    else chunk_loss(h, y)
                )
            primary = torch.stack(parts).sum() / (targets != -100).sum().clamp_min(1)
            total = primary + aux
        if return_logits is None:
            return_logits = labels is None
        logits = self.lm_head(x) if return_logits else None
        return TAMEOutput(
            loss=total,
            logits=logits,
            past_key_values=caches,
            aux_losses={
                "primary": None if primary is None else primary.detach(),
                "router_balance": aux.detach(),
            },
            attention_stats=stats,
        )

    # Generation integration

    @classmethod
    def _supports_default_dynamic_cache(cls):
        # A plain DynamicCache has no tile-routing state.
        return False

    def _prepare_cache_for_generation(
        self, generation_config, model_kwargs, *args, **kwargs
    ):
        implementation = generation_config.cache_implementation
        if implementation not in (None, "dynamic"):
            raise ValueError(
                "TAMELM supports its dynamic KV/routing cache; static, offloaded and quantized caches are unsupported"
            )
        cache = model_kwargs.get("past_key_values")
        if cache is not None and not isinstance(cache, TAMECache):
            raise TypeError("TAMELM requires TAMECache to preserve tile-routing state")

    def _validate_generated_length(
        self, generation_config, input_ids_length, has_default_max_length
    ):
        super()._validate_generated_length(
            generation_config, input_ids_length, has_default_max_length
        )
        if generation_config.max_length > self.config.max_seq_len:
            raise ValueError(
                f"Requested generation exceeds max_seq_len={self.config.max_seq_len}; shorten the prompt or max_new_tokens"
            )

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        use_cache=True,
        next_sequence_length=None,
        position_ids=None,
        **kwargs,
    ):
        seen = past_key_values.get_seq_length() if past_key_values is not None else 0
        if seen:
            if input_ids.shape[1] > seen:
                input_ids = input_ids[:, seen:]
            elif next_sequence_length is not None:
                input_ids = input_ids[:, -next_sequence_length:]
        if position_ids is not None:
            position_ids = position_ids[:, -input_ids.shape[1] :]
        return {
            "input_ids": input_ids.contiguous(),
            "past_key_values": past_key_values,
            "attention_mask": attention_mask,
            "position_ids": position_ids,
            "use_cache": use_cache,
            "logits_to_keep": 1,
        }

    # Public forward path with padding and cache handling

    def forward(
        self,
        input_ids,
        labels=None,
        attention_mask=None,
        use_cache=False,
        past_key_values=None,
        start_pos=None,
        return_logits=None,
        return_dict=None,
        position_ids=None,
        logits_to_keep=0,
        token_type_ids=None,
        **kwargs,
    ):
        if input_ids.ndim != 2 or input_ids.shape[1] < 1:
            raise ValueError("input_ids must be a nonempty [batch, tokens] tensor")
        if token_type_ids is not None and bool(token_type_ids.any()):
            raise ValueError("TAMELM does not use token type embeddings")
        if kwargs.get("output_attentions") or kwargs.get("output_hidden_states"):
            raise NotImplementedError(
                "Attention/hidden-state collection is not implemented"
            )
        if (use_cache or past_key_values is not None) and (
            labels is not None or torch.is_grad_enabled()
        ):
            raise ValueError(
                "KV cache requires inference without labels under torch.no_grad()"
            )
        if past_key_values is not None and not isinstance(past_key_values, TAMECache):
            raise TypeError("past_key_values must be a TAMECache")
        if past_key_values is not None and not use_cache:
            raise ValueError("Set use_cache=True when providing past_key_values")
        cache = past_key_values
        if use_cache and cache is None:
            cache = TAMECache(len(self.blocks), self.config.tile)
        seen = cache.get_seq_length() if cache is not None else 0
        if start_pos is not None and start_pos != seen:
            raise ValueError("start_pos disagrees with the cache token count")
        batch, length = input_ids.shape
        if labels is not None and labels.shape != input_ids.shape:
            raise ValueError("labels must match input_ids")
        if attention_mask is None:
            current_mask = torch.ones_like(input_ids, dtype=torch.bool)
            mask = (
                torch.cat((cache.mask, current_mask), dim=1) if seen else current_mask
            )
        else:
            if attention_mask.shape != (batch, seen + length):
                raise ValueError("attention_mask must cover cached and new tokens")
            if not bool(((attention_mask == 0) | (attention_mask == 1)).all()):
                raise ValueError("attention_mask must contain only zeros and ones")
            mask = attention_mask.to(device=input_ids.device, dtype=torch.bool)
            if seen and not torch.equal(mask[:, :seen], cache.mask):
                raise ValueError("attention_mask changes already cached tokens")
            current_mask = mask[:, seen:]
        if cache is not None:
            if not cache.rows:
                cache.rows = [[{} for _ in self.blocks] for _ in range(batch)]
            if len(cache.rows) != batch:
                raise ValueError(
                    "Cache batch size differs; reorder or expand it explicitly"
                )
        if (
            attention_mask is None
            and cache is None
            and length > self.config.max_seq_len
        ) or (
            (attention_mask is not None or cache is not None)
            and bool((mask.sum(1) > self.config.max_seq_len).any())
        ):
            raise ValueError("Real token count exceeds max_seq_len")
        if position_ids is not None:
            expected = mask.long().cumsum(-1)[:, seen:] - 1
            if position_ids.shape != input_ids.shape or not torch.equal(
                position_ids[current_mask], expected[current_mask]
            ):
                raise ValueError("position_ids must follow unpadded token positions")
        # Packed training and unpadded full forwards retain the vectorized path.
        if cache is None and (attention_mask is None or bool(current_mask.all())):
            out = self._forward_unpadded(
                input_ids, labels=labels, return_logits=return_logits
            )
        else:
            if return_logits is None:
                return_logits = labels is None
            logits = None
            weighted_loss, weighted_primary, weighted_aux, counts = [], [], [], []
            stats = []
            for i in range(batch):
                positions = current_mask[i].nonzero(as_tuple=True)[0]
                if not positions.numel():
                    raise ValueError(
                        "Each sequence must provide at least one real token"
                    )
                ids = input_ids[i : i + 1].index_select(1, positions)
                target = (
                    labels[i : i + 1].index_select(1, positions)
                    if labels is not None
                    else None
                )
                offset = cache.row_length(i) if cache is not None else 0
                native = self._forward_unpadded(
                    ids,
                    labels=target,
                    use_cache=cache is not None,
                    past_key_values=cache.rows[i] if cache is not None else None,
                    start_pos=offset,
                    return_logits=return_logits,
                )
                if native.logits is not None:
                    if logits is None:
                        logits = native.logits.new_zeros(
                            batch, length, self.config.vocab_size
                        )
                    logits[i, positions] = native.logits[0]
                if labels is not None:
                    count = (target[:, 1:] != -100).sum()
                    counts.append(count)
                    weighted_loss.append(native.loss * count)
                    weighted_primary.append(native.aux_losses["primary"] * count)
                    weighted_aux.append(native.aux_losses["router_balance"] * count)
                stats = native.attention_stats
            if cache is not None:
                cache.mask = mask.detach().clone()
            denom = torch.stack(counts).sum().clamp_min(1) if counts else None
            out = TAMEOutput(
                loss=torch.stack(weighted_loss).sum() / denom if counts else None,
                logits=logits,
                past_key_values=cache,
                aux_losses={
                    "primary": (
                        torch.stack(weighted_primary).sum() / denom if counts else None
                    ),
                    "router_balance": (
                        torch.stack(weighted_aux).sum() / denom
                        if counts
                        else input_ids.new_zeros((), dtype=torch.float32)
                    ),
                },
                attention_stats=None,
            )  # per-layer routing summaries are defined for packed batches
        if logits_to_keep and out.logits is not None:
            out.logits = out.logits[:, -logits_to_keep:]
        return out if return_dict is not False else out.to_tuple()

    # Portable Hugging Face export

    def export_hf_code(self, directory):
        from pathlib import Path
        import inspect
        import sys

        directory = Path(directory)
        directory.mkdir(parents=True, exist_ok=True)
        source = _NOTEBOOK_SOURCES.get("model")
        if source is None:
            source = inspect.getsource(sys.modules[type(self).__module__])
        # Registration is notebook-local; portable loading uses auto_map.
        source = source.split("\nAutoConfig.register(TAMEConfig.model_type")[0]
        source = source.split("\n# Notebook-only source capture")[0]
        (directory / "modeling_tamelm.py").write_text(source, encoding="utf-8")
        self.config.auto_map = {
            "AutoConfig": "modeling_tamelm.TAMEConfig",
            "AutoModelForCausalLM": "modeling_tamelm.TAMELM",
        }
        self.config.architectures = ["TAMELM"]
        self.config.save_pretrained(directory)
        self.generation_config.save_pretrained(directory)

    def save_pretrained(self, save_directory, *args, **kwargs):
        self.export_hf_code(save_directory)
        result = super().save_pretrained(save_directory, *args, **kwargs)
        if self.tokenizer is not None:
            self.tokenizer.save_pretrained(save_directory)
        return result