Upload TAMELM

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,47 @@
+{
+  "architectures": [
+    "TAMELM"
+  ],
+  "attention_backend": "flex",
+  "attention_top_k": 2,
+  "auto_map": {
+    "AutoConfig": "modeling_tamelm.TAMEConfig",
+    "AutoModelForCausalLM": "modeling_tamelm.TAMELM"
+  },
+  "balance_coef": 0.01,
+  "bos_token_id": 1,
+  "depth_dim": 768,
+  "depth_gate_init": 0.01,
+  "depth_heads": 8,
+  "depth_mode": "attention",
+  "depth_slots": 16,
+  "dtype": "float32",
+  "embed_dim": 1024,
+  "eos_token_id": 2,
+  "expert_backend": "grouped",
+  "expert_hidden_dim": 4096,
+  "expert_top_k": 2,
+  "graph_recent": 256,
+  "graph_top_r": 2,
+  "init_scheme": "v1",
+  "local_window": 256,
+  "loss_chunk_size": 256,
+  "max_seq_len": 1024,
+  "model_type": "tamelm_two_axis",
+  "n_kv_heads": 4,
+  "n_q_heads": 16,
+  "num_experts": 4,
+  "num_layers": 12,
+  "pad_token_id": 0,
+  "rms_norm_eps": 1e-06,
+  "rope_type": "warped",
+  "router_init_std": 0.01,
+  "strided_every": 4,
+  "strided_recent": 128,
+  "tie_word_embeddings": true,
+  "tile": 128,
+  "tokenizer_sha256": "0d93d8f943b1a6627d62395c2fd583894cfe577b8aca1ca946495f7abba50958",
+  "transformers_version": "5.5.0",
+  "use_cache": true,
+  "vocab_size": 49152
+}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 0,
+  "transformers_version": "5.5.0",
+  "use_cache": true
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47febb594dd8e9e95a739ef251a0ab49dd3cc85a4b07445cbd9a5f5ce67db20a
+size 2828580568

modeling_tamelm.py

@@ -0,0 +1,1846 @@
+#@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
+