#@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 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