Upload 6 files
7a936fb verified 3 days ago
99.6 kB
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