diff --git "a/model.py" "b/model.py" new file mode 100644--- /dev/null +++ "b/model.py" @@ -0,0 +1,2217 @@ +import functools +import math +import os +import warnings +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.utils.checkpoint +from torch.nn import functional as F + +_native_rms_norm = getattr(F, "rms_norm", None) +_HAS_NATIVE_RMSNORM = _native_rms_norm is not None +if not _HAS_NATIVE_RMSNORM: + raise ImportError("model.py requires torch>=2.4 (F.rms_norm)") + +try: + import config as alpha_config +except ImportError: # pragma: no cover - fallback for package-style imports + from . import config as alpha_config + +try: + import kernel +except ImportError: # pragma: no cover - fallback for package-style imports + from . import kernel + +# -- Optional Triton fusions (triton_kernel.py) -------------------------------- +# Auto-detected, no config flag. Kill switch for debugging / A-B timing: +# ALPHA_DISABLE_TRITON=1 +# Every op is verified once against the eager reference (outputs AND grads) on a +# small slice of the first real call, then used; any failure falls back to eager +# for that (op, dtype/shape-key) and warns once. See _triton_attempt below. +# Config switch: config.py `use_triton_fusions` (default False). Off = no triton_kernel import and every +# `if _HAS_TRITON` branch is dead code under dynamo. +def _triton_fusions_on_in_config() -> bool: + return bool(getattr(alpha_config, "use_triton_fusions", False)) + + +_triton_kernel = None +if os.environ.get("ALPHA_DISABLE_TRITON") != "1" and _triton_fusions_on_in_config(): + try: + import triton_kernel as _triton_kernel + except ImportError: # pragma: no cover - fallback for package-style imports + try: + from . import triton_kernel as _triton_kernel + except ImportError: + _triton_kernel = None + except Exception as _e: # noqa: BLE001 -- file present but broken: warn, run eager + warnings.warn(f"[model.py/triton] triton_kernel failed to import ({_e!r}) -- running eager.") + _triton_kernel = None + + +def _triton_ready() -> bool: + if _triton_kernel is None: + return False + for _name in ("is_available", "_ok"): + _fn = getattr(_triton_kernel, _name, None) + if callable(_fn): + try: + return bool(_fn()) + except Exception: # noqa: BLE001 + return False + return True + + +# Static module constant: when False dynamo drops every Triton branch at trace +# time (no graph breaks, no overhead). +_HAS_TRITON = _triton_ready() +_TRITON_OK: set = set() +_TRITON_BAD: set = set() +_triton_warned: set = set() + +try: + from ngram import Engram, NgramHasher +except ImportError: # pragma: no cover - fallback for package-style imports + try: + from .ngram import Engram, NgramHasher + except ImportError: + Engram = None + NgramHasher = None + + + +_alpha_config_warned: set[str] = set() + +_TORCHAO_AVAILABLE: bool = False +_NVFP4_RHT_AVAILABLE: bool = False +_INT8_AVAILABLE: bool = False +# BUGFIX: these two were only bound inside the `try` body. If `import torchao` +# succeeded but a later import in the same try block raised, they stayed +# undefined and apply_int8_to_model() hit a NameError instead of the intended +# RuntimeError. Bind safe defaults up front. +_INT8_GROUPED_AVAILABLE: bool = False +_Int8WOConfig = None # type: ignore[assignment] +try: + import torchao as _torchao_probe # noqa: F401 + _TORCHAO_AVAILABLE = True + from torchao.quantization import quantize_ as _torchao_quantize + from torchao.quantization import ( + Int8DynamicActivationInt8WeightConfig as _Int8Config, + ) + + try: + from torchao.quantization import Int8WeightOnlyConfig as _Int8WOConfig + _INT8_GROUPED_AVAILABLE = True + except ImportError: + _Int8WOConfig = None # type: ignore[assignment] + _INT8_GROUPED_AVAILABLE = False + _INT8_AVAILABLE = True + try: + from torchao.prototype.mx_formats import ( + NVFP4DynamicActivationNVFP4WeightConfig as _NVFP4Config, + ) + _NVFP4_RHT_AVAILABLE = True + except ImportError: + _NVFP4Config = None # type: ignore[assignment] +except ImportError: + _torchao_quantize = None # type: ignore[assignment] + _Int8Config = None # type: ignore[assignment] + _NVFP4Config = None # type: ignore[assignment] + +_NVFP4_DISABLED_REASON: "str | None" = None + + +_NVFP4_SKIP_KEYWORDS: tuple[str, ...] = ( + "wte", "lm_head", "ln_", "ln_f", "norm", + "router", "attn_gate_proj", "expert_bias", "engram", +) + + +def _nvfp4_filter(module: nn.Module, fqn: str) -> bool: + """torchao filter_fn: True = quantize this module.""" + if not isinstance(module, nn.Linear): + return False + for kw in _NVFP4_SKIP_KEYWORDS: + if kw in fqn: + return False + + out_f, in_f = module.weight.shape + if in_f < 32 or out_f < 64 or in_f % 16 != 0: + return False + return True + + +def _nvfp4_enabled() -> bool: + global _NVFP4_DISABLED_REASON + if not getattr(alpha_config, 'use_nvfp4', False): + return False + if not torch.cuda.is_available(): + _NVFP4_DISABLED_REASON = "nvFP4: no CUDA device" + return False + major, _minor = torch.cuda.get_device_capability() + if major < 10: + msg = ( + f"nvFP4 requires Blackwell (sm_100+); found sm_{major}{_minor} " + f"({torch.cuda.get_device_name()}) -- disabled." + ) + if msg not in _alpha_config_warned: + _alpha_config_warned.add(msg) + warnings.warn(msg) + _NVFP4_DISABLED_REASON = msg + return False + if not _NVFP4_RHT_AVAILABLE: + _NVFP4_DISABLED_REASON = ( + "nvFP4: torchao or torchao.prototype.mx_formats not available. " + "Install: pip install torchao --pre" + ) + warnings.warn(_NVFP4_DISABLED_REASON) + return False + return True + + +def apply_nvfp4_to_model(model: nn.Module) -> int: + + before_ids = {fqn: id(m.weight) for fqn, m in model.named_modules() + if _nvfp4_filter(m, fqn)} + _torchao_quantize(model, _NVFP4Config(), filter_fn=_nvfp4_filter) + # BUGFIX: torchao swaps the *weight* for a tensor subclass and leaves the + # module a plain nn.Linear, so the old `type(m).__name__ != "Linear"` test + # always counted 0. Compare weight identity instead (same test the INT8 + # path already used). + current = dict(model.named_modules()) + count = 0 + for fqn, old_id in before_ids.items(): + m = current.get(fqn) + if m is not None and getattr(m, "weight", None) is not None and id(m.weight) != old_id: + count += 1 + model._quant_mode = "nvfp4" + model._use_nvfp4 = True + model._use_int8 = False + return count + + + +# -- FP8 *training* (torchao.float8.Float8Linear) ---------------------------------- +# Master weights stay in their own dtype and the state_dict is unchanged, so FP8-trained +# checkpoints are ordinary checkpoints (_quant_mode stays None). Applied explicitly from +# train.py (not GPT.__init__), so inference never sees Float8Linear. +_FP8_AVAILABLE: bool = False +try: + from torchao.float8 import convert_to_float8_training as _convert_to_float8_training + from torchao.float8 import Float8LinearConfig as _Float8LinearConfig + _FP8_AVAILABLE = True +except Exception: # noqa: BLE001 -- torchao missing/too old: FP8 simply unavailable + _convert_to_float8_training = None # type: ignore[assignment] + _Float8LinearConfig = None # type: ignore[assignment] + +# nn.Linear layers that always stay in BF16 (note: "gate" alone would also hit SwiGLU's W_gate). +_FP8_SKIP_KEYWORDS: tuple[str, ...] = ( + "wte", "lm_head", "ln_", "ln_f", "norm", "router", "expert_bias", "engram", "attn_gate_proj", "mtp", + "routed_experts", +) +_FP8_BLOCK_RE = __import__("re").compile(r"(?:^|\.)h\.(\d+)\.") +_FP8_RECIPES = ("tensorwise", "rowwise", "rowwise_with_gw_hp") + + +def _fp8_enabled() -> bool: + """use_fp8=True and the hardware/library can do it. Resolved once per run (train.py).""" + if not getattr(alpha_config, 'use_fp8', False): + return False + msg = None + if not torch.cuda.is_available(): + msg = "FP8 training: no CUDA device -- running BF16." + elif torch.cuda.get_device_capability() < (8, 9): + msg = (f"FP8 training needs sm_89+ (Ada/Hopper/Blackwell); found " + f"sm_{''.join(map(str, torch.cuda.get_device_capability()))} -- running BF16.") + elif not _FP8_AVAILABLE: + msg = "FP8 training: torchao.float8 not importable (pip install torchao) -- running BF16." + if msg is not None: + if msg not in _alpha_config_warned: + _alpha_config_warned.add(msg) + warnings.warn(msg) + return False + return True + + +def apply_fp8_to_model(model: nn.Module, n_layer: int, recipe: str = "rowwise_with_gw_hp", + bf16_tail_frac: float = 0.15, emulate: bool = False) -> dict: + """Swap eligible nn.Linear -> torchao Float8Linear in place (must run BEFORE torch.compile). + Keeps the last max(1, round(frac * n_layer)) blocks and everything in _FP8_SKIP_KEYWORDS in + BF16. Eligible = no bias, both dims multiples of 16 (scaled_mm) and >= 64 (tiny GEMMs lose). + `emulate` runs the matmul in high precision (tests on non-FP8 hardware only).""" + import dataclasses as _dc + if recipe not in _FP8_RECIPES: + raise ValueError(f"fp8 recipe {recipe!r} not in {_FP8_RECIPES}") + if not _FP8_AVAILABLE: + raise RuntimeError("torchao.float8 not importable") + n_tail = max(1, round(bf16_tail_frac * n_layer)) if bf16_tail_frac > 0 else 0 + first_bf16 = n_layer - n_tail + kept_tail = [0] + + def _filter(mod: nn.Module, fqn: str) -> bool: + if type(mod) is not nn.Linear or mod.bias is not None: + return False + if any(k in fqn for k in _FP8_SKIP_KEYWORDS): + return False + out_f, in_f = mod.weight.shape + if in_f % 16 or out_f % 16 or in_f < 64 or out_f < 64: + return False + mb = _FP8_BLOCK_RE.search(fqn) + if mb is not None and int(mb.group(1)) >= first_bf16: + kept_tail[0] += 1 + return False + return True + + cfg = _Float8LinearConfig.from_recipe_name(recipe) + if emulate: + cfg = _dc.replace(cfg, emulate=True) + before = {fqn for fqn, m in model.named_modules() if type(m) is nn.Linear} + _convert_to_float8_training(model, module_filter_fn=_filter, config=cfg) + converted = sum(1 for fqn, m in model.named_modules() + if fqn in before and type(m).__name__ == "Float8Linear") + model._fp8_train = converted > 0 + return {"converted": converted, "kept_bf16_tail": kept_tail[0], "tail_blocks": n_tail, "recipe": recipe} + + +_INT8_SKIP_KEYWORDS: tuple[str, ...] = _NVFP4_SKIP_KEYWORDS + +_INT8_DISABLED_REASON: "str | None" = None + + +def _int8_filter(module: nn.Module, fqn: str) -> bool: + """torchao filter_fn for INT8: True = quantize this module.""" + if not isinstance(module, nn.Linear): + return False + for kw in _INT8_SKIP_KEYWORDS: + if kw in fqn: + return False + + out_f, in_f = module.weight.shape + if in_f < 64 or out_f < 32: + return False + return True + + +def _int8_enabled() -> bool: + global _INT8_DISABLED_REASON + if not getattr(alpha_config, 'use_int8', False): + return False + if not torch.cuda.is_available(): + _INT8_DISABLED_REASON = "INT8: no CUDA device" + return False + if not _INT8_AVAILABLE: + _INT8_DISABLED_REASON = ( + "INT8: torchao not installed or Int8DynamicActivationInt8WeightConfig " + "not found. Install: pip install torchao" + ) + warnings.warn(_INT8_DISABLED_REASON) + return False + + if _nvfp4_enabled(): + msg = ( + "INT8 disabled: use_nvfp4=True also set and this is SM100+ hardware -- " + "nvFP4 takes priority. Set use_nvfp4=False to use INT8 instead." + ) + if msg not in _alpha_config_warned: + _alpha_config_warned.add(msg) + warnings.warn(msg) + _INT8_DISABLED_REASON = msg + return False + return True + + +def apply_int8_to_model(model: nn.Module, group_size: int = 128) -> int: + """Quantize eligible Linear layers to INT8 (Jetfire-style, dynamic activations). + + Config selection priority: + 1. Int8DynamicActivationInt8WeightConfig() -- stable torchao, per-channel + weights + dynamic per-tensor INT8 activations. Best for T4 throughput. + 2. Int8WeightOnlyConfig(group_size=N) -- weight-only, activations stay fp16. + Only used as fallback if (1) is unavailable (shouldn't happen on stable). + + Works on T4 (sm_75+). Fused kernel emitted by torch.compile via Inductor. + Must be called BEFORE torch.compile(). + """ + def _filter_with_gs(module: nn.Module, fqn: str) -> bool: + if not _int8_filter(module, fqn): + return False + + in_f = module.weight.shape[1] + if in_f % 32 != 0: # 32 = minimum alignment for any INT8 kernel + return False + return True + + + # torchao replaces module.weight in place and keeps the nn.Linear object, so + # identity must be tracked on the weight (same test apply_nvfp4_to_model uses). + before_ids = {fqn: id(m.weight) for fqn, m in model.named_modules() if _filter_with_gs(m, fqn)} + + + if _INT8_AVAILABLE and _Int8Config is not None: + config_obj = _Int8Config() + elif _INT8_GROUPED_AVAILABLE and _Int8WOConfig is not None: + warnings.warn( + f"[INT8] Int8DynamicActivationInt8WeightConfig unavailable -- " + f"falling back to Int8WeightOnlyConfig(group_size={group_size}). " + f"Activations will stay fp16 (weight-only quant).", + ) + config_obj = _Int8WOConfig(group_size=group_size) + else: + raise RuntimeError("[INT8] No usable INT8 config found in torchao -- pip install torchao") + + _torchao_quantize(model, config_obj, filter_fn=_filter_with_gs) + + + current = dict(model.named_modules()) + count = 0 + for fqn, old_id in before_ids.items(): + mod = current.get(fqn) + if mod is not None and getattr(mod, "weight", None) is not None and id(mod.weight) != old_id: + count += 1 + model._quant_mode = "int8" + model._use_int8 = True + model._use_nvfp4 = False + return count + +def _alpha_config_attr(name: str, default): + if not hasattr(alpha_config, name): + warnings.warn( + f"config module has no attribute {name!r} -- falling back to default " + f"{default!r}. If this is unexpected, check for a casing mismatch " + "or rename in your config.py." + ) + return default + return getattr(alpha_config, name) + +def _triton_warn_once(key, msg: str) -> None: + if key not in _triton_warned: + _triton_warned.add(key) + warnings.warn(f"[model.py/triton] {msg}", stacklevel=3) + + +def triton_status() -> dict: + """Which (op, key) pairs were verified and engaged vs fell back to eager.""" + return {"available": _HAS_TRITON, "verified": sorted(map(str, _TRITON_OK)), + "fell_back": sorted(map(str, _TRITON_BAD))} + + +def _as_list(out): + return list(out) if isinstance(out, (tuple, list)) else [out] + + +def _probe_precision(args, outputs=None): + """Lowest floating-point precision among probe inputs and fused outputs.""" + dts = [a.dtype for a in args if torch.is_tensor(a) and a.is_floating_point()] + if outputs is not None: + outs = outputs if isinstance(outputs, (list, tuple)) else [outputs] + dts += [o.dtype for o in outs if torch.is_tensor(o) and o.is_floating_point()] + for d in (torch.bfloat16, torch.float16): + if d in dts: + return d + return None + + +def _triton_close(a: torch.Tensor, b: torch.Tensor, what: str, grad: bool = False, + prec_dtype=None) -> None: + dt = prec_dtype if prec_dtype is not None else b.dtype + tol = {torch.float32: 1e-3, torch.float16: 1e-2, torch.bfloat16: 3e-2}.get(dt, 3e-2) + if grad: + tol *= 2 + af, bf = a.float(), b.float() + d = (af - bf).abs().max().item() + scale = bf.abs().max().item() + if not (d <= tol * scale + 1e-6): # also catches NaN + raise AssertionError( + f"{what} mismatch vs eager: max|diff|={d:.3e}, " + f"ref max|.|={scale:.3e}, tol={tol} (dt={dt})" + ) + + +def _triton_probe(tri_fn, ref_fn, probe_args) -> None: + """One-time check on a small slice: shape, dtype, values, grad flow, grad values.""" + def mk(a): + if torch.is_tensor(a): + b = a.detach().clone() + if a.requires_grad and a.is_floating_point() and torch.is_grad_enabled(): + b.requires_grad_(True) + return b + return a + a_t = [mk(a) for a in probe_args] + a_r = [mk(a) for a in probe_args] + out_t, out_r = _as_list(tri_fn(*a_t)), _as_list(ref_fn(*a_r)) + prec = _probe_precision(a_t, out_t) + if len(out_t) != len(out_r): + raise AssertionError(f"output arity {len(out_t)} != eager {len(out_r)}") + for i, (x, y) in enumerate(zip(out_t, out_r)): + if x.shape != y.shape: + raise AssertionError(f"output[{i}] shape {tuple(x.shape)} != eager {tuple(y.shape)}") + if x.dtype != y.dtype: + raise AssertionError(f"output[{i}] dtype {x.dtype} != eager {y.dtype}") + if y.is_floating_point(): # integer outputs (top-k indices / counts): shape+dtype only (ties may differ) + _triton_close(x, y, f"output[{i}]", prec_dtype=prec) + wrt = [i for i, a in enumerate(a_t) if torch.is_tensor(a) and a.requires_grad] + if not wrt: + return + out_t = [o for o in out_t if o.is_floating_point()] + out_r = [o for o in out_r if o.is_floating_point()] + if not all(o.requires_grad for o in out_t): + raise AssertionError("output has no grad_fn although inputs require grad (no backward?)") + gs = [] + for y in out_r: + n = y.numel() + gs.append(torch.cos(torch.arange(n, device=y.device, dtype=torch.float32)).reshape(y.shape).to(y.dtype)) + g_t = torch.autograd.grad(out_t, [a_t[i] for i in wrt], gs, allow_unused=True) + g_r = torch.autograd.grad(out_r, [a_r[i] for i in wrt], gs, allow_unused=True) + for i, x, y in zip(wrt, g_t, g_r): + if x is None and y is None: + continue + if x is None or y is None: + raise AssertionError(f"grad[arg {i}] present on only one side (triton={x is not None})") + _triton_close(x, y, f"grad[arg {i}]", grad=True, prec_dtype=prec) + + +@torch.compiler.disable +def _triton_attempt(op, key, tri_fn, ref_fn, args, probe_args): + """Run tri_fn(*args) or return None (caller then runs its eager path). + probe_args may be a callable (evaluated only on the first call for this key).""" + tag = (op, key) + if tag in _TRITON_BAD: + return None + try: + if tag not in _TRITON_OK: + _triton_probe(tri_fn, ref_fn, probe_args() if callable(probe_args) else probe_args) + _TRITON_OK.add(tag) + kernel._record_backend(f"triton_{op}_verified") + out = tri_fn(*args) + if torch.is_grad_enabled() and any(torch.is_tensor(a) and a.requires_grad for a in args): + if not all(o.requires_grad for o in _as_list(out) if o.is_floating_point()): + raise RuntimeError("output has no grad_fn although inputs require grad") + return out + except Exception as e: # noqa: BLE001 -- must never crash training + _TRITON_BAD.add(tag) + kernel._record_backend(f"triton_{op}_fallback:{type(e).__name__}") + _triton_warn_once( + tag, + f"triton {op} {key} disabled -> eager fallback ({type(e).__name__}: {e}). " + f"If this is a signature mismatch, fix the adapter next to _tri_{op}.", + ) + return None + + +# -- Adapters: the ONLY place that knows triton_kernel's call signatures -------- +def _tri_rmsnorm(x, weight, eps): + return _triton_kernel.rmsnorm(x, weight, eps) + + +def _tri_swiglu(gate, value): + return _triton_kernel.swiglu(gate, value) + + +def _tri_rope_qk(q, k, cos, sin, rope_dim): + return _triton_kernel.rope_qk(q, k, cos, sin, rope_dim) + + +def _tri_fused_ce(h, weight, t): + return _triton_kernel.fused_linear_cross_entropy(h, weight, t, ignore_index=-1) + + +def _rmsnorm_eager(x, weight, eps): + # weight cast to the input dtype: F.rms_norm falls back to a slow composite path on a dtype mismatch. + return _native_rms_norm(x, (x.shape[-1],), weight.to(x.dtype), eps) + + +def _swiglu_eager(gate, value): + return F.silu(gate) * value + + +class RMSNorm(nn.Module): + """Root Mean Square Layer Normalization (used in modern transformers)""" + def __init__(self, dim, eps=1e-5): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + self.dim = dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if _HAS_TRITON and x.is_cuda: + out = _triton_attempt( + "rmsnorm", (x.dtype, x.ndim, x.is_contiguous()), + _tri_rmsnorm, _rmsnorm_eager, + (x, self.weight, self.eps), (x[:1], self.weight, self.eps), + ) + if out is not None: + return out + return _native_rms_norm(x, (self.dim,), self.weight.to(x.dtype), self.eps) +def rotate_half(x): + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +@dataclass +class AttentionOutput: + """Uniform return type so callers can optionally inspect weights / cache cost.""" + + output: torch.Tensor + attention_weights: Optional[torch.Tensor] = None + kv_cache_bytes: Optional[int] = None + selected_indices: Optional[torch.Tensor] = None + + +def module_output_tensor(x): + + return x.output if isinstance(x, AttentionOutput) else x + + +def reshape_heads(x: torch.Tensor, num_heads: int, head_dim: int) -> torch.Tensor: + + b, t, _ = x.shape + return x.view(b, t, num_heads, head_dim).transpose(1, 2) + + +def merge_heads(x: torch.Tensor) -> torch.Tensor: + + b, h, t, d = x.shape + return x.transpose(1, 2).contiguous().view(b, t, h * d) + + +def expand_kv_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor: + + b, h_kv, t, d = x.shape + if h_kv == num_heads: + return x + if num_heads % h_kv != 0: + raise ValueError(f"num_heads ({num_heads}) must be divisible by kv heads ({h_kv})") + reps = num_heads // h_kv + return x.repeat_interleave(reps, dim=1) + + +@functools.lru_cache(maxsize=64) +def _rope_cache(seq_len: int, dim: int, device, dtype, base: float = 10000.0): + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim)) + t = torch.arange(seq_len, device=device, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos().to(dtype), emb.sin().to(dtype) + + +def build_rope_cache(seq_len: int, dim: int, base: float = 10000.0): + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + t = torch.arange(seq_len, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos(), emb.sin() + + +def apply_rope( + x: torch.Tensor, + positions: torch.Tensor | None = None, + rope_dim: int | None = None, + max_position: int | None = None, + cos_full: torch.Tensor | None = None, + sin_full: torch.Tensor | None = None, +) -> torch.Tensor: + + b, h, t, d = x.shape + rope_dim = d if rope_dim is None else rope_dim + + is_sequential = positions is None + if positions is None: + positions = torch.arange(t, device=x.device) + + if cos_full is None or sin_full is None: + if max_position is None: + if is_sequential: + max_position = t + else: + max_position = int(positions.max().item()) + 1 if positions.numel() > 0 else 1 + max_pos = int(max_position) + cos_full, sin_full = _rope_cache(max_pos, rope_dim, x.device, x.dtype) + + + if is_sequential: + if t > cos_full.shape[0]: + raise ValueError( + f"RoPE table too small: highest requested position=" + f"{t - 1}, table size={cos_full.shape[0]}. If this " + f"happened during generate(), set gen_headroom in config.py " + f"to however many extra positions past CONTEXT (block_size) " + f"generation needs -- see GPT.__init__'s gen_headroom " + f"wiring. This is separate from max_gen_tokens (a hard cap " + f"on tokens produced per generate() call, unrelated to " + f"RoPE table size). Default gen_headroom is 0 if unset." + ) + elif positions is not None and positions.numel() > 0: + highest_pos = int(positions.max().item()) + if highest_pos >= cos_full.shape[0]: + raise ValueError( + f"RoPE table too small: highest requested position=" + f"{highest_pos}, table size={cos_full.shape[0]}. If this " + f"happened during generate(), set gen_headroom in config.py " + f"to however many extra positions past CONTEXT (block_size) " + f"generation needs -- see GPT.__init__'s gen_headroom " + f"wiring. This is separate from max_gen_tokens (a hard cap " + f"on tokens produced per generate() call, unrelated to " + f"RoPE table size). Default gen_headroom is 0 if unset." + ) + + + if is_sequential: + cos = cos_full[:t] + sin = sin_full[:t] + else: + cos = cos_full[positions] + sin = sin_full[positions] + if cos.dtype != x.dtype: + cos = cos.to(x.dtype) + sin = sin.to(x.dtype) + cos = cos.unsqueeze(0).unsqueeze(0) + sin = sin.unsqueeze(0).unsqueeze(0) + + x_rot, x_pass = x[..., :rope_dim], x[..., rope_dim:] + x_rot = (x_rot * cos) + (rotate_half(x_rot) * sin) + return torch.cat([x_rot, x_pass], dim=-1) if x_pass.shape[-1] > 0 else x_rot + + +def _eager_rope_qk(q, k, cos, sin, rope_dim): + return ( + apply_rope(q, positions=None, rope_dim=rope_dim, cos_full=cos, sin_full=sin), + apply_rope(k, positions=None, rope_dim=rope_dim, cos_full=cos, sin_full=sin), + ) + + +def apply_rope_qk( + q: torch.Tensor, + k: torch.Tensor, + cos_full: torch.Tensor, + sin_full: torch.Tensor, + rope_dim: int, + positions: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + + b, h, t, d = q.shape + hk = k.shape[1] + + + if positions is None and _HAS_TRITON and q.is_cuda: + cos_t = cos_full[:t].to(q.dtype) + sin_t = sin_full[:t].to(q.dtype) + result = _triton_attempt( + "rope_qk", (q.dtype, h == hk), _tri_rope_qk, _eager_rope_qk, + (q, k, cos_t, sin_t, rope_dim), (q[:1], k[:1], cos_t, sin_t, rope_dim), + ) + if result is not None: + return tuple(result) + + q_rot = apply_rope(q, positions=positions, rope_dim=rope_dim, cos_full=cos_full, sin_full=sin_full) + k_rot = apply_rope(k, positions=positions, rope_dim=rope_dim, cos_full=cos_full, sin_full=sin_full) + return q_rot, k_rot + + +def chunked_cross_entropy( + hidden: torch.Tensor, + weight: torch.Tensor, + targets: torch.Tensor, + ignore_index: int = -1, + chunk_size: int = 4096, +) -> torch.Tensor: + """Triton fused linear-CE when verified, else the chunked eager path below.""" + if _HAS_TRITON and hidden.is_cuda and ignore_index == -1: + h = hidden.reshape(-1, hidden.size(-1)) + t = targets.reshape(-1) + valid = t != ignore_index + if bool(valid.any()): # one sync; the eager path syncs on n_valid too + + def _probe(): + # 192 valid + up to 64 ignored rows, so ignore_index handling is verified too. + vi = valid.nonzero(as_tuple=True)[0][:192] + ii = (~valid).nonzero(as_tuple=True)[0][:64] + idx = torch.cat([vi, ii]) + return (h[idx], weight, t[idx]) + + out = _triton_attempt( + "fused_ce", (h.dtype,), _tri_fused_ce, + lambda hh, ww, tt: _chunked_cross_entropy_eager(hh, ww, tt, ignore_index=-1), + (h, weight, t), _probe, + ) + if out is not None: + return out + if torch.is_grad_enabled() and (hidden.requires_grad or weight.requires_grad): + return _fused_linear_ce(hidden, weight, targets, ignore_index, chunk_size) + return _chunked_cross_entropy_eager(hidden, weight, targets, ignore_index, chunk_size) + + +def _chunked_cross_entropy_eager( + hidden: torch.Tensor, # (B, T, D) or (N, D) already flattened + weight: torch.Tensor, # (vocab_size, D) + targets: torch.Tensor, # (B, T) or (N,) already flattened + ignore_index: int = -1, + chunk_size: int = 4096, +) -> torch.Tensor: + + h = hidden.reshape(-1, hidden.size(-1)) # (N, D) + t = targets.reshape(-1) # (N,) + N = h.size(0) + total_loss = torch.zeros((), device=h.device, dtype=torch.float32) + n_valid = (t != ignore_index).sum().float() + # No data-dependent early return: an all-ignored batch still gets zero loss with a grad edge, + # because every checkpointed chunk stays in the graph. The clamp below avoids 0/0. + + def _chunk_loss(h_chunk: torch.Tensor, t_chunk: torch.Tensor) -> torch.Tensor: + # Recomputed during backward via checkpoint -- only the scalar loss is + # retained, not the (chunk_size, vocab) fp32 logit tensor. + logits_chunk = F.linear(h_chunk, weight).float() + return F.cross_entropy(logits_chunk, t_chunk, ignore_index=ignore_index, reduction="sum") + + for start in range(0, N, chunk_size): + end = min(start + chunk_size, N) + h_chunk = h[start:end] # (C, D) + t_chunk = t[start:end] # (C,) + # checkpoint recomputes the F.linear in backward; only the scalar + # loss_chunk is kept alive in the autograd graph between forward and + # backward, so the (C, vocab) fp32 logit tensor is never retained. + loss_chunk = torch.utils.checkpoint.checkpoint( + _chunk_loss, h_chunk, t_chunk, use_reentrant=False + ) + total_loss = total_loss + loss_chunk + return total_loss / n_valid.clamp(min=1) + + +class _FusedLinearCE(torch.autograd.Function): + """Sum of token cross-entropies of `F.linear(h, weight)`, chunked over tokens, with the gradients computed + DURING the forward chunk loop. Only grad_h (N, D) and grad_W (V, D) are kept -- never the (chunk, V) logits -- + and backward is just a scalar rescale. The checkpointed version recomputed every chunk's lm_head matmul in + backward (one extra N*D*V matmul, ~5% of a step at D=384/V=12288) and put 64 checkpoint regions in the graph. + The matmuls run under the ambient autocast, exactly as the autograd-generated backward would.""" + + @staticmethod + def forward(ctx, h, weight, t, ignore_index, chunk_size): + N = h.size(0) + total = torch.zeros((), device=h.device, dtype=torch.float32) + grad_h = torch.empty(h.shape, device=h.device, dtype=torch.float32) + grad_w = torch.zeros(weight.shape, device=h.device, dtype=torch.float32) + for start in range(0, N, chunk_size): + end = min(start + chunk_size, N) + hc, tc = h[start:end], t[start:end] + valid = (tc != ignore_index).unsqueeze(1) + safe_t = torch.where(valid, tc.unsqueeze(1), torch.zeros_like(tc.unsqueeze(1))) + logits = F.linear(hc, weight).float() + lse = torch.logsumexp(logits, dim=-1, keepdim=True) + total = total + ((lse - logits.gather(1, safe_t)) * valid).sum() + p = torch.exp(logits - lse) # softmax + p.scatter_add_(1, safe_t, -valid.float()) # - onehot(target) on valid rows + p = p * valid # ignored rows contribute nothing + grad_h[start:end] = torch.mm(p, weight).float() + grad_w += torch.mm(p.t(), hc).float() + ctx.h_dtype, ctx.w_dtype = h.dtype, weight.dtype + ctx.save_for_backward(grad_h, grad_w) + return total + + @staticmethod + def backward(ctx, grad_out): + grad_h, grad_w = ctx.saved_tensors + return (grad_h * grad_out).to(ctx.h_dtype), (grad_w * grad_out).to(ctx.w_dtype), None, None, None + + +def _fused_linear_ce(hidden, weight, targets, ignore_index: int = -1, chunk_size: int = 4096) -> torch.Tensor: + h = hidden.reshape(-1, hidden.size(-1)) + t = targets.reshape(-1) + n_valid = (t != ignore_index).sum().float() + return _FusedLinearCE.apply(h, weight, t, ignore_index, chunk_size) / n_valid.clamp(min=1) + + +def make_causal_mask(q_len: int, k_len: int, device, q_offset: int = 0) -> torch.Tensor: + + q_idx = (torch.arange(q_len, device=device).view(q_len, 1) + q_offset) + k_idx = torch.arange(k_len, device=device).view(1, k_len) + return torch.where(k_idx <= q_idx, torch.zeros(1, device=device), torch.full((1,), float("-inf"), device=device)) + + +def make_sliding_window_causal_mask(q_len: int, k_len: int, window_size: int, device, q_offset: int = 0) -> torch.Tensor: + + q_idx = (torch.arange(q_len, device=device).view(q_len, 1) + q_offset) + k_idx = torch.arange(k_len, device=device).view(1, k_len) + visible = (k_idx <= q_idx) & (k_idx > q_idx - window_size) + return torch.where(visible, torch.zeros(1, device=device), torch.full((1,), float("-inf"), device=device)) + + + +def masked_softmax(scores: torch.Tensor, mask: torch.Tensor | None, dim: int = -1) -> torch.Tensor: + if mask is not None: + scores = scores + mask.to(scores.dtype) + + return torch.softmax(scores.float(), dim=dim).to(scores.dtype) + + +def estimate_kv_cache_bytes( + batch_size, seq_len, num_kv_heads, head_dim, dtype, + compression_ratio: int = 1, window_size: int | None = None, +) -> int: + + elem_size = torch.zeros(1, dtype=dtype).element_size() + effective_len = max(1, -(-seq_len // compression_ratio)) # ceil div + if window_size is not None: + effective_len = min(effective_len, window_size) + return int(batch_size * num_kv_heads * effective_len * head_dim * elem_size * 2) + + + + +class KVCache: + """Per-layer K/V buffers for incremental decoding. + + Buffers are allocated lazily on first write (dtype/device/batch follow k) and + written in place, so a decode step costs O(T) attention reads and no torch.cat. + Invariant: after a forward, `pos` == number of tokens already cached. + + `extra` holds shared non-layer state such as the Engram hasher history. + Per-layer Engram short-convolution history lives in each layer dictionary. + """ + + def __init__(self, n_layers: int, capacity: int): + self.capacity = int(capacity) + self.pos = 0 + self.extra: dict = {} + self._layers = [{"capacity": self.capacity} for _ in range(n_layers)] + + def layer(self, i: int) -> dict: + return self._layers[i] + + def reset(self) -> None: + self.pos = 0 # stale data past pos is never read + self.extra.clear() + for layer in self._layers: + layer.pop("sc_hist_x", None) + layer.pop("sc_hist_doc", None) + + +# -- Static architectural constants ------------------------------------------ +# XSA, QK-Norm and Gated Attention are always on. No config flags: the +# checkpoint shape is fully determined by (n_embd, n_head). +# +# Attn-gate input window: how many dims of the post-norm block input the +# per-head gate reads. Smaller = fewer params. 16 is the speedrun value; +# 0 means "use full n_embd" (dense gate). +ATTN_GATE_WINDOW_DIMS = 16 + +# XSA smoothing constant. The self-value projection is (c.v) v / (|v|^2 + XSA_EPS), which equals the +# hard-normalised (c.v_hat) v_hat whenever |v| >> sqrt(XSA_EPS) (~1e-3), but fades to zero as |v| -> 0. +# The old v/|v|.clamp(1e-6) form had a 1/|v| gradient gain: one near-zero value vector next to a large +# context gave a gradient spike (measured ~200x at |v|=1e-5) that, under global grad clipping, shrank +# the whole step. +XSA_EPS = 1e-6 + + +class DenseMHA(nn.Module): + + + def __init__( + self, + hidden_size: int, + num_heads: int, + head_dim: int | None = None, + num_kv_heads: int | None = None, + dropout: float = 0.0, + use_rope: bool = True, + causal: bool = True, + bias: bool = True, + max_seq_len: int = 4096, + window_size: int | None = None, + tp_size: int = 1, + rope_dim: int | None = None, + ) -> None: + super().__init__() + if head_dim is None: + if hidden_size % num_heads != 0: + raise ValueError("hidden_size must be divisible by num_heads when head_dim is omitted") + head_dim = hidden_size // num_heads + num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads + if num_heads % num_kv_heads != 0: + raise ValueError(f"num_heads ({num_heads}) must be divisible by num_kv_heads ({num_kv_heads})") + if tp_size > 1: + raise NotImplementedError( + "tensor_parallel_size > 1 is not implemented: per-rank weights are " + "independently initialized (never sharded from a canonical weight), " + "so DDP gradient averaging corrupts them. Set tensor_parallel_size=1. " + "Real TP needs: init full weights once, slice per rank, and a " + "skip_out_proj guard that honors tp_size." + ) + if num_heads % tp_size != 0: + raise ValueError( + f"tensor_parallel_size ({tp_size}) must divide num_heads ({num_heads}) " + f"for this layer -- reduce tensor_parallel_size or increase n_head." + ) + + local_num_heads = num_heads // tp_size + + if num_kv_heads < tp_size: + local_num_kv_heads = num_kv_heads + else: + local_num_kv_heads = num_kv_heads // tp_size if num_kv_heads % tp_size == 0 else num_kv_heads + + self.hidden_size = hidden_size + self.num_heads = local_num_heads + self.num_kv_heads = local_num_kv_heads + self.head_dim = head_dim + self.inner_dim = local_num_heads * head_dim + self.use_rope = use_rope + self.causal = causal + if dropout > 0.0 or not causal: + raise ValueError("FlexAttention-only kernel: attention dropout and non-causal attention are not " + f"supported (got dropout={dropout}, causal={causal}). Set dropout=0.0.") + self.tp_size = tp_size + self.tp_group = None + self.window_size = window_size + + self.rope_dim = head_dim if rope_dim is None else rope_dim + + + + self._q_width = self.inner_dim + self._k_width = self.num_kv_heads * head_dim + self._v_width = self.num_kv_heads * head_dim + self.qkv_proj = nn.Linear(hidden_size, self._q_width + self._k_width + self._v_width, bias=bias) + self.out_proj = nn.Linear(self.inner_dim, hidden_size, bias=bias) + self.dropout = nn.Dropout(dropout) + + # -- Always-on QK-Norm -------------------------------------------------- + # Per-head RMSNorm over head_dim on q and k, before the rope split. + # Bounds attention-logit scale so softmax doesn't saturate. + self.q_norm = RMSNorm(self.head_dim, eps=1e-6) + self.k_norm = RMSNorm(self.head_dim, eps=1e-6) + + # -- Always-on per-head XSA strength gate ------------------------------- + # tanh(alpha) in (-1, 1). Zero-init -> gate=0 -> step-0 forward is + # numerically identical to vanilla SA. + self.xsa_alpha = nn.Parameter(torch.zeros(self.num_heads)) + self._xsa_v_reps = local_num_heads // local_num_kv_heads + + # -- Always-on Gated Attention (Qiu et al., NeurIPS 2025) --------------- + # Per-head sigmoid gate on the SDPA output, before XSA/out_proj. Reads + # the first ATTN_GATE_WINDOW_DIMS dims of the post-norm block input. + # Zero-init -> 2*sigmoid(0) = 1.0 -> identity at step 0. + # NOTE: GPT.__init__ re-zeroes this after self.apply(_init_weights), + # which would otherwise overwrite it with N(0, 0.02). + self.attn_gate_window = min(ATTN_GATE_WINDOW_DIMS, hidden_size) or hidden_size + self.attn_gate_proj = nn.Linear(self.attn_gate_window, self.num_heads, bias=False) + nn.init.zeros_(self.attn_gate_proj.weight) + + if use_rope: + rope_cos, rope_sin = build_rope_cache(max_seq_len, self.rope_dim) + # Store in the training activation dtype to avoid casting the tables on every call. + self.register_buffer("rope_cos", rope_cos.to(torch.bfloat16), persistent=False) + self.register_buffer("rope_sin", rope_sin.to(torch.bfloat16), persistent=False) + + # XSA pos-0 guard mask: registered once at init so forward() never + # allocates a fresh (T,) tensor. Non-persistent: re-created from + # max_seq_len on load, costs nothing in the state_dict. + self.register_buffer( + "_xsa_pos_mask", + torch.ones(max_seq_len), + persistent=False, + ) + # Slot 0 is always masked off (position-0 has no context beyond itself). + self._xsa_pos_mask[0] = 0.0 + + def forward( + self, + hidden_states: torch.Tensor, + output_attentions: bool = False, + return_kv_cache_estimate: bool = False, + skip_out_proj: bool = False, + positions: torch.Tensor | None = None, + kv_cache: dict | None = None, + cache_pos: int = 0, + attn_masks: dict | None = None, + xsa_pos_mask: torch.Tensor | None = None, + ) -> torch.Tensor | AttentionOutput: + batch_size, seq_len, _ = hidden_states.shape + if kv_cache is not None: + if output_attentions or positions is not None or not self.causal: + raise ValueError("kv_cache requires causal attention, positions=None, output_attentions=False") + if cache_pos > 0 and seq_len != 1: + raise NotImplementedError("cached forward: one full prefill at position 0, then one token per step") + if cache_pos + seq_len > kv_cache["capacity"]: + raise ValueError(f"KV cache capacity {kv_cache['capacity']} exceeded at position {cache_pos + seq_len}") + + qkv = self.qkv_proj(hidden_states) + if self.use_rope: + if kv_cache is not None: + rope_c = self.rope_cos[cache_pos:cache_pos + seq_len] + rope_s = self.rope_sin[cache_pos:cache_pos + seq_len] + else: + rope_c, rope_s = self.rope_cos, self.rope_sin + + q_flat, k_flat, v_flat = qkv.split([self._q_width, self._k_width, self._v_width], dim=-1) + q = reshape_heads(q_flat, self.num_heads, self.head_dim) + k = reshape_heads(k_flat, self.num_kv_heads, self.head_dim) + v = reshape_heads(v_flat, self.num_kv_heads, self.head_dim) + + # QK-Norm (always on): per-head RMSNorm over head_dim, before rope. + q = self.q_norm(q) + k = self.k_norm(k) + + if self.use_rope: + q, k = apply_rope_qk( + q, + k, + cos_full=rope_c, + sin_full=rope_s, + rope_dim=self.rope_dim, + positions=positions, + ) + + if kv_cache is not None: + _kb = kv_cache.get("k") + if _kb is None or (cache_pos == 0 and ( + _kb.shape[0] != batch_size or _kb.dtype != k.dtype or _kb.device != k.device)): + cap = kv_cache["capacity"] + kv_cache["k"] = k.new_empty(batch_size, self.num_kv_heads, cap, self.head_dim) + kv_cache["v"] = v.new_empty(batch_size, self.num_kv_heads, cap, self.head_dim) + kv_cache["k"][:, :, cache_pos:cache_pos + seq_len] = k + kv_cache["v"][:, :, cache_pos:cache_pos + seq_len] = v + + if output_attentions: + + if self.num_kv_heads != self.num_heads: + reps = self.num_heads // self.num_kv_heads + k_eager = k.repeat_interleave(reps, dim=1) + v_eager = v.repeat_interleave(reps, dim=1) + else: + k_eager, v_eager = k, v + # fp32 logits: bf16 rounding of raw scores (|s|~10 -> ~0.06 abs error) is amplified by exp(). + scores = torch.matmul(q.float() * (1.0 / math.sqrt(self.head_dim)), k_eager.float().transpose(-2, -1)) + mask = None + if self.window_size is not None: + mask = make_sliding_window_causal_mask( + seq_len, seq_len, self.window_size, hidden_states.device + ).view(1, 1, seq_len, seq_len) + elif self.causal: + mask = make_causal_mask(seq_len, seq_len, hidden_states.device).view(1, 1, seq_len, seq_len) + weights = masked_softmax(scores, mask, dim=-1).to(q.dtype) + + weights = self.dropout(weights) + context = torch.matmul(weights, v_eager) + elif kv_cache is not None and cache_pos > 0: + # Single-token decode: the new token sees every cached key (or the last + # `window_size` of them for sliding-window layers), so no mask is needed. + # GQA without copying K/V: heads of one KV group become the "query rows". + end = cache_pos + 1 + start = 0 if self.window_size is None else max(0, end - self.window_size) + k_all = kv_cache["k"][:, :, start:end].to(q.dtype) + v_all = kv_cache["v"][:, :, start:end].to(q.dtype) + reps = self.num_heads // self.num_kv_heads + q_g = q.reshape(batch_size, self.num_kv_heads, reps, self.head_dim) + ctx = F.scaled_dot_product_attention(q_g, k_all, v_all) + context = ctx.reshape(batch_size, self.num_heads, 1, self.head_dim) + weights = None + else: + context = kernel.fused_attention(q, k, v, attn_masks[self.window_size]) + weights = None + + # ── Gated Attention (Qiu et al., NeurIPS 2025 Best Paper) ─────────────── + # Per-head sigmoid gate modulates the SDPA output before XSA/out_proj. + # Gate input: first `attn_gate_window` dims of the post-norm block input + # (hidden_states). Sparse window keeps the gate head lightweight. + # 2*sigmoid(z) is used so the zero-init projection gives gate=1 (identity) + # at step 0; the head learns to suppress or amplify over training. + gate_in = hidden_states[..., :self.attn_gate_window] # (B, T, gw) + gate = self.attn_gate_proj(gate_in) # (B, T, num_heads) + + gate = 2.0 * torch.sigmoid(gate) # (0, 2), =1.0 at init + gate = gate.transpose(1, 2).unsqueeze(-1) # (B, H, T, 1) + context = context * gate + # ──────────────────────────────────────────────────────────────────────── + + if context is not None: + v_xsa = v + if self._xsa_v_reps > 1: + # GQA/MQA-aware broadcast: mathematically identical to + # v.repeat_interleave(self._xsa_v_reps, dim=1) (verified via + # torch.equal), but avoids materializing a full copy of v. + Bv, Hkv, Tv, Dv = v.shape + v_xsa = ( + v.reshape(Bv, Hkv, 1, Tv, Dv) + .expand(Bv, Hkv, self._xsa_v_reps, Tv, Dv) + .reshape(Bv, Hkv * self._xsa_v_reps, Tv, Dv) + ) + + # Accumulate in fp32 without materializing a full [B,H,T,D] fp32 copy. + v_sq = (v_xsa * v_xsa).sum(dim=-1, keepdim=True, dtype=torch.float32) + v_dot = (context * v_xsa).sum(dim=-1, keepdim=True, dtype=torch.float32) + proj_coeff = torch.tanh(v_dot / (v_sq + XSA_EPS)).to(context.dtype) + # Learnable per-head gate (modded-nanoGPT style): tanh(alpha) in (-1, 1). + # alpha starts at 0 -> gate=0 -> forward pass is numerically identical + # to vanilla SA at step 0. Each head learns how much (if any) self-value + # component to remove, rather than a hard, deterministic subtraction. + gate = torch.tanh(self.xsa_alpha).to(context.dtype).view(1, -1, 1, 1) + correction = gate * proj_coeff * v_xsa + + # Position-0 guard: causal attention at position 0 has exactly one + # reachable key (itself), so context[0] == v[0] identically, at every + # step of training. XSA's projection therefore removes 100% of the + # context for that head, for any gate > 0. This is structural, not + # a training artifact -- there is no "context minus self" at + # position 0 because there is no context other than self there. + # Mask the correction off so the first token always keeps a + # full-strength value contribution. + T = context.shape[2] + # Slice the pre-built buffer; no allocation per forward. Applied for T == 1 as well: + # a 1-token input IS position 0 and must be masked exactly as it is in training. + mask = self._xsa_pos_mask[cache_pos:cache_pos + T].to(dtype=context.dtype) + if xsa_pos_mask is not None: + mask = mask.unsqueeze(0) * xsa_pos_mask.to(dtype=context.dtype) + else: + mask = mask.unsqueeze(0) + mask = mask.unsqueeze(1).unsqueeze(-1) + correction = correction * mask + + context = context - correction + + merged_context = merge_heads(context) + # ───────────────────────────────────────────────────────────────────── + + if skip_out_proj and self.tp_size == 1 and not output_attentions and not return_kv_cache_estimate: + return merged_context + + output = self.out_proj(merged_context) + + + if self.tp_size > 1: + if self.tp_group is None: + raise RuntimeError( + "DenseMHA was built with tp_size>1 but tp_group was never " + "set -- call GPT.set_tp_group(group) after construction, " + "before the first forward pass." + ) + torch.distributed.all_reduce(output, op=torch.distributed.ReduceOp.SUM, group=self.tp_group) + + if output_attentions or return_kv_cache_estimate: + cache_bytes = None + if return_kv_cache_estimate: + cache_bytes = estimate_kv_cache_bytes( + batch_size, seq_len, self.num_kv_heads, self.head_dim, + hidden_states.dtype, window_size=self.window_size, + ) + return AttentionOutput(output=output, attention_weights=weights if output_attentions else None, kv_cache_bytes=cache_bytes) + return output + + + +class SwiGLU_FFN(nn.Module): + """ + SwiGLU Feed Forward Network. + + + """ + + def __init__(self, d_model, ffn_mult=8/3, dropout=0.0, tp_size: int = 1): + super().__init__() + + raw_hidden = d_model * ffn_mult + hidden_dim = int(((raw_hidden + 63) // 64) * 64) + if hidden_dim % tp_size != 0: + raise ValueError( + f"tensor_parallel_size ({tp_size}) must divide the FFN hidden_dim " + f"({hidden_dim} = round({d_model} * {ffn_mult}) -> next mult of 64) " + f"-- adjust ffn_mult or tensor_parallel_size." + ) + local_hidden_dim = hidden_dim // tp_size + self.tp_size = tp_size + self.tp_group = None # set post-construction via GPT.set_tp_group() + self.W_gate = nn.Linear(d_model, local_hidden_dim, bias=False) + self.W_value = nn.Linear(d_model, local_hidden_dim, bias=False) + self.W_out = nn.Linear(local_hidden_dim, d_model, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x): + if (self.tp_size == 1 and type(self.W_gate.weight) is nn.Parameter + and type(self.W_value.weight) is nn.Parameter): + # One GEMM for both projections: [W_gate; W_value] stacked on the output dim, then split. + # Parameters and state_dict keys are unchanged; only the cat (2*H*D weights) is extra work. + gate_raw, value = F.linear( + x, torch.cat((self.W_gate.weight, self.W_value.weight), dim=0) + ).chunk(2, dim=-1) + else: + gate_raw = self.W_gate(x) + value = self.W_value(x) + + hidden = None + if _HAS_TRITON and gate_raw.is_cuda: + hidden = _triton_attempt( + "swiglu", (gate_raw.dtype,), _tri_swiglu, _swiglu_eager, + (gate_raw.contiguous(), value.contiguous()), (gate_raw[:1].contiguous(), value[:1].contiguous()), + ) + if hidden is None: + hidden = F.silu(gate_raw) * value + + hidden = self.dropout(hidden) + out = self.W_out(hidden) + if self.tp_size > 1: + if self.tp_group is None: + raise RuntimeError( + "SwiGLU_FFN was built with tp_size>1 but tp_group was never " + "set -- call GPT.set_tp_group(group) after construction, " + "before the first forward pass." + ) + torch.distributed.all_reduce(out, op=torch.distributed.ReduceOp.SUM, group=self.tp_group) + return out # Return delta only +class AryaSparseMoE(nn.Module): + """ + DeepSeek-style Sparse Mixture of Experts. + Expects input to be ALREADY normalized (Pre-LN). + Returns the combined FFN delta. + """ + + def __init__(self, d_model, n_experts=8, n_shared=1, top_k=2, ffn_mult=8/3, dropout=0.0, + use_capacity_routing=False, capacity_factor=1.25, fixed_capacity: int | None = None): + super().__init__() + self.n_experts = n_experts + self.n_routed = n_experts - n_shared + self.top_k = top_k + self.d_model = d_model + self.use_capacity_routing = use_capacity_routing + self.capacity_factor = capacity_factor + self.fixed_capacity = fixed_capacity + + if self.n_routed < top_k: + raise ValueError( + f"AryaSparseMoE config error: n_experts={n_experts} - n_shared={n_shared} " + f"= {self.n_routed} routed experts, but top_k={top_k} routing needs at " + f"least {top_k} routed experts to choose from." + ) + + self.shared_expert = SwiGLU_FFN(d_model, ffn_mult, dropout=dropout) + self.routed_experts = nn.ModuleList([ + SwiGLU_FFN(d_model, ffn_mult, dropout=dropout) + for _ in range(self.n_routed) + ]) + + self.router = nn.Linear(d_model, self.n_routed, bias=False) + self.expert_bias = nn.Parameter(torch.zeros(self.n_routed), requires_grad=False) + self.aux_loss_weight = float(getattr(alpha_config, 'moe_aux_loss_weight', 0.01)) + self.last_aux_loss = torch.zeros((), dtype=torch.float32) + + def forward(self, x): + B, T, D = x.shape + + + shared_out = self.shared_expert(x) + + + router_logits = self.router(x) + if self.training and torch.is_grad_enabled(): + biased_logits = router_logits.float() + self.expert_bias.float() + _, topk_indices = biased_logits.topk(self.top_k, dim=-1) + full_probs = F.softmax(router_logits.float(), dim=-1) + topk_probs = full_probs.gather(-1, topk_indices).to(router_logits.dtype) + topk_probs = topk_probs / topk_probs.sum(dim=-1, keepdim=True).clamp_min(1e-9) + one_hot = F.one_hot(topk_indices, num_classes=self.n_routed).float() + f = one_hot.mean(dim=(0, 1)) + P = full_probs.mean(dim=(0, 1)) + self.last_aux_loss = self.aux_loss_weight * self.n_routed * (f * P).sum() + # Do not mutate expert_bias here. Checkpoint recomputation runs + # this forward a second time, so an in-forward update would apply + # twice per optimizer step and make routing state nondeterministic. + else: + full_probs = F.softmax(router_logits.float(), dim=-1) + topk_probs, topk_indices = full_probs.topk(self.top_k, dim=-1) + topk_probs = (topk_probs / topk_probs.sum(dim=-1, keepdim=True).clamp_min(1e-9)).to(router_logits.dtype) + self.last_aux_loss = torch.zeros((), device=x.device, dtype=torch.float32) + gate_weights = topk_probs + + x_flat = x.view(B * T, D) + gate_flat = gate_weights.view(B * T, self.top_k) + idx_flat = topk_indices.view(B * T, self.top_k) + + if self.use_capacity_routing: + routed_out = self._forward_capacity_routed(x_flat, gate_flat, idx_flat) + routed_out = routed_out.view(B, T, D) + else: + routed_out = self._eager_loop(x_flat, gate_flat, idx_flat).view(B, T, D) + + out = shared_out + routed_out + return out, self.last_aux_loss + + def _eager_loop(self, x_flat, gate_flat, idx_flat): + """Per-expert Python loop -- the reference implementation and fallback.""" + routed_out = torch.zeros_like(x_flat) + expert_masks = F.one_hot(idx_flat, num_classes=self.n_routed).bool().any(dim=1) + for expert_id in range(self.n_routed): + expert = self.routed_experts[expert_id] + token_mask = expert_masks[:, expert_id] + selected_tokens = x_flat[token_mask] + if selected_tokens.shape[0] == 0: + continue + expert_out = expert(selected_tokens) + weight = ( + (idx_flat[token_mask] == expert_id).float() * gate_flat[token_mask] + ).sum(dim=-1) + routed_out[token_mask] += weight.unsqueeze(-1) * expert_out + return routed_out + + def _forward_capacity_routed(self, x_flat, gate_flat, idx_flat): + n_tokens, D = x_flat.shape + n_routed, capacity = self.n_routed, self._capacity(n_tokens) + + flat_expert_ids = idx_flat.reshape(-1) + flat_gates = gate_flat.reshape(-1) + flat_token_ids = torch.arange(n_tokens, device=x_flat.device).unsqueeze(1).expand(-1, self.top_k).reshape(-1) + + + one_hot = F.one_hot(flat_expert_ids, num_classes=n_routed).to(torch.float32) + position_in_expert = (one_hot.cumsum(dim=0) * one_hot).sum(dim=1).long() - 1 + keep = position_in_expert < capacity + safe_position = position_in_expert.clamp(min=0, max=capacity - 1) + + flat_slot = flat_expert_ids * capacity + safe_position + + gathered_x = x_flat[flat_token_ids] * keep.unsqueeze(-1).to(x_flat.dtype) + dispatch = torch.zeros(n_routed * capacity, D, device=x_flat.device, dtype=x_flat.dtype) + dispatch.index_add_(0, flat_slot, gathered_x) + dispatch = dispatch.view(n_routed, capacity, D) + + expert_out = torch.stack( + [expert(dispatch[e]) for e, expert in enumerate(self.routed_experts)], + dim=0, + ).view(n_routed * capacity, D) + + contrib = expert_out[flat_slot] * keep.unsqueeze(-1).to(x_flat.dtype) * flat_gates.unsqueeze(-1) + routed_out = torch.zeros(n_tokens, D, device=x_flat.device, dtype=x_flat.dtype) + routed_out.index_add_(0, flat_token_ids, contrib) + return routed_out + + def _capacity(self, n_tokens) -> int: + if self.fixed_capacity is not None: + return self.fixed_capacity + total_assignments = n_tokens * self.top_k + avg_load = -(-total_assignments // self.n_routed) + numer = int(round(self.capacity_factor * 1000)) + capacity = -(-(avg_load * numer) // 1000) + return capacity + 1 + + +def build_attention_layers( + num_layers: int, + hidden_size: int, + num_heads: int, + dropout: float = 0.0, + bias: bool = True, + num_kv_heads: int | None = None, + pattern: str = "dense", + max_seq_len: int = 4096, + sliding_window_size: int = 256, + tp_size: int = 1, + rope_dim: int | None = None, +): + + if num_kv_heads is None: + num_kv_heads = num_heads + + layers = [] + layer_types = [] + + if pattern == "dense": + for _ in range(num_layers): + layers.append(DenseMHA( + hidden_size=hidden_size, num_heads=num_heads, num_kv_heads=num_kv_heads, + dropout=dropout, bias=bias, max_seq_len=max_seq_len, + window_size=None, tp_size=tp_size, rope_dim=rope_dim, + )) + layer_types.append("dense") + + elif pattern == "pyramid_swa": + + if num_kv_heads <= 1: + pyramid_kv_heads = [1, 1, 1] + else: + import math as _math + + mqa_floor = max(2, _math.ceil(num_kv_heads / 4)) + mqa_floor = min(mqa_floor, num_kv_heads) # never exceed full heads + pyramid_kv_heads = sorted(set([ + mqa_floor, + max(mqa_floor, _math.ceil(num_kv_heads / 2)), + num_kv_heads, + ])) + + while len(pyramid_kv_heads) < 3: + pyramid_kv_heads.append(num_kv_heads) + pyramid_kv_heads = pyramid_kv_heads[:3] + + for i in range(num_layers): + pos_in_group = i % 4 + is_last_layer = (i == num_layers - 1) + + if pos_in_group == 3 or is_last_layer: + layers.append(DenseMHA( + hidden_size=hidden_size, num_heads=num_heads, num_kv_heads=num_kv_heads, + dropout=dropout, bias=bias, max_seq_len=max_seq_len, + window_size=None, tp_size=tp_size, rope_dim=rope_dim, + )) + layer_types.append("global") + else: + kv_heads_here = pyramid_kv_heads[pos_in_group] + layers.append(DenseMHA( + hidden_size=hidden_size, num_heads=num_heads, num_kv_heads=kv_heads_here, + dropout=dropout, bias=bias, max_seq_len=max_seq_len, + window_size=sliding_window_size, tp_size=tp_size, rope_dim=rope_dim, + )) + layer_types.append("swa") + elif pattern == "mod": + + for i in range(num_layers): + layers.append(DenseMHA( + hidden_size=hidden_size, num_heads=num_heads, num_kv_heads=num_kv_heads, + dropout=dropout, bias=bias, max_seq_len=max_seq_len, + window_size=None, tp_size=tp_size, rope_dim=rope_dim, + )) + is_full_dense = (i % 4 == 3) or (i == num_layers - 1) + layer_types.append("mod_dense" if is_full_dense else "mod") + else: + raise ValueError( + f"Unknown attention pattern {pattern!r} -- supported patterns: " + f"'dense', 'pyramid_swa', 'mod'. " + f"(BigBird/conv-hybrid/LSA/HCA patterns have been removed.)" + ) + + return nn.ModuleList(layers), layer_types + + +class AryaBlock(nn.Module): + """Single transformer block: Layer Norm → Attention → MLP (with residual connections). + + `attn_module` is DenseMHA, + pre-built by GPT via build_attention_layers. Every layer is fully + standalone -- no cross-layer attention/KV sharing of any kind. + `role` is purely a label ("dense" | "lsa") for inspection/debugging. + """ + + def __init__(self, config, attn_module: nn.Module, role: str = "dense"): + super().__init__() + self.ln_1 = RMSNorm(config.n_embd, eps=1e-5) + # Pre-LN for the MLP branch: ensure router and experts see the + # identical normalized tensor the MLPs expect. + self.ln_2 = RMSNorm(config.n_embd, eps=1e-5) + self.attn = attn_module + self.role = role # "dense" | "lsa" + + + self.use_moe = getattr(config, 'use_moe', False) + if self.use_moe: + if getattr(config, 'tp_size', 1) > 1: + raise ValueError( + "tensor_parallel_size > 1 is not supported together with use_moe=True -- " + "MoE needs its own (unimplemented here) expert-parallel sharding scheme, " + "column/row-parallel head splitting does not apply to expert routing. " + "Set use_moe=False or tensor_parallel_size=1." + ) + self.ffn = AryaSparseMoE( + d_model=config.n_embd, + n_experts=config.n_experts, + n_shared=config.n_shared, + top_k=getattr(config, 'moe_top_k', 2), + ffn_mult=config.ffn_mult, + dropout=config.dropout, + use_capacity_routing=False, + fixed_capacity=getattr(config, 'moe_fixed_capacity', None), + ) + else: + self.ffn = SwiGLU_FFN(config.n_embd, config.ffn_mult, dropout=config.dropout, tp_size=getattr(config, 'tp_size', 1)) + self._ffn_residual = self._moe_residual if self.use_moe else self._dense_residual + + def forward(self, x, kv_cache=None, cache_pos=0, attn_masks=None, + xsa_pos_mask=None): + """Every attention layer (DenseMHA, LSA) is standalone and + computes its own K/V from `x` alone -- + no external_kv, no cross-layer sharing, no produced_kv to pass on. + """ + normed = self.ln_1(x) + if kv_cache is None: + attn_delta = module_output_tensor(self.attn( + normed, attn_masks=attn_masks, xsa_pos_mask=xsa_pos_mask)) + else: + attn_delta = module_output_tensor(self.attn(normed, kv_cache=kv_cache, cache_pos=cache_pos, + attn_masks=attn_masks, xsa_pos_mask=xsa_pos_mask)) + x = x + attn_delta + ffn_in = self.ln_2(x) + return self._ffn_residual(x, ffn_in) + + def _dense_residual(self, x, ffn_in): + return x + self.ffn(ffn_in) + + def _moe_residual(self, x, ffn_in): + ffn_delta, moe_aux = self.ffn(ffn_in) + return x + ffn_delta, moe_aux + + +class MoDBlock(nn.Module): + + def __init__(self, config, inner: nn.Module, capacity: float = 0.125): + super().__init__() + self.inner = inner + self.capacity = capacity + self.keep_target = float(capacity) + self._budget_coeff = float(getattr(config, 'mod_budget_coeff', 0.05)) + self._entropy_coeff = float(getattr(config, 'mod_gate_entropy_coeff', 0.0)) + self.hard_eval = bool(getattr(config, 'mod_hard_eval', False)) + + self._router_norm = RMSNorm(config.n_embd, eps=1e-5) + self.router = nn.Linear(config.n_embd, 1, bias=True) + self.last_gate_aux = None + self._call_inner = self._call_moe if self.inner.use_moe else self._call_dense + + def forward(self, x, kv_cache=None, cache_pos=0, attn_masks=None, + xsa_pos_mask=None): + gate = torch.sigmoid(self.router(self._router_norm(x))) + + aux = torch.zeros((), device=x.device, dtype=torch.float32) + if self.training: + g = gate.float() + if self._budget_coeff > 0.0: + aux = aux + self._budget_coeff * (g.mean() - self.keep_target) ** 2 + if self._entropy_coeff > 0.0: + gc = g.clamp(1e-6, 1.0 - 1e-6) + aux = aux + self._entropy_coeff * ( + -(gc * gc.log() + (1 - gc) * (1 - gc).log()).mean() + ) + # BUGFIX: `aux` must stay attached to the graph -- returning + # aux.detach() made the MoD budget/entropy regularizer contribute + # exactly zero gradient to the router, i.e. it was inert. + self.last_gate_aux = aux.detach() # detached copy, for logging only + else: + aux = torch.zeros((), device=x.device, dtype=torch.float32) + self.last_gate_aux = aux + + kw = ({"attn_masks": attn_masks, "xsa_pos_mask": xsa_pos_mask} if kv_cache is None else + {"kv_cache": kv_cache, "cache_pos": cache_pos, "attn_masks": attn_masks, + "xsa_pos_mask": xsa_pos_mask}) + inner_out, moe_aux = self._call_inner(x, kw) + if gate.dtype == x.dtype == inner_out.dtype: + out = torch.lerp(x, inner_out, gate) # == x + gate * (inner_out - x), single fused kernel + else: + out = x + gate * (inner_out - x) + return (out, moe_aux), aux + + def _call_dense(self, x, kw): + return self.inner(x, **kw), torch.zeros((), device=x.device, dtype=torch.float32) + + def _call_moe(self, x, kw): + return self.inner(x, **kw) + + def extra_repr(self) -> str: + return ( + f"gate_based=True,keep_target={self.keep_target}," + f"budget_coeff={self._budget_coeff},entropy_coeff={self._entropy_coeff}" + ) + + +class MTPHead(nn.Module): + """Residual shared-head auxiliary for a future-token prediction depth.""" + + def __init__(self, hidden_size: int, ln_f_weight: torch.Tensor): + super().__init__() + self.weight = nn.Parameter(ln_f_weight.detach().clone()) + self.proj = nn.Linear(hidden_size, hidden_size, bias=False) + nn.init.zeros_(self.proj.weight) + self.eps = 1e-5 + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + # F.rms_norm handles the dtype internally -- no manual fp32 upcast + # allocation, no (B, T, D) intermediate tensor, free for any dtype. + h_norm = _native_rms_norm(hidden, (self.weight.shape[0],), self.weight.to(hidden.dtype), self.eps) + return h_norm + self.proj(h_norm) + + +class MultiTokenPrediction(nn.Module): + + def __init__(self, config, depth: int = 1, ln_f_weight: torch.Tensor | None = None): + super().__init__() + self.depth = depth + initial_weight = ( + ln_f_weight if ln_f_weight is not None else torch.ones(config.n_embd) + ) + self.heads = nn.ModuleList([ + MTPHead(config.n_embd, initial_weight) for _ in range(depth) + ]) + self.last_raw = None # unweighted mean MTP loss from the last compute_loss (for logging) + + def compute_loss( + self, + x: torch.Tensor, # (B, T, D) -- pre-lm_head hidden states + targets: torch.Tensor, # (B, T) -- already the +1 shifted targets + lm_head_weight: torch.Tensor, # (vocab_size, D) -- tied weight + mtp_lambda: float = 0.3, + ) -> torch.Tensor: + + + total = torch.zeros((), device=x.device, dtype=torch.float32) + valid_depths = 0 + for k, head in enumerate(self.heads): + offset = k + 1 # additional token offset beyond the main +1 shift + Tk = x.size(1) - offset + if Tk <= 0: + continue + h = head(x[:, :Tk]) # (B, Tk, D) + tgt_k = targets[:, offset:] # (B, Tk) + # Keep sum reduction: all-ignored chunks must contribute zero, not NaN. + loss_k = chunked_cross_entropy(h, lm_head_weight, tgt_k, ignore_index=-1) + total = total + loss_k + valid_depths += 1 + if valid_depths == 0: + self.last_raw = torch.zeros((), device=x.device, dtype=torch.float32) + return self.last_raw + self.last_raw = (total / valid_depths).detach() + return mtp_lambda * (total / valid_depths) + + +@dataclass +class GPTConfig: + """Configuration for GPT model.""" + block_size: int = _alpha_config_attr('CONTEXT', 1024) + vocab_size: int = _alpha_config_attr('vocab_size', 50304) + n_layer: int = _alpha_config_attr('numberoflayers', 12) + n_head: int = _alpha_config_attr('numberofheads', 12) + n_embd: int = _alpha_config_attr('D_MODEL', 768) + dropout: float = _alpha_config_attr('dropout', 0.0) + + bias: bool = _alpha_config_attr('bias', False) + ffn_mult: float = _alpha_config_attr('ffn_mult', 4.0) + + d_rope: int | None = _alpha_config_attr('d_rope', None) + + top_k: int = _alpha_config_attr('top_k', 64) + n_experts: int = _alpha_config_attr('n_experts', 8) + n_shared: int = _alpha_config_attr('n_shared', 1) + use_moe: bool = _alpha_config_attr('use_moe', False) + + moe_top_k: int = _alpha_config_attr('moe_top_k', 2) + + gradient_checkpointing: bool = _alpha_config_attr('gradient_checkpointing', False) # trade ~20-30% more compute for substantially less activation memory (grows with n_layer -- see GPT.forward) + num_kv_heads: int | None = _alpha_config_attr('num_kv_heads', None) + + pattern: str = _alpha_config_attr('pattern', 'pyramid_swa') + + sliding_window_size: int = _alpha_config_attr('sliding_window_size', 256) + + tp_size: int = _alpha_config_attr('tensor_parallel_size', 1) + + + gen_headroom: int = _alpha_config_attr('gen_headroom', 0) + + max_gen_tokens: int | None = _alpha_config_attr('max_gen_tokens', None) + + + mod_capacity: float = _alpha_config_attr('mod_capacity', 0.125) + + + use_mtp: bool = _alpha_config_attr('use_mtp', False) + mtp_depth: int = _alpha_config_attr('mtp_depth', 1) + mtp_lambda: float = _alpha_config_attr('mtp_lambda', 0.3) + + + use_nvfp4: bool = _alpha_config_attr('use_nvfp4', False) + + + use_int8: bool = _alpha_config_attr('use_int8', False) + int8_group_size: int = _alpha_config_attr('int8_group_size', 128) + + mod_gate_entropy_coeff: float = _alpha_config_attr('mod_gate_entropy_coeff', 0.01) + # Read via getattr(config, ...) in MoDBlock / AryaBlock; without fields here config.py values were ignored. + mod_budget_coeff: float = getattr(alpha_config, 'mod_budget_coeff', 0.05) + mod_hard_eval: bool = getattr(alpha_config, 'mod_hard_eval', False) + moe_fixed_capacity: int | None = getattr(alpha_config, 'moe_fixed_capacity', None) + + # Engram static n-gram memory (arXiv:2601.07372). Tuples, not lists: + # dataclass rejects mutable defaults. + use_engram: bool = _alpha_config_attr('use_engram', False) + engram_layer_ids: tuple = tuple(_alpha_config_attr('engram_layer_ids', (1,))) + engram_max_ngram: int = _alpha_config_attr('engram_max_ngram', 3) + engram_vocab_size: tuple = tuple(_alpha_config_attr('engram_vocab_size', (8192, 8192))) + engram_embed_per_ngram: int = _alpha_config_attr('engram_embed_per_ngram', 128) + engram_n_heads: int = _alpha_config_attr('engram_n_heads', 4) + engram_kernel_size: int = _alpha_config_attr('engram_kernel_size', 4) + engram_seed: int = _alpha_config_attr('engram_seed', 0) + # Set automatically from the tokenizer lookup at build time; saved in + # checkpoints so GPT(config) can be rebuilt without a tokenizer. + engram_compressed_vocab: int = _alpha_config_attr('engram_compressed_vocab', 0) + +class GPT(nn.Module): + """Full GPT language model""" + + def __init__(self, config, engram_lookup=None): + super().__init__() + if config.vocab_size is None: + raise ValueError("config.vocab_size must be set") + if config.block_size is None: + raise ValueError("config.block_size must be set") + self.config = config + + + + head_dim = config.n_embd // config.n_head + rope_dim = head_dim if getattr(config, 'd_rope', None) is None else config.d_rope + + gen_headroom = max(0, int(getattr(config, 'gen_headroom', 0))) + rope_table_len = config.block_size + gen_headroom + + attn_layers, self.layer_types = build_attention_layers( + num_layers=config.n_layer, + hidden_size=config.n_embd, + num_heads=config.n_head, + dropout=config.dropout, + bias=config.bias, + num_kv_heads=config.num_kv_heads, + pattern=config.pattern, + max_seq_len=rope_table_len, + sliding_window_size=config.sliding_window_size, + tp_size=config.tp_size, + rope_dim=rope_dim, + ) + + + mod_capacity = float(getattr(config, 'mod_capacity', 0.125)) + blocks = [] + for i in range(config.n_layer): + block = AryaBlock(config, attn_layers[i], role=self.layer_types[i]) + if self.layer_types[i] == "mod": + + block = MoDBlock(config, block, capacity=mod_capacity) + blocks.append(block) + + # Hoist once: avoid an isinstance(block, MoDBlock) check (and the + # use_moe lookup that decides whether a block yields a moe_aux + # tensor) on every block on every forward. + self._no_kw = [{} for _ in blocks] # shared empty kwargs: no per-layer dict alloc + # Document masking (set from train.py; None = plain causal, which is what inference uses). + self.doc_mask_bos_id = None + _ws = [] + for _b in blocks: + _w = (_b.inner if isinstance(_b, MoDBlock) else _b).attn.window_size + if _w not in _ws: + _ws.append(_w) + self._mask_windows = _ws # distinct attention window sizes (None = full causal) + self._causal_mask_cache: dict = {} # (T, device) -> causal masks; see build_attn_masks() + self._staged_masks = None # set via stage_attn_masks(); a module attribute, not a forward argument + self._block_flags = [ + (isinstance(b, MoDBlock), b.inner.use_moe if isinstance(b, MoDBlock) else b.use_moe) + for b in blocks + ] + + self.transformer = nn.ModuleDict(dict( + wte=nn.Embedding(config.vocab_size, config.n_embd), + drop=nn.Dropout(config.dropout), + h=nn.ModuleList(blocks), + ln_f=RMSNorm(config.n_embd, eps=1e-5), + )) + self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) + + # ── Engram (optional) ──────────────────────────────────────────── + self.engram_hasher = None + self.engram_layers = nn.ModuleDict() + self._engram_layer_ids = frozenset() + if bool(getattr(config, 'use_engram', False)): + if Engram is None: + raise ImportError("use_engram=True but engram.py could not be imported") + ids = tuple(int(l) for l in config.engram_layer_ids) + if not ids or any(not (0 <= l < config.n_layer) for l in ids) or len(set(ids)) != len(ids): + raise ValueError( + f"engram_layer_ids={ids} must be unique and within [0, {config.n_layer})") + if engram_lookup is not None: + config.engram_compressed_vocab = int(engram_lookup.max()) + 1 + if int(getattr(config, 'engram_compressed_vocab', 0)) <= 0: + raise ValueError( + "use_engram=True needs engram_lookup (built from the tokenizer via " + "engram.CompressedTokenizer) or a config with engram_compressed_vocab set") + self.engram_hasher = NgramHasher( + layer_ids=ids, + max_ngram=config.engram_max_ngram, + vocab_size_per_ngram=config.engram_vocab_size, + n_heads=config.engram_n_heads, + compressed_vocab=config.engram_compressed_vocab, + raw_vocab_size=config.vocab_size, + seed=config.engram_seed, + lookup=engram_lookup, + ) + self.engram_layers = nn.ModuleDict({ + str(l): Engram( + hidden_size=config.n_embd, + head_sizes=self.engram_hasher.head_sizes[i], + max_ngram=config.engram_max_ngram, + embed_per_ngram=config.engram_embed_per_ngram, + n_heads=config.engram_n_heads, + kernel_size=config.engram_kernel_size, + ) + for i, l in enumerate(ids) + }) + self._engram_layer_ids = frozenset(ids) + + + use_mtp = bool(getattr(config, 'use_mtp', False)) + mtp_depth = int(getattr(config, 'mtp_depth', 1)) + self.mtp = MultiTokenPrediction(config, depth=mtp_depth) if use_mtp else None + self._mtp_loss = self._compute_mtp_loss if use_mtp else self._skip_mtp_loss + + self._mtp_lambda = float(getattr(config, 'mtp_lambda', 0.3)) + + + self.apply(self._init_weights) + for _eg in self.engram_layers.values(): + _eg.reset_parameters() # tables std=0.01 (apply() above overwrote it with 0.02) + + + # apply(_init_weights) above re-inits every nn.Linear to N(0, 0.02), + # including the always-on attn_gate_proj. Restore its zero-init so the + # gate is identity (2*sigmoid(0) = 1) at step 0. + for block in self.transformer.h: + real_block = block.inner if isinstance(block, MoDBlock) else block + if isinstance(real_block.attn, DenseMHA): + nn.init.zeros_(real_block.attn.attn_gate_proj.weight) + + for block in self.transformer.h: + if isinstance(block, MoDBlock): + with torch.no_grad(): + block.router.bias.fill_(2.0) + + if self.mtp is not None: + with torch.no_grad(): + for head in self.mtp.heads: + head.weight.copy_(self.transformer.ln_f.weight) + head.proj.weight.zero_() + + + self.transformer.wte.weight = self.lm_head.weight + + + residual_std = 0.02 / math.sqrt(config.n_layer) + for name, p in self.named_parameters(): + if name.endswith("out_proj.weight") or name.endswith("W_out.weight"): + torch.nn.init.normal_(p, mean=0.0, std=residual_std) + + + self._use_nvfp4: bool = False + self._use_int8: bool = False + self._quant_mode = None + + if _nvfp4_enabled(): + n_quantized = apply_nvfp4_to_model(self) + self._use_nvfp4 = True + self._quant_mode = "nvfp4" + warnings.warn( + f"[nvFP4/RHT] torchao quantized {n_quantized} Linear layers " + f"(Blackwell SM100+). Weight: FP4 e2m1 per-block-16 RHT. " + f"Activation: FP4 dynamic per-tensor scale (FP8 intermediate). " + f"Fused GEMM kernel emitted by torch.compile. " + f"Skipped: wte, lm_head, norms, router, attn gate, small/misaligned linears. " + f"Verify val_loss vs bf16 baseline before committing to a long run.", + stacklevel=2, + ) + elif _int8_enabled(): + gs = int(getattr(config, 'int8_group_size', 128)) + n_quantized = apply_int8_to_model(self, group_size=gs) + self._use_int8 = True + self._quant_mode = "int8" + warnings.warn( + f"[INT8/Jetfire] torchao quantized {n_quantized} Linear layers " + f"(T4/Turing+). Weight: INT8 per-group-{gs} symmetric. " + f"Activation: INT8 dynamic per-tensor scale (runtime, no calibration). " + f"Fused kernel emitted by torch.compile via Inductor. " + f"Skipped: wte, lm_head, norms, router, attn gate, small linears. " + f"Verify val_loss vs fp16 baseline before a long run.", + stacklevel=2, + ) + + + + self.transformer.wte.weight = self.lm_head.weight + assert self.transformer.wte.weight is self.lm_head.weight + + + + + def set_tp_group(self, group) -> None: + + if self.config.tp_size <= 1: + return + for block in self.transformer.h: + + real_block = block.inner if isinstance(block, MoDBlock) else block + real_block.attn.tp_group = group + if hasattr(real_block, "ffn") and real_block.ffn is not None: + real_block.ffn.tp_group = group + + + if torch.distributed.is_initialized() and group is not None: + for block in self.transformer.h: + real_block = block.inner if isinstance(block, MoDBlock) else block + attn = real_block.attn + if not isinstance(attn, DenseMHA): + continue + if attn.num_kv_heads >= attn.tp_size: + continue + with torch.no_grad(): + # Full qkv projection tensor is the canonical object to align. + torch.distributed.broadcast(attn.qkv_proj.weight.data, src=0, group=group) + if attn.qkv_proj.bias is not None: + torch.distributed.broadcast(attn.qkv_proj.bias.data, src=0, group=group) + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + + def get_num_params(self, non_embedding=True): + """Count number of parameters""" + n_params = sum(p.numel() for p in self.parameters()) + if non_embedding: + n_params -= self.transformer.wte.weight.numel() + return n_params + + def _doc_ids(self, idx): + # Documents are [BOS ... EOS], bos_id == eos_id: a new document starts at a BOS that directly follows another BOS. + is_bos = idx == self.doc_mask_bos_id + starts = torch.cat([torch.zeros_like(is_bos[:, :1]), is_bos[:, 1:] & is_bos[:, :-1]], dim=1) + return starts.to(torch.int32).cumsum(dim=1) + + def build_attn_masks(self, idx): + """Attention masks for a plain (kv_cache=None) forward. Call this OUTSIDE torch.compile and pass the result as + forward(..., attn_masks=masks): kernel.doc_masks / kernel.causal_masks are @torch.compiler.disable'd, so + building them inside the compiled forward splits the graph in two on every step. Causal masks depend only on + (T, device) and are cached, so they are built once; doc masks depend on the batch and are built per call.""" + if self.doc_mask_bos_id is not None: + return kernel.doc_masks(self._doc_ids(idx.long()), self._mask_windows) + key = (idx.size(1), idx.device) + masks = self._causal_mask_cache.get(key) + if masks is None: + masks = self._causal_mask_cache[key] = kernel.causal_masks(idx.size(1), self._mask_windows, idx.device) + return masks + + def stage_attn_masks(self, idx): + """Build this batch's masks OUTSIDE torch.compile and stage them on the module. Passing them as a forward + argument to a compiled module dropped the loss's grad in testing; staging as an attribute did not.""" + self._staged_masks = self.build_attn_masks(idx) + + def forward(self, idx, targets=None, need_logits=True, kv_cache=None, attn_masks=None): + + idx = idx.long() + if targets is not None: + targets = targets.long() + b, t = idx.size() + # Plain check, not torch._check with a message closure (that closure forces a graph break). + if t > self.config.block_size: + raise ValueError(f"Sequence too long: {t} > {self.config.block_size}") + cache_pos = 0 + if kv_cache is not None: + if not self.supports_kv_cache: + raise NotImplementedError("this model configuration does not support kv_cache") + cache_pos = kv_cache.pos + + tok_emb = self.transformer.wte(idx) # (B, T, n_embd) + x = self.transformer.drop(tok_emb) + + + mod_gate_loss = torch.zeros((), device=x.device, dtype=torch.float32) + moe_aux_loss = torch.zeros((), device=x.device, dtype=torch.float32) + + # Document ids (documents are [BOS ... EOS], bos_id == eos_id: a new document starts at a BOS that + # directly follows another BOS). Shared by the attention masks and by Engram's hash / short conv. + doc_ids = self._doc_ids(idx) if (self.doc_mask_bos_id is not None and kv_cache is None) else None + xsa_pos_mask = None + if doc_ids is not None: + doc_starts = torch.cat( + (torch.ones_like(doc_ids[:, :1], dtype=torch.bool), + doc_ids[:, 1:] != doc_ids[:, :-1]), + dim=1, + ) + xsa_pos_mask = (~doc_starts).to(dtype=torch.float32) + + engram_cache = kv_cache.extra if kv_cache is not None else None + engram_hashes = ( + self.engram_hasher(idx, doc_ids, cache=engram_cache) + if self.engram_hasher is not None else None + ) + + # Static-per-call decisions, resolved once instead of once per layer. + engram_on = engram_hashes is not None + use_ckpt = self.config.gradient_checkpointing and self.training + # Attention masks (FlexAttention BlockMasks): built once per forward, shared by every layer. + # Selection below uses static Python values only (kv_cache / cache_pos / doc_mask_bos_id), so it + # is resolved at trace time -- no data-dependent branch. + # Documents are [BOS ... EOS] and bos_id == eos_id, so a document starts at a BOS that directly + # follows another BOS. RoPE scores depend only on relative offsets: positions are NOT reset. + if kv_cache is not None and cache_pos > 0: + attn_masks = None # single-token decode: DenseMHA runs SDPA over the cache, no mask + elif attn_masks is None and self._staged_masks is not None: + attn_masks = self._staged_masks # staged by stage_attn_masks() outside torch.compile + elif attn_masks is None: # not prebuilt by build_attn_masks(): build here (graph break under compile) + if doc_ids is not None: + attn_masks = kernel.doc_masks(doc_ids, self._mask_windows) + else: + attn_masks = kernel.causal_masks(t, self._mask_windows, idx.device) + if kv_cache is None: + kws = [{"attn_masks": attn_masks, "xsa_pos_mask": xsa_pos_mask}] * len(self._no_kw) + else: + kws = [{"kv_cache": kv_cache.layer(i), "cache_pos": cache_pos, + "attn_masks": attn_masks, "xsa_pos_mask": xsa_pos_mask} + for i in range(len(self._no_kw))] + ckw = {"attn_masks": attn_masks, "xsa_pos_mask": xsa_pos_mask} + + for layer_i, (block, (is_mod, has_moe), kw) in enumerate(zip(self.transformer.h, self._block_flags, kws)): + if engram_on and layer_i in self._engram_layer_ids: + # Official placement: h = h + Engram(h, ids) BEFORE the block's attention. + hash_ids = engram_hashes.pop(layer_i) + engram_layer_cache = kv_cache.layer(layer_i) if kv_cache is not None else None + x = x + self.engram_layers[str(layer_i)]( + x, hash_ids, doc_ids, cache=engram_layer_cache) + del hash_ids + if use_ckpt: + if is_mod: + (x, moe_aux), gate_aux = torch.utils.checkpoint.checkpoint( + block, x, use_reentrant=False, **ckw + ) + mod_gate_loss = mod_gate_loss + gate_aux.float() + elif has_moe: + x, moe_aux = torch.utils.checkpoint.checkpoint( + block, x, use_reentrant=False, **ckw + ) + else: + x = torch.utils.checkpoint.checkpoint( + block, x, use_reentrant=False, **ckw + ) + moe_aux = None + elif is_mod: + (x, moe_aux), gate_aux = block(x, **kw) + mod_gate_loss = mod_gate_loss + gate_aux.float() + elif has_moe: + x, moe_aux = block(x, **kw) + else: + x = block(x, **kw) + moe_aux = None + if moe_aux is not None: + moe_aux_loss = moe_aux_loss + moe_aux.float() + + pre_ln_f = x + x = self.transformer.ln_f(x) + + if targets is not None: + mtp_loss = self._mtp_loss(pre_ln_f, targets) + + if not need_logits: + + ce_loss = chunked_cross_entropy(x, self.lm_head.weight, targets, ignore_index=-1) + logits = None + else: + logits = F.linear(x, self.lm_head.weight) + # sum / clamp(n_valid, 1) == the mean, but an all-ignored batch gives 0 instead of NaN + # (same behaviour as the chunked path). + ce_loss = F.cross_entropy( + logits.reshape(-1, logits.size(-1)).float(), + targets.reshape(-1), + ignore_index=-1, + reduction="sum", + ) / (targets != -1).sum().clamp_min(1) + + total_loss = ce_loss + mtp_loss + mod_gate_loss + moe_aux_loss + self.last_ce_loss = ce_loss.detach() + self.last_moe_aux_loss = moe_aux_loss.detach() + self.last_total_loss = total_loss.detach() + else: + logits = F.linear(x[:, [-1], :], self.lm_head.weight) # (B, 1, vocab_size) + ce_loss = None + total_loss = None + + if kv_cache is not None: + kv_cache.pos += t + return logits, total_loss + + def _compute_mtp_loss(self, hidden, targets): + loss = self.mtp.compute_loss( + hidden, targets, self.lm_head.weight, mtp_lambda=self._mtp_lambda + ) + self.last_mtp_loss = self.mtp.last_raw + return loss + + def _skip_mtp_loss(self, hidden, targets): + self.last_mtp_loss = None + return torch.zeros((), device=hidden.device, dtype=torch.float32) + + @property + def supports_kv_cache(self) -> bool: + return all( + isinstance((b.inner if isinstance(b, MoDBlock) else b).attn, DenseMHA) + and (b.inner if isinstance(b, MoDBlock) else b).attn.causal + for b in self.transformer.h + ) + + def new_kv_cache(self, max_len: int | None = None) -> "KVCache": + """Cache sized for `max_len` tokens (default and upper bound: block_size).""" + cap = self.config.block_size if max_len is None else min(int(max_len), self.config.block_size) + return KVCache(len(self.transformer.h), cap) + + @torch.no_grad() + def generate(self, idx, max_new_tokens=None, temperature=1.0, top_k=None, use_cache=True): + + if max_new_tokens is None: + max_new_tokens = getattr(self.config, 'max_gen_tokens', None) + if max_new_tokens is None: + raise ValueError( + "generate() needs max_new_tokens, either passed " + "directly or set as max_gen_tokens in config.py." + ) + if temperature is None or temperature <= 0: + raise ValueError( + f"generate(): temperature must be > 0, got {temperature!r}. " + f"For greedy decoding use top_k=1 with a small temperature." + ) + # BUGFIX: @torch.no_grad() does NOT put the module in eval mode, so + # sampling from a model left in .train() ran with dropout active (and + # kept mutating MoE expert_bias). Force eval, restore afterwards. + was_training = self.training + self.eval() + try: + return self._generate_loop(idx, max_new_tokens, temperature, top_k, use_cache) + finally: + if was_training: + self.train() + + def _generate_loop(self, idx, max_new_tokens, temperature, top_k, use_cache=True): + block_size = self.config.block_size + cached = bool(use_cache) and self.supports_kv_cache + cache = None + for _ in range(max_new_tokens): + if cached and idx.size(1) <= block_size: + # Full context still fits: KV cache gives the same result as re-encoding it. + if cache is None: + cache = self.new_kv_cache() + logits, _ = self(idx, kv_cache=cache) # prefill + else: + logits, _ = self(idx[:, -1:], kv_cache=cache) # one token + else: + # Past block_size the original sliding-window re-encode applies + # (positions re-based each step), which a cache can't reproduce. + cached, cache = False, None + idx_cond = idx if idx.size(1) <= block_size else idx[:, -block_size:] + logits, _ = self(idx_cond) + + + logits = logits[:, -1, :] / temperature + + + if top_k is not None: + v, _ = torch.topk(logits, min(top_k, logits.size(-1))) + logits[logits < v[:, [-1]]] = -float('Inf') + + + probs = F.softmax(logits.float(), dim=-1) # (B, vocab_size) + + # Sample next token + idx_next = torch.multinomial(probs, num_samples=1) # (B, 1) + + # Append to sequence + idx = torch.cat((idx, idx_next), dim=1) # (B, T+1) + + return idx \ No newline at end of file