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