Upload 4 files

config.py

@@ -0,0 +1,136 @@
+
+use_liger=True
+xla_on_gpu=True
+# ---------------------------------------------------------------------------
+# Core shape -- chosen for GPU efficiency, not just "a number that fits"
+# ---------------------------------------------------------------------------
+CONTEXT = 512
+max_gen_tokens = 256
+yarn_scale = 1.0 # neutralized YaRN extension ramp for clean RoPE sanity checks
+gen_headroom = 128
+vocab_size = 16384
+
+# Missing config fields expected by the shared GPTConfig dataclass and the
+# model factory path. Keeping them explicit here removes the fallback-warning
+# cascade shown in the traceback and stabilizes object construction.
+bias = False
+
+
+
+gradient_checkpointing = False
+num_kv_heads = 4
+pattern = "dense"
+sliding_window_size = 512
+# tensor_parallel_size is the knob the config dataclass reads.
+tensor_parallel_size = 1
+
+# YaRN config defaults from the model constructor.
+
+mod_capacity = 0.25
+use_mtp    = False # set True to activate
+mtp_depth  = 1      # 1 = predict 2 tokens ahead; 2 = also predict 3 ahead
+mtp_lambda = 0.3 
+use_nvfp4 = False
+use_int8 = False
+int8_group_size = 128
+mod_gate_entropy_coeff = 0.01
+# Experimental hardware-specific paths. These are deliberately opt-in and
+# are read live by model.py/prefetch.py so post-import CLI overrides work.
+
+numberoflayers = 12     # 4 complete pyramid_swa groups (4x4) -- no odd
+                           # trailing layer forced into an incomplete group
+numberofheads = 12
+D_MODEL = 192         # multiple of 64 -> head_dim = 192/12 = 16, a clean
+                           # Tensor Core tile-friendly size (fp16 wants
+                           # multiples of 16; 16 divides evenly with no
+                           # padding waste on T4/A100 GEMM tiling)
+      # AdamW-scale LR for a model this size; only
+                           # takes effect if use_static_lr=False below
+
+# ---------------------------------------------------------------------------
+# Attention
+# ---------------------------------------------------------------------------
+num_kv_heads = 4          # 2x GQA compression on global layers (8->4 kv
+                           # heads); pyramid_swa steps kv heads [1,2,4]
+                           # across each group's 3 SWA layers (MQA->GQA),
+                           # verified against the real build_attention_layers
+                           # logic in model.py, not assumed
+
+sliding_window_size = 512  # explicit (was previously left to fall back to
+                           # the default with a warning) -- 1/4 of CONTEXT,
+                           # a reasonable local-attention window for this
+                           # sequence length
+use_sdpa=False
+optimizer_8bit = True
+# ---------------------------------------------------------------------------
+# FFN / MoE
+# ---------------------------------------------------------------------------
+ffn_mult = 8 / 3           # SwiGLU hidden_dim multiplier. model.py rounds
+                           # d_model*ffn_mult UP to the nearest multiple of
+                           # 64 (256 * 8/3 = 682.67 -> 704), avoiding the
+                           # prime-dimension GEMM padding waste a plain
+                           # round() would produce (683 is prime).
+use_moe = False            # plain dense SwiGLU_FFN every layer. AryaSparseMoE
+                           # currently has NO auxiliary load-balancing loss
+                           # (expert_bias is trainable but nothing pushes it
+                           # toward balanced routing) -- leave this False
+                           # unless/until that's added, especially at this
+                           # small a scale where routing collapse compounds
+                           # fastest across depth.
+n_experts = 8
+n_shared = 1
+# T4 path
+
+# (use_int8 ignored on SM100+ -- nvFP4 takes priority)
+# ---------------------------------------------------------------------------
+# Precision / batching
+# ---------------------------------------------------------------------------
+PRECISION = 'fp16'         # T4 has no bf16 Tensor Cores -- fp16 is correct
+                           # here. On A100+, bf16 is the better choice (same
+                           # speed, wider dynamic range, no GradScaler
+                           # needed) -- change this if/when you move hardware.
+GLOBAL_DTYPE = PRECISION
+USE_ASYNC_LAYER_PREFETCH = True
+DISABLE_MMAP = False
+use_flash_outproj_add_rmsnorm = False
+use_flash_rope = False
+use_flash_swiglu = False
+use_fused_add_rmsnorm = False
+batch_size = 64 # starting point, not a measured optimum -- train.py
+                           # now logs `vram X/Y GB` on every step (added
+                           # recently). Watch that number on your first run
+                           # and raise batch_size toward your actual VRAM
+                           # headroom; this value has NOT been validated
+                           # against a real GPU run for this exact config.
+weight_decay = 0.01        # train.py now splits params into decay/no-decay
+                           # groups automatically (biases, RMSNorm gains,
+                           # expert_bias get weight_decay=0.0 regardless of
+                           # this value) -- this number only applies to
+                           # actual weight matrices + the tied embedding.
+
+# ---------------------------------------------------------------------------
+# Schedule
+# ---------------------------------------------------------------------------
+# Defaults to warmup+cosine (NOT static) for this fresh config -- unlike
+# your existing config.py, this file has no prior "it works, don't touch it"
+# history behind a flat LR, so there's no reason to ship it with the known
+# no-warmup risk baked in. WARMUP_STEPS=500 and MAX_STEPS=20000 are
+# hardcoded in train.py's main() (not read from config.py at all) -- change
+# them there directly if this run needs a different horizon.
+use_static_lr =True
+static_lr = 4e-3        # only used if you flip use_static_lr back to True
+learning_rate = 3e-3
+# ---------------------------------------------------------------------------
+# Parallelism / misc
+# ---------------------------------------------------------------------------
+
+
+data_parallel_only = True
+dropout = 0.0                   # explicit (was previously left to the
+                                  # default with a warning)
+use_muon = True
+# ---------------------------------------------------------------------------
+# Data
+# ---------------------------------------------------------------------------
+DATA_PATH = './training_data'
+data_path = DATA_PATH

inference.py

@@ -0,0 +1,361 @@
+"""
+inference.py -- load a train.py checkpoint and generate text with it.
+
+    python inference.py --checkpoint ./checkpoints/latest.pt --prompt "Once upon a time"
+    python inference.py --checkpoint ./checkpoints/latest.pt --interactive
+
+The checkpoint carries its own architecture config (see train.py's
+save_checkpoint), so you don't need config.py to match the training run --
+the model is reconstructed exactly as it was trained, every time.
+
+KNOWN LIMITATION: model.py's attention modules have no KV cache. Every
+generated token re-runs a full forward pass over the whole sequence so far
+(O(T) attention passes, each O(T) or O(T*window) itself -- no incremental
+state is carried between steps). Fine for short completions and for
+sanity-checking a checkpoint; a real KV cache is the natural next piece of
+work if you want fast interactive generation, and would need to be threaded
+through CompressedSparseAttention / HeavilyCompressedAttention /
+SlidingWindowAttention individually since each keeps different per-branch
+state (compressed blocks vs. raw recent tokens).
+"""
+
+import argparse
+import contextlib
+import sys
+
+import torch
+
+from model import GPT, GPTConfig
+from tokenizer import TokenizerWrapper
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Device selection
+# ─────────────────────────────────────────────────────────────────────────────
+
+def pick_device(requested: str | None) -> torch.device:
+    if requested is not None:
+        return torch.device(requested)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    try:
+        import torch_xla.core.xla_model as xm  # noqa: F401
+        return torch.device("xla")
+    except ImportError:
+        pass
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Checkpoint loading
+# ─────────────────────────────────────────────────────────────────────────────
+
+def load_model(checkpoint_path: str, device: torch.device) -> tuple[GPT, GPTConfig]:
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+
+    if "model_config" not in checkpoint:
+        raise RuntimeError(
+            f"'{checkpoint_path}' has no saved model_config -- it was written "
+            "before save_checkpoint() started saving the architecture config "
+            "alongside the weights. Either retrain (new checkpoints save it "
+            "automatically), or construct GPTConfig(...) by hand here to "
+            "match whatever config.py looked like when this checkpoint was "
+            "produced, and call GPT(that_config) instead of this function."
+        )
+
+    model_config = GPTConfig(**checkpoint["model_config"])
+    model = GPT(model_config)
+
+    state_dict = checkpoint["model_state_dict"]
+    
+   
+    if any(k.startswith("_orig_mod.") for k in state_dict):
+        state_dict = {
+            (k[len("_orig_mod."):] if k.startswith("_orig_mod.") else k): v
+            for k, v in state_dict.items()
+        }
+        print(
+            "[Inference] Checkpoint was saved from a torch.compile()'d model "
+            "(keys prefixed with '_orig_mod.') -- stripped the prefix to "
+            "load onto this plain GPT instance.",
+            file=sys.stderr,
+        )
+
+    model.load_state_dict(state_dict)
+    model.to(device)
+    model.eval()
+
+    n_params = model.get_num_params() / 1e6
+    step = checkpoint.get("step", "?")
+    print(f"[Inference] Loaded checkpoint (step {step}, {n_params:.2f}M params) onto {device}", file=sys.stderr)
+    return model, model_config
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Sampling
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _apply_repetition_penalty(logits: torch.Tensor, generated_ids: list[int], penalty: float) -> torch.Tensor:
+    """CTRL-style penalty, COUNT-SCALED: a token seen N times gets divided
+    by penalty**N (if positive) / multiplied by penalty**N (if negative),
+    not a single flat divide regardless of occurrence count.
+
+    The original flat version (divide once no matter how many times a
+    token has recurred) is too weak on small/repetitive-domain models
+    (e.g. TinyStories-style data where "so so so happy!" is a genuinely
+    common training pattern): once the model locks onto a token, the
+    penalty never gets any stronger no matter how many times it fires,
+    so a token that's only mildly favored over the runner-up never gets
+    pushed low enough to escape the loop. Scaling by occurrence count
+    means each additional repeat compounds the penalty, so a loop that's
+    3-4 tokens deep gets meaningfully suppressed even if a single
+    application wouldn't have been enough.
+    """
+    if penalty == 1.0 or not generated_ids:
+        return logits
+    from collections import Counter
+    counts = Counter(generated_ids)
+    seen = torch.tensor(sorted(counts.keys()), device=logits.device, dtype=torch.long)
+    reps = torch.tensor([counts[t] for t in sorted(counts.keys())], device=logits.device, dtype=logits.dtype)
+    vals = logits[0, seen]
+    factor = penalty ** reps
+    logits[0, seen] = torch.where(vals > 0, vals / factor, vals * factor)
+    return logits
+
+
+def _block_repeated_ngrams(logits: torch.Tensor, generated_ids: list[int], ngram_size: int) -> torch.Tensor:
+    """Hard n-gram repeat block (HF's `no_repeat_ngram_size`, same idea).
+
+    Look at the last (ngram_size - 1) generated tokens. If that same
+    prefix has appeared earlier in generated_ids, whatever token followed
+    it there is banned from being chosen again right now (logit -> -inf).
+    This is a HARD constraint, unlike repetition_penalty which only
+    nudges probabilities down -- it's what actually stops "so so so so
+    so..." dead rather than just making it less likely each step. Only
+    kicks in once at least ngram_size-1 tokens have been generated.
+    """
+    if ngram_size <= 0 or len(generated_ids) < ngram_size - 1:
+        return logits
+    prefix_len = ngram_size - 1
+    if prefix_len == 0:
+        return logits
+    current_prefix = tuple(generated_ids[-prefix_len:])
+    banned = set()
+    for i in range(len(generated_ids) - prefix_len):
+        if tuple(generated_ids[i:i + prefix_len]) == current_prefix:
+            banned.add(generated_ids[i + prefix_len])
+    if banned:
+        banned_idx = torch.tensor(sorted(banned), device=logits.device, dtype=torch.long)
+        logits[0, banned_idx] = float("-inf")
+    return logits
+
+
+def _top_k_filter(logits: torch.Tensor, top_k: int) -> torch.Tensor:
+    k = min(top_k, logits.size(-1))
+    values, _ = torch.topk(logits, k)
+    logits[logits < values[:, [-1]]] = float("-inf")
+    return logits
+
+
+def _top_p_filter(logits: torch.Tensor, top_p: float) -> torch.Tensor:
+    """Nucleus sampling: keep the smallest prefix of sorted tokens whose
+    cumulative probability crosses top_p, drop the rest."""
+    sorted_logits, sorted_idx = torch.sort(logits, descending=True)
+    probs = torch.softmax(sorted_logits, dim=-1)
+    cum_probs = torch.cumsum(probs, dim=-1)
+    # remove token i if the cumulative prob BEFORE it already exceeds top_p
+    # (i.e. it wasn't needed to cross the threshold) -- matches the standard
+    # HF TopPLogitsWarper convention.
+    remove = (cum_probs - probs) > top_p
+    sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
+    return torch.full_like(logits, float("-inf")).scatter(1, sorted_idx, sorted_logits)
+
+
+@torch.no_grad()
+def generate(
+    model: GPT,
+    tokenizer: TokenizerWrapper,
+    prompt: str,
+    device: torch.device,
+    max_new_tokens: int = 200,
+    temperature: float = 1.0,
+    top_k: int | None = None,
+    top_p: float | None = None,
+    repetition_penalty: float = 1.0,
+    no_repeat_ngram_size: int = 0,
+    stream: bool = True,
+    seed: int | None = None,
+    num_samples: int = 1,
+) -> list[str]:
+    if seed is not None:
+        torch.manual_seed(seed)
+
+    prompt_ids = tokenizer.encode(prompt, add_bos=True, add_eos=False)
+    idx = torch.tensor([prompt_ids] * num_samples, dtype=torch.long, device=device)
+    block_size = model.config.block_size
+    eos_id = getattr(tokenizer, "eos_id", None)
+
+    generated_ids = [list(prompt_ids) for _ in range(num_samples)]
+    finished = [False] * num_samples
+    # Decode the FULL sequence each step and print only the new suffix,
+    # rather than decoding one new token at a time: BPE/subword tokens don't
+    # each cleanly map to standalone text (a token can be half a multi-byte
+    # character or mid-word piece), so decoding incrementally token-by-token
+    # can emit garbled text at token boundaries. Re-decoding the whole
+    # sequence and diffing against what's already been printed sidesteps
+    # that regardless of the tokenizer's internals. O(T) decode cost per
+    # step, which is fine at CLI-generation scale.
+    prev_text = [tokenizer.decode(g) for g in generated_ids] if stream else None
+    if stream and num_samples == 1:
+        print(prev_text[0], end="", flush=True)
+
+    for _ in range(max_new_tokens):
+        if all(finished):
+            break
+        idx_cond = idx if idx.size(1) <= block_size else idx[:, -block_size:]
+        logits, _ = model(idx_cond)
+        logits = logits[:, -1, :].float()  # last position, fp32 for stable sampling math regardless of training dtype
+
+        for b in range(num_samples):
+            if finished[b]:
+                continue
+            row = logits[b:b + 1]
+            row = _apply_repetition_penalty(row, generated_ids[b], repetition_penalty)
+            row = _block_repeated_ngrams(row, generated_ids[b], no_repeat_ngram_size)
+            row = row / max(temperature, 1e-5)
+            if top_k is not None:
+                row = _top_k_filter(row, top_k)
+            if top_p is not None:
+                row = _top_p_filter(row, top_p)
+            logits[b:b + 1] = row
+
+        probs = torch.softmax(logits, dim=-1)
+        next_ids = torch.multinomial(probs, num_samples=1)  # (num_samples, 1)
+
+        # Once a sequence hits EOS, keep feeding it its own last token so the
+        # batch stays rectangular, but stop appending to its recorded output.
+        for b in range(num_samples):
+            if finished[b]:
+                next_ids[b, 0] = idx[b, -1]
+                continue
+            tid = next_ids[b, 0].item()
+            generated_ids[b].append(tid)
+            if eos_id is not None and tid == eos_id:
+                finished[b] = True
+
+        idx = torch.cat([idx, next_ids], dim=1)
+
+        if stream and num_samples == 1:
+            new_text = tokenizer.decode(generated_ids[0])
+            print(new_text[len(prev_text[0]):], end="", flush=True)
+            prev_text[0] = new_text
+
+    if stream and num_samples == 1:
+        print()
+
+    return [tokenizer.decode(g) for g in generated_ids]
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# CLI
+# ─────────────────────────────────────────────────────────────────────────────
+
+def build_arg_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description="Generate text from a trained Arya/Veylon checkpoint.")
+    p.add_argument("--checkpoint", required=True, help="Path to a checkpoint saved by train.py")
+    p.add_argument("--tokenizer", default="./tokenizer.model", help="Path passed to TokenizerWrapper")
+    p.add_argument("--prompt", default=None, help="Prompt text (omit if using --interactive)")
+    p.add_argument("--interactive", action="store_true", help="Drop into a REPL: one prompt per line")
+    p.add_argument("--max-new-tokens", type=int, default=200)
+    p.add_argument("--temperature", type=float, default=0.8)
+    p.add_argument("--top-k", type=int, default=50)
+    p.add_argument("--top-p", type=float, default=None, help="Nucleus sampling threshold, e.g. 0.9. Combine freely with --top-k.")
+    p.add_argument("--repetition-penalty", type=float, default=1.3,
+                    help="CTRL-style penalty on already-generated tokens, now count-scaled "
+                         "(penalty**occurrences). 1.0 disables it -- small models loop into "
+                         "degenerate repetition ('so so so so...') without some penalty here, "
+                         "so this defaults ON.")
+    p.add_argument("--no-repeat-ngram-size", type=int, default=3,
+                    help="Hard-ban repeating any n-gram of this size (HF-style). 0 disables it. "
+                         "This is what actually stops infinite word loops -- repetition-penalty "
+                         "alone only makes them less likely, this makes them impossible.")
+    p.add_argument("--num-samples", type=int, default=1, help="Generate N completions per prompt (batched)")
+    p.add_argument("--seed", type=int, default=None)
+    p.add_argument("--device", default=None, help="cuda / cpu / mps / xla -- auto-detected if omitted")
+    p.add_argument("--dtype", default="auto", choices=["auto", "fp32", "fp16", "bf16"],
+                    help="Inference precision. 'auto' uses fp16 on CUDA, fp32 elsewhere.")
+    return p
+
+
+def _resolve_dtype(dtype_arg: str, device: torch.device) -> torch.dtype | None:
+    if device.type not in ("cuda", "cpu"):
+        # torch.autocast is only exercised here against cuda/cpu; MPS and XLA
+        # handle precision differently (XLA in particular usually wants
+        # XLA_USE_BF16 or torch_xla's own autocast, not torch.autocast).
+        # Guessing at that without real hardware to test against would be
+        # worse than just running those in fp32 and saying so.
+        if dtype_arg != "auto":
+            print(f"[Inference] --dtype={dtype_arg} requested but autocast isn't wired up for device "
+                  f"'{device.type}' -- running in fp32.", file=sys.stderr)
+        return None
+    if dtype_arg == "fp32":
+        return None  # no autocast
+    if dtype_arg == "fp16":
+        return torch.float16
+    if dtype_arg == "bf16":
+        return torch.bfloat16
+    # auto
+    return torch.float16 if device.type == "cuda" else None
+
+
+def main() -> None:
+    args = build_arg_parser().parse_args()
+    if not args.prompt and not args.interactive:
+        print("Provide --prompt \"...\" or pass --interactive.", file=sys.stderr)
+        sys.exit(1)
+
+    device = pick_device(args.device)
+    model, model_config = load_model(args.checkpoint, device)
+    tokenizer = TokenizerWrapper(args.tokenizer)
+    autocast_dtype = _resolve_dtype(args.dtype, device)
+
+    def run(prompt: str, stream: bool) -> list[str]:
+        ctx = torch.autocast(device_type=device.type, dtype=autocast_dtype) if autocast_dtype else contextlib.nullcontext()
+        with ctx:
+            return generate(
+                model, tokenizer, prompt, device,
+                max_new_tokens=args.max_new_tokens,
+                temperature=args.temperature,
+                top_k=args.top_k,
+                top_p=args.top_p,
+                repetition_penalty=args.repetition_penalty,
+                no_repeat_ngram_size=args.no_repeat_ngram_size,
+                stream=stream,
+                seed=args.seed,
+                num_samples=args.num_samples,
+            )
+
+    if args.interactive:
+        print("[Inference] Interactive mode. Ctrl-C or empty line to exit.", file=sys.stderr)
+        while True:
+            try:
+                prompt = input("\n>>> ")
+            except (EOFError, KeyboardInterrupt):
+                break
+            if not prompt.strip():
+                break
+            outputs = run(prompt, stream=(args.num_samples == 1))
+            if args.num_samples > 1:
+                for i, text in enumerate(outputs):
+                    print(f"\n--- sample {i} ---\n{text}")
+    else:
+        outputs = run(args.prompt, stream=(args.num_samples == 1))
+        if args.num_samples > 1:
+            for i, text in enumerate(outputs):
+                print(f"\n--- sample {i} ---\n{text}")
+
+
+if __name__ == "__main__":
+    main()

model.py

@@ -0,0 +1,1813 @@
+import functools
+import math
+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
+
+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
+
+try:
+    import fused_kernels
+except ImportError:  # pragma: no cover - fallback for package-style imports
+    try:
+        from . import fused_kernels
+    except ImportError:
+        fused_kernels = None
+
+try:
+    import flash_ops
+except ImportError:  # pragma: no cover - fallback for package-style imports
+    try:
+        from . import flash_ops
+    except ImportError:
+        flash_ops = None
+
+
+
+
+# ── nvFP4 via torchao RHT (Blackwell SM100+ only) ───────────────────────────
+# Uses NVFP4DynamicActivationNVFP4WeightConfig from torchao.prototype.mx_formats.
+# This replaces the old manual block-quantization implementation.
+# Dynamic activation scaling: per-tensor FP8 scale computed at runtime.
+# Weight scaling: per-block-16 FP4 scale (RHT = Reciprocal-Half-Tile layout).
+# Both scales are fused into a single Blackwell GEMM kernel via torch.compile.
+# No manual weight surgery needed -- torchao rewrites nn.Linear in-place via quantize_().
+
+_TORCHAO_AVAILABLE: bool = False
+_NVFP4_RHT_AVAILABLE: bool = False
+_INT8_AVAILABLE: bool = False
+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,
+    )
+    # Per-group INT8 weight quant: available as Int8WeightOnlyConfig with
+    # group_size, paired with dynamic int8 activations via a composed config.
+    # Stable API: try the grouped variant; fall back to per-channel if absent.
+    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
+
+# Layers to keep in full precision.
+# lm_head / wte: embedding + output projection -- high vocab sensitivity.
+# norm / ln_:    RMSNorm scale weights -- tiny tensors, quantizing them is pointless.
+# router / gate: MoE router logits drive routing decisions; FP4 routing = collapsed experts.
+# expert_bias:   per-expert scalar bias, not a Linear weight.
+_NVFP4_SKIP_KEYWORDS: tuple[str, ...] = (
+    "wte", "lm_head", "ln_", "ln_f", "norm",
+    "router", "gate", "expert_bias",
+)
+
+
+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
+    # Minimum size guard: torchao's RHT kernel requires in_features % 16 == 0
+    # and out_features >= 64. Skip tiny linears (e.g. MoD router D->1).
+    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:
+    """Quantize eligible Linear layers to NVFP4 (RHT, dynamic activations).
+
+    Uses torchao.quantization.quantize_() with
+    NVFP4DynamicActivationNVFP4WeightConfig, which applies:
+      - Weight:     FP4 e2m1 with per-block-16 FP8 scales (RHT layout)
+      - Activation: FP4 with per-tensor FP32 dynamic scale (computed at runtime)
+
+    The fused GEMM kernel is emitted by torch.compile on Blackwell SM100+.
+    Must be called BEFORE torch.compile(); quantize_() rewrites modules in-place.
+
+    Returns the number of quantized Linear layers.
+    """
+    before = {fqn for fqn, m in model.named_modules() if _nvfp4_filter(m, fqn)}
+    _torchao_quantize(model, _NVFP4Config(), filter_fn=_nvfp4_filter)
+    count = 0
+    for fqn, m in model.named_modules():
+        if fqn in before and isinstance(m, nn.Module):
+            if type(m).__name__ != "Linear":
+                count += 1
+    return count
+
+
+# ── Jetfire-style INT8 via torchao (T4 / Turing sm_75 and any CUDA device) ──
+# Int8DynamicActivationInt8WeightConfig applies:
+#   Weight:     INT8 symmetric per-group quantization (group_size from config).
+#               group_size=128 is the Jetfire default; smaller (32) tightens
+#               outlier containment at the cost of more scale overhead.
+#   Activation: INT8 dynamic per-tensor scale computed at runtime (no calibration).
+# Inductor fuses the INT8 GEMM + dequant into a single Triton or CUTLASS kernel.
+# Works on T4 (sm_75) and any device with INT8 tensor core support (Turing+).
+# Much cheaper than FP16 GEMM on T4 where INT8 tensor cores run at 2× the rate.
+
+# Reuse same skip list as FP4 -- same sensitive layers apply.
+_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
+    # group_size constraint: in_features must be divisible by group_size.
+    # We can't check the configured group_size here without a closure, so
+    # enforce the stricter of the two defaults (128). apply_int8_to_model
+    # re-checks with the actual group_size before calling quantize_().
+    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
+    # Mutual exclusion: nvFP4 and INT8 cannot both be active.
+    # nvFP4 only runs on SM100+; if the user enabled both on a Blackwell machine
+    # nvFP4 takes priority (higher throughput).
+    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
+        # group_size divisibility only matters for grouped weight configs;
+        # Int8DynamicActivationInt8WeightConfig is per-channel so no constraint.
+        # Keep the check anyway to skip genuinely misaligned tiny layers.
+        in_f = module.weight.shape[1]
+        if in_f % 32 != 0:  # 32 = minimum alignment for any INT8 kernel
+            return False
+        return True
+
+    # Snapshot module ids before -- torchao replaces nn.Linear with a subclass
+    # in-place; the fqn stays the same but id() changes. Class-name check is
+    # unreliable because torchao's subclass is also called "Linear" in some versions.
+    before_ids = {fqn: id(m) for fqn, m in model.named_modules() if _filter_with_gs(m, fqn)}
+
+    # Prefer Int8DynamicActivationInt8WeightConfig: dynamic INT8 on both weight
+    # and activation, fused into a single CUTLASS/Triton INT8 GEMM by Inductor.
+    # This is the correct stable API -- no constructor args needed.
+    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)
+
+    # Count by id() change: torchao replaces the module object in-place on the
+    # parent, so the same fqn now resolves to a different object.
+    count = sum(
+        1 for fqn, old_id in before_ids.items()
+        if id(dict(model.named_modules()).get(fqn)) != old_id
+    )
+    return count
+
+
+
+
+
+_alpha_config_warned: set[str] = set()
+
+
+def _alpha_config_attr(name: str, default):
+    
+    if not hasattr(alpha_config, name):
+        if name not in _alpha_config_warned:
+            _alpha_config_warned.add(name)
+            warnings.warn(
+                f"config module has no attribute '{name}' -- falling back to "
+                f"default {default!r}. If this is unexpected, check for a "
+                f"casing mismatch or rename in your config.py.",
+                stacklevel=2,
+            )
+    return getattr(alpha_config, name, default)
+
+
+try:
+    from torch.nn.functional import rms_norm as _native_rms_norm
+except (ImportError, AttributeError):
+    _native_rms_norm = None
+
+
+
+try:
+    from liger_kernel.ops.rms_norm import LigerRMSNormFunction as _LigerRMSNormFn
+except ImportError:
+    _LigerRMSNormFn = None
+
+try:
+    from liger_kernel.ops.rope import LigerRopeFunction as _LigerRopeFn
+except ImportError:
+    _LigerRopeFn = None
+
+try:
+    from liger_kernel.transformers.fused_linear_cross_entropy import (
+        LigerFusedLinearCrossEntropyLoss as _LigerFusedLinearCrossEntropyLoss,
+    )
+except ImportError:
+    _LigerFusedLinearCrossEntropyLoss = None
+
+_liger_warned: set[str] = set()
+
+
+def _liger_warn_once(key: str, msg: str) -> None:
+    if key not in _liger_warned:
+        _liger_warned.add(key)
+        warnings.warn(f"[model.py/liger] {msg}", stacklevel=3)
+
+
+def _liger_cuda_gate(x: torch.Tensor) -> bool:
+    
+    return x.device.type == "cuda"
+
+
+def _use_liger_config() -> bool:
+    
+    return bool(_alpha_config_attr('use_liger', False))
+
+
+def _use_flash_ops_master() -> bool:
+   
+    return bool(_alpha_config_attr('use_flash_ops', True))
+
+
+def _use_flash_rope_config() -> bool:
+    
+    return _use_flash_ops_master() and bool(_alpha_config_attr('use_flash_rope', True))
+
+
+def _use_flash_outproj_config() -> bool:
+    
+    return _use_flash_ops_master() and bool(_alpha_config_attr('use_flash_outproj_add_rmsnorm', True))
+
+
+def _use_flash_swiglu_config() -> bool:
+    
+    return _use_flash_ops_master() and bool(_alpha_config_attr('use_flash_swiglu', True))
+
+
+
+def _use_fused_add_rmsnorm_config() -> bool:
+    
+    
+    return bool(_alpha_config_attr('use_fused_add_rmsnorm', True))
+
+
+@torch.compiler.disable
+def _liger_rmsnorm_attempt(x: torch.Tensor, weight: torch.Tensor, eps: float):
+    
+    try:
+        out = _LigerRMSNormFn.apply(x, weight, eps)
+        kernel._record_backend("liger_rmsnorm_success")
+        return out
+    except Exception as e:  # noqa: BLE001 -- must never crash training
+        kernel._record_backend(f"liger_rmsnorm_fallback:{type(e).__name__}")
+        _liger_warn_once(
+            f"rmsnorm_fail_{type(e).__name__}",
+            f"Liger RMSNorm failed on {x.device} ({e!r}) -- falling back to native/eager RMSNorm.",
+        )
+        return None
+
+
+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):
+        
+        if _use_liger_config() and _LigerRMSNormFn is not None and _liger_cuda_gate(x):
+            out = _liger_rmsnorm_attempt(x, self.weight, self.eps)
+            if out is not None:
+                return out
+        if _native_rms_norm is not None:
+           
+            return _native_rms_norm(x, (self.dim,), self.weight, self.eps)
+        
+        x_fp32 = x.float()
+        rms = torch.sqrt(x_fp32.pow(2).mean(-1, keepdim=True) + self.eps)
+        weight = self.weight.to(x.dtype)
+        return ((x_fp32 / rms).to(x.dtype)) * weight
+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)
+
+
+def _yarn_find_correction_dim(num_rotations, dim, base=10000.0, orig_max_pos=2048):
+    
+    return (dim * math.log(orig_max_pos / (num_rotations * 2 * math.pi))) / (2 * math.log(base))
+
+
+def _yarn_find_correction_range(low_rot, high_rot, dim, base=10000.0, orig_max_pos=2048):
+    low = math.floor(_yarn_find_correction_dim(low_rot, dim, base, orig_max_pos))
+    high = math.ceil(_yarn_find_correction_dim(high_rot, dim, base, orig_max_pos))
+    return max(low, 0), min(high, dim - 1)
+
+
+def _yarn_linear_ramp_mask(low, high, dim):
+    
+    if low > high:
+        low, high = high, low
+    if low == high:
+        high += 0.001  # avoid a divide-by-zero when the ramp is degenerate
+    ramp = (torch.arange(dim, dtype=torch.float32) - low) / (high - low)
+    return torch.clamp(ramp, 0, 1)
+
+
+def yarn_attention_temperature(scale: float) -> float:
+    
+    if scale <= 1.0:
+        return 1.0
+    t = 0.1 * math.log(scale) + 1.0
+    return 1.0 / t
+
+
+@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,
+    yarn_scale: float = 1.0,
+    yarn_orig_max_pos: int | None = None,
+    yarn_alpha: float = 1.0,
+    yarn_beta: float = 32.0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if yarn_scale <= 1.0:
+     inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    else:
+        if yarn_orig_max_pos is None:
+            raise ValueError(
+                "build_rope_cache: yarn_scale > 1.0 requires yarn_orig_max_pos "
+                "(the ORIGINAL trained context length, i.e. CONTEXT/block_size "
+                "from config.py before scaling) to compute the NTK-by-parts "
+                "ramp correctly -- it can't be inferred from seq_len alone, "
+                "since seq_len here is already the SCALED target length."
+            )
+        pos_freqs = base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)
+        inv_freq_extrapolation = 1.0 / pos_freqs          # gamma=1 region: untouched
+        inv_freq_interpolation = 1.0 / (yarn_scale * pos_freqs)  # gamma=0 region: NTK/linear-interpolated
+        low, high = _yarn_find_correction_range(
+            yarn_alpha, yarn_beta, dim, base, yarn_orig_max_pos
+        )
+        if low > high:
+            warnings.warn(
+                f"YaRN correction range inverted (low={low} > high={high}) "
+                f"for dim={dim} with yarn_alpha={yarn_alpha}, yarn_beta="
+                f"{yarn_beta} -- these defaults were tuned for LLaMA-family "
+                f"head_dim=128 (per Peng et al. 2023), and can invert at "
+                f"smaller head_dim. Auto-corrected by swapping low/high so "
+                f"the ramp stays monotonic, but the resulting correction "
+                f"window differs from LLaMA's -- if generation quality at "
+                f"the extended length looks off, try tuning yarn_alpha/"
+                f"yarn_beta in config.py rather than trusting these "
+                f"defaults blindly for this model's shape."
+            )
+        
+        inv_freq_mask = 1.0 - _yarn_linear_ramp_mask(low, high, dim // 2)
+        inv_freq = (
+            inv_freq_interpolation * (1 - inv_freq_mask)
+            + inv_freq_extrapolation * inv_freq_mask
+        )
+    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:
+            max_position = t
+        max_pos = int(max_position) if max_position is not None else int(positions.max().item()) + 1 if positions.numel() > 0 else 1
+        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].to(x.dtype).unsqueeze(0).unsqueeze(0)
+        sin = sin_full[:t].to(x.dtype).unsqueeze(0).unsqueeze(0)
+    else:
+        cos = cos_full[positions].to(x.dtype).unsqueeze(0).unsqueeze(0)
+        sin = sin_full[positions].to(x.dtype).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 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 _use_flash_rope_config() and flash_ops is not None and q.is_cuda:
+        cos_t = cos_full[:t].to(q.dtype)
+        sin_t = sin_full[:t].to(q.dtype)
+        result = flash_ops.fused_rope_qk(q, k, cos_t, sin_t, rope_dim)
+        if result is not None:
+            return result
+
+    if (
+        _use_liger_config()
+        and _LigerRopeFn is not None
+        and _liger_cuda_gate(q)
+        and rope_dim == d
+        and h == hk  
+    ):
+        result = _liger_rope_attempt(q, k, cos_full[:t], sin_full[:t])
+        if result is not None:
+            return 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
+
+
+@torch.compiler.disable
+def _liger_rope_attempt(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    
+    try:
+        cos = cos.to(q.dtype)
+        sin = sin.to(q.dtype)
+        q_rot, k_rot = _LigerRopeFn.apply(q, k, cos, sin)
+        kernel._record_backend("liger_rope_success")
+        return q_rot, k_rot
+    except Exception as e:  # noqa: BLE001 -- must never crash training
+        kernel._record_backend(f"liger_rope_fallback:{type(e).__name__}")
+        _liger_warn_once(
+            f"rope_fail_{type(e).__name__}",
+            f"Liger RoPE failed on {q.device} ({e!r}) -- falling back to eager RoPE.",
+        )
+        return None
+
+
+@torch.compiler.disable
+def _liger_fused_ce_attempt(hidden: torch.Tensor, weight: torch.Tensor, targets: torch.Tensor):
+    
+    try:
+        loss_fn = _LigerFusedLinearCrossEntropyLoss(ignore_index=-1)
+        # Liger's fused linear+CE kernel expects the activation tensor first
+        # and the output projection weight second; the previous order was
+        # swapped and would pass an invalid tensor layout to the kernel.
+        loss = loss_fn(hidden.reshape(-1, hidden.size(-1)), weight, targets.reshape(-1))
+        kernel._record_backend("liger_fused_ce_success")
+        return loss
+    except Exception as e:  # noqa: BLE001 -- must never crash training
+        kernel._record_backend(f"liger_fused_ce_fallback:{type(e).__name__}")
+        _liger_warn_once(
+            f"fused_ce_fail_{type(e).__name__}",
+            f"Liger fused linear CE failed on {hidden.device} ({e!r}) -- "
+            f"falling back to eager lm_head + F.cross_entropy.",
+        )
+        return None
+
+
+def make_causal_mask(q_len: int, k_len: int, device) -> torch.Tensor:
+    
+    q_idx = torch.arange(q_len, device=device).view(q_len, 1)
+    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,), -1e9, device=device))
+
+
+def make_sliding_window_causal_mask(q_len: int, k_len: int, window_size: int, device) -> torch.Tensor:
+    
+    q_idx = torch.arange(q_len, device=device).view(q_len, 1)
+    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,), -1e9, 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 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,
+        yarn_scale: float = 1.0,
+        yarn_orig_max_pos: int | None = None,
+        yarn_alpha: float = 1.0,
+        yarn_beta: float = 32.0,
+    ) -> 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 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
+        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)
+
+        if use_rope:
+            
+            rope_cos, rope_sin = build_rope_cache(
+                max_seq_len, self.rope_dim,
+                yarn_scale=yarn_scale, yarn_orig_max_pos=yarn_orig_max_pos,
+                yarn_alpha=yarn_alpha, yarn_beta=yarn_beta,
+            )
+            self.register_buffer("rope_cos", rope_cos, persistent=False)
+            self.register_buffer("rope_sin", rope_sin, persistent=False)
+            
+            self.yarn_temp_factor = yarn_attention_temperature(yarn_scale)
+
+    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,
+    ) -> torch.Tensor | AttentionOutput:
+        batch_size, seq_len, _ = hidden_states.shape
+
+        qkv = self.qkv_proj(hidden_states)
+        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)
+
+        if self.use_rope:
+            q, k = apply_rope_qk(
+                q,
+                k,
+                cos_full=self.rope_cos,
+                sin_full=self.rope_sin,
+                rope_dim=self.rope_dim,
+                positions=positions,
+            )
+            if self.yarn_temp_factor != 1.0:
+               
+                q = q * self.yarn_temp_factor
+
+        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
+            scores = torch.matmul(q, k_eager.transpose(-2, -1)) / math.sqrt(self.head_dim)
+            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)
+            
+            weights = self.dropout(weights)
+            context = torch.matmul(weights, v_eager)
+        else:
+            context = kernel.fused_attention(
+                q,
+                k,
+                v,
+                causal=self.causal,
+                window_size=self.window_size,
+                dropout_p=self.dropout.p,
+                training=self.training,
+            )
+            weights = None
+
+        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):
+        gate_raw = self.W_gate(x)
+        value = self.W_value(x)
+
+        
+        hidden = None
+        if _use_flash_swiglu_config() and flash_ops is not None and gate_raw.is_cuda:
+            hidden = flash_ops.fused_swiglu(gate_raw, value)
+        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))
+
+    def forward(self, x):
+        B, T, D = x.shape
+
+        
+        shared_out = self.shared_expert(x)
+
+        
+        router_logits = self.router(x)
+        selection_scores = router_logits + self.expert_bias
+        full_probs = F.softmax(selection_scores.float(), dim=-1)
+        topk_probs, topk_indices = full_probs.topk(self.top_k, dim=-1)
+        gate_weights = topk_probs.to(router_logits.dtype)
+
+        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)
+            out = shared_out + routed_out
+            return out
+
+        routed_out = torch.zeros_like(x_flat)
+        for expert_id in range(self.n_routed):
+            expert = self.routed_experts[expert_id]
+            token_mask = (idx_flat == expert_id).any(dim=-1)
+            if not token_mask.any():
+                continue
+
+            selected_tokens = x_flat[token_mask]
+            expert_out = expert(selected_tokens)  # Returns delta only
+            weight = (
+                (idx_flat[token_mask] == expert_id).float() * gate_flat[token_mask]
+            ).sum(dim=-1)
+
+            routed_out[token_mask] += weight.unsqueeze(-1) * expert_out
+
+        routed_out = routed_out.view(B, T, D)
+        out = shared_out + routed_out
+        return 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,
+    yarn_scale: float = 1.0,
+    yarn_orig_max_pos: int | None = None,
+    yarn_alpha: float = 1.0,
+    yarn_beta: float = 32.0,
+):
+    
+    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,
+                yarn_scale=yarn_scale, yarn_orig_max_pos=yarn_orig_max_pos,
+                yarn_alpha=yarn_alpha, yarn_beta=yarn_beta,
+            ))
+            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,
+                    yarn_scale=yarn_scale, yarn_orig_max_pos=yarn_orig_max_pos,
+                    yarn_alpha=yarn_alpha, yarn_beta=yarn_beta,
+                ))
+                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,
+                    yarn_scale=yarn_scale, yarn_orig_max_pos=yarn_orig_max_pos,
+                    yarn_alpha=yarn_alpha, yarn_beta=yarn_beta,
+                ))
+                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,
+                yarn_scale=yarn_scale, yarn_orig_max_pos=yarn_orig_max_pos,
+                yarn_alpha=yarn_alpha, yarn_beta=yarn_beta,
+            ))
+            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=True,
+                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))
+
+    def forward(self, x):
+        """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)
+
+        
+        fused_out = None
+        raw_ctx = None
+        attn_delta = None
+       
+        use_fused_outproj = (
+            _use_flash_outproj_config()
+            and flash_ops is not None
+            and self.attn.out_proj.bias is None
+        )
+        if use_fused_outproj:
+            attn_raw = self.attn(normed, skip_out_proj=True)
+            candidate = module_output_tensor(attn_raw)
+            if candidate.is_cuda and candidate.shape[-1] == self.attn.out_proj.weight.shape[1]:
+                raw_ctx = candidate
+                
+                gemm_dtype = raw_ctx.dtype
+                fused_out = flash_ops.fused_outproj_add_rmsnorm(
+                    raw_ctx.reshape(-1, raw_ctx.shape[-1]),
+                    self.attn.out_proj.weight.to(dtype=gemm_dtype),
+                    x.reshape(-1, x.shape[-1]).to(dtype=gemm_dtype),
+                    self.ln_2.weight.to(dtype=gemm_dtype),
+                    self.ln_2.eps,
+                )
+                if fused_out is not None:
+                    residual_out, normed_out = fused_out
+                    
+                    x = residual_out.to(dtype=x.dtype).reshape(x.shape)
+                    ffn_in = normed_out.reshape(x.shape)
+            else:
+                
+                skip_out_proj_honored = self.attn.tp_size == 1
+                if skip_out_proj_honored:
+                    raw_ctx = candidate
+                    attn_delta = self.attn.out_proj(raw_ctx)
+                else:
+                    attn_delta = candidate
+
+        if fused_out is None:
+            if attn_delta is None:
+                if raw_ctx is not None:
+                   
+                    attn_delta = self.attn.out_proj(raw_ctx)
+                else:
+                    
+                    attn_raw = self.attn(normed)
+                    attn_delta = module_output_tensor(attn_raw)
+           
+
+            
+            tier2_out = None
+            if _use_fused_add_rmsnorm_config() and fused_kernels is not None:
+                tier2_out = fused_kernels.fused_add_rmsnorm(
+                    x, attn_delta, self.ln_2.weight.to(dtype=x.dtype), self.ln_2.eps
+                )
+            if tier2_out is not None:
+                x, ffn_in = tier2_out
+            else:
+                # Tier 3: fully eager, byte-for-byte what this block
+                # always ran before any fusion work this session.
+                x = x + attn_delta
+                ffn_in = self.ln_2(x)
+
+        
+        ffn_delta = self.ffn(ffn_in)
+        x = x + ffn_delta
+        return x
+
+
+class MoDBlock(nn.Module):
+    """Mixture-of-Depths via a continuous sigmoid gate -- no topk, no gather, no scatter.
+
+    WILL THE GATE ACTUALLY SELECT TOKENS?
+    ───���──────────────────────────────────
+    Short answer: not without explicit sparsity pressure.
+
+    A plain sigmoid gate starts near 0.5 (for zero-mean inputs after RMSNorm)
+    and stays fractional -- it applies 40% of the delta to some tokens, 70%
+    to others, 30% to others. The gradient never forces it toward 0 or 1
+    unless there is a cost to using non-binary values. Two mechanisms push
+    the gate toward genuine token selection:
+
+    1. mod_gate_entropy_coeff (config.py, default 0.01):
+       Adds a per-step auxiliary loss term:
+           L_entropy = -coeff * mean( gate*log(gate) + (1-gate)*log(1-gate) )
+       This is the binary entropy H(gate), negated -- maximizing binary entropy
+       means pushing gate toward 0 or 1, where H=0 (the gate is certain).
+       A gate at 0.5 has H=1.0 (maximum uncertainty); a gate at 0.05 or 0.95
+       has H≈0.29. The coefficient controls how hard this regularizer pushes:
+         0.0  -> pure sigmoid, stays fractional, no selection
+         0.01 -> gentle push; gates typically reach 0.1/0.9 range by step 2000
+         0.1  -> strong push; gates near 0/1 within ~500 steps but can destabilize
+       This term is computed in forward() and returned alongside the block output.
+       train.py adds it to the main loss, multiplied by the coefficient.
+
+    2. Natural gradient pressure from the main loss:
+       If one token position's delta reliably INCREASES the loss regardless of
+       context, the model learns to set gate≈0 there (the identity is better).
+       If another position's delta reliably DECREASES the loss, gate≈1 emerges.
+       This happens without any regularization, just more slowly -- typically
+       ~5000-10000 steps before meaningful gate bimodality appears.
+
+    The entropy loss (mechanism 1) dramatically speeds up mechanism 2 by making
+    the gate pay a cost for ambiguity. Use mod_gate_entropy_coeff=0.01 as a
+    starting point and inspect a histogram of gate values at step 1000 and 5000.
+    If the distribution is bimodal (peaks near 0 and 1), the gate is selecting.
+    If it's unimodal near 0.5, increase the coefficient.
+
+    
+
+    GATE DESIGN
+    ───────────
+        gate  = sigmoid(router(x))            # (B, T, 1) -- D->1 matmul + sigmoid
+        delta = inner(x) - x                  # (B, T, D) -- net block update
+        out   = x + gate * delta              # elementwise, fixed shape, fused
+        aux   = entropy_coeff * H_binary(gate) # scalar, added to training loss
+
+    Parameter cost: D floats (router.weight). At D=256 that's 256 params total.
+    """
+
+    def __init__(self, config, inner: nn.Module, capacity: float = 0.125):
+        super().__init__()
+        self.inner = inner
+        self.capacity = capacity  
+        self._entropy_coeff = float(getattr(config, 'mod_gate_entropy_coeff', 0.01))
+
+       
+        self.router = nn.Linear(config.n_embd, 1, bias=False)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        
+       
+        gate = torch.sigmoid(self.router(x))   # (B, T, 1)
+
+        
+        inner_out = self.inner(x)              # (B,T,D)
+        out = x + gate * (inner_out - x)       # (B,T,D); fuses delta into out
+
+        
+        if self.training and self._entropy_coeff > 0.0:
+            g = gate.float().clamp(1e-6, 1.0 - 1e-6)   # fp32 for log stability
+            
+            ent = -(g * torch.log(g) + (1.0 - g) * torch.log(1.0 - g))  # H per element
+            gate_aux = self._entropy_coeff * ent.mean()
+        else:
+            gate_aux = torch.zeros((), device=x.device, dtype=torch.float32)
+
+        return out, gate_aux
+
+    def extra_repr(self) -> str:
+        return f"gate_based=True,entropy_coeff={self._entropy_coeff}"
+
+
+class MultiTokenPrediction(nn.Module):
+    """Multi-Token Prediction auxiliary heads (Gloeckle et al. / DeepSeek-V3
+    style): predict tokens at offsets 2, 3, ... beyond the main next-token
+    prediction, using the same pre-lm_head hidden state.
+
+    Each head k predicts token t+k+1 from hidden state h_t via a separate
+    RMSNorm (so each depth can learn its own scale) followed by the shared
+    lm_head weight (weight tying: same projection used for the main head),
+    keeping parameter count increase minimal.
+
+    Loss contribution: mtp_lambda * mean(CE_1, CE_2, ..., CE_depth).
+    The main model's loss is unaffected; MTP loss is added on top.
+
+    
+    """
+
+    def __init__(self, config, depth: int = 1):
+        super().__init__()
+        self.depth = depth
+        # One RMSNorm per MTP depth; no separate linear (share lm_head.weight).
+        self.norms = nn.ModuleList([
+            RMSNorm(config.n_embd, eps=1e-5) for _ in range(depth)
+        ])
+
+    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, norm in enumerate(self.norms):
+            offset = k + 1   # additional token offset beyond the main +1 shift
+            Tk = x.size(1) - offset
+            if Tk <= 0:
+                continue
+            
+            h = norm(x[:, :Tk])                      # (B, Tk, D)
+            
+            logits_k = F.linear(h, lm_head_weight)   # (B, Tk, vocab_size)
+            tgt_k = targets[:, offset:]               # (B, Tk)
+            loss_k = F.cross_entropy(
+                logits_k.reshape(-1, logits_k.size(-1)).float(),
+                tgt_k.reshape(-1).long(),
+                ignore_index=-1,
+            )
+            total = total + loss_k    # fp32 + fp32 -- no precision loss
+            valid_depths += 1
+
+        if valid_depths == 0:
+            
+            return torch.zeros((), device=x.device, dtype=torch.float32)
+        return mtp_lambda * (total / valid_depths)   # fp32 scalar
+
+
+@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)
+    
+    use_liger: bool = _alpha_config_attr('use_liger', 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)
+
+    
+    yarn_scale: float = _alpha_config_attr('yarn_scale', 1.0)
+    yarn_orig_max_pos: int | None = _alpha_config_attr('yarn_orig_max_pos', None)
+    yarn_alpha: float = _alpha_config_attr('yarn_alpha', 1.0)
+    yarn_beta: float = _alpha_config_attr('yarn_beta', 32.0)
+
+   
+    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)
+
+class GPT(nn.Module):
+    """Full GPT language model"""
+
+    def __init__(self, config):
+        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
+
+        
+        yarn_scale = float(getattr(config, 'yarn_scale', 1.0))
+        yarn_orig_max_pos_cfg = getattr(config, 'yarn_orig_max_pos', None)
+        yarn_orig_max_pos = int(yarn_orig_max_pos_cfg) if yarn_orig_max_pos_cfg is not None else config.block_size
+        yarn_alpha = float(getattr(config, 'yarn_alpha', 1.0))
+        yarn_beta = float(getattr(config, 'yarn_beta', 32.0))
+        effective_max_seq_len = int(rope_table_len * yarn_scale) if yarn_scale > 1.0 else rope_table_len
+
+        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=effective_max_seq_len,
+            sliding_window_size=config.sliding_window_size,
+            tp_size=config.tp_size,
+            rope_dim=rope_dim,
+            yarn_scale=yarn_scale,
+            yarn_orig_max_pos=yarn_orig_max_pos,
+            yarn_alpha=yarn_alpha,
+            yarn_beta=yarn_beta,
+        )
+
+        
+        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":
+                # Wrap in MoDBlock: 3/4 of layers route only top-capacity
+                # fraction of tokens through the full AryaBlock.
+                block = MoDBlock(config, block, capacity=mod_capacity)
+            blocks.append(block)
+
+        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)
+
+        
+        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_lambda = float(getattr(config, 'mtp_lambda', 0.3))
+
+        
+        self.apply(self._init_weights)
+
+        
+        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)
+
+        # ── Quantization (mutually exclusive; nvFP4 takes priority on SM100+) ──
+        # Both quantize_() calls must happen BEFORE torch.compile() in train.py.
+        # _int8_enabled() already checks _nvfp4_enabled() and refuses if both are set.
+        self._use_nvfp4: bool = False
+        self._use_int8: bool = False
+
+        if _nvfp4_enabled():
+            n_quantized = apply_nvfp4_to_model(self)
+            self._use_nvfp4 = True
+            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, 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
+            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, gate, small linears. "
+                f"Verify val_loss vs fp16 baseline before a long run.",
+                stacklevel=2,
+            )
+
+        # Weight tying can be broken by quantized module replacement.
+        # Reassert the canonical tie after either quantization pass.
+        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 forward(self, idx, targets=None, need_logits=True):
+        
+        idx = idx.long()
+        if targets is not None:
+            targets = targets.long()
+        b, t = idx.size()
+        torch._check(
+            t <= self.config.block_size,
+            
+            lambda: f"Sequence too long: {t} > {self.config.block_size}",
+        )
+
+        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)
+
+        for block in self.transformer.h:
+            is_mod = isinstance(block, MoDBlock)
+            if self.config.gradient_checkpointing and self.training:
+                
+                out = torch.utils.checkpoint.checkpoint(block, x, use_reentrant=False)
+                if is_mod:
+                    x, gate_aux = out
+                    mod_gate_loss = mod_gate_loss + gate_aux.float()
+                else:
+                    x = out
+            else:
+                if is_mod:
+                    x, gate_aux = block(x)
+                    mod_gate_loss = mod_gate_loss + gate_aux.float()
+                else:
+                    x = block(x)
+
+        x = self.transformer.ln_f(x)
+
+        if targets is not None:
+            mtp_loss = torch.zeros((), device=x.device, dtype=torch.float32)
+            if self.mtp is not None:
+                mtp_loss = self.mtp.compute_loss(
+                    x, targets, self.lm_head.weight, mtp_lambda=self._mtp_lambda
+                )
+
+            # Keep the primary language-model CE loss separate from auxiliary MTP/MoD penalties.
+            # This lets caller-side logging and validation track the true CE signal without
+            # the auxiliary terms masking the baseline perplexity curve.
+            if not need_logits and _use_liger_config() and _LigerFusedLinearCrossEntropyLoss is not None and _liger_cuda_gate(x):
+                liger_ce = _liger_fused_ce_attempt(x, self.lm_head.weight, targets)
+                if liger_ce is not None:
+                    total_loss = liger_ce + mtp_loss + mod_gate_loss
+                    self.last_ce_loss = liger_ce.detach()
+                    self.last_total_loss = total_loss.detach()
+                    return None, total_loss
+
+            logits = F.linear(x, self.lm_head.weight)
+            ce_loss = F.cross_entropy(
+                logits.reshape(-1, logits.size(-1)).float(),
+                targets.reshape(-1),
+                ignore_index=-1
+            )
+
+            # Keep the loss scalar finite by avoiding a patched post-hoc
+            # nan_to_num that stamps out the gradient entirely. Repair the
+            # overflow sources (logit clamp/initialization) out front instead.
+            total_loss = ce_loss + mtp_loss + mod_gate_loss
+            self.last_ce_loss = ce_loss.detach()
+            self.last_total_loss = total_loss.detach()
+
+            if not need_logits:
+                logits = None
+        else:
+            logits = F.linear(x[:, [-1], :], self.lm_head.weight)  # (B, 1, vocab_size)
+            ce_loss = None
+            total_loss = None
+
+        return logits, total_loss
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens=None, temperature=1.0, top_k=None):
+        
+        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."
+                )
+        for _ in range(max_new_tokens):
+            # Crop to block size if needed
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            
+            # Forward pass
+            logits, _ = self(idx_cond)  # (B, T, vocab_size)
+            
+            # Get logits for next token (only last position)
+            logits = logits[:, -1, :] / temperature  # (B, vocab_size)
+            
+            # Top-k filtering
+            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
+
+
+def audit_fused_kernels(model, device, dtype=torch.float16, verbose=True):
+   
+    config = model.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
+    b, t = 2, 8
+    results = {}
+
+    if flash_ops is None:
+        if verbose:
+            print(
+                "[FusionAudit] flash_ops extension failed to load entirely -- "
+                "all three flash_* fusions (rope/outproj_add_rmsnorm/swiglu) "
+                "are running eager fallback for the WHOLE run, not just under "
+                "some condition. Check flash.cu compile output / nvcc "
+                "availability above for the real cause."
+            )
+        return {"flash_ops_loaded": False}
+    results["flash_ops_loaded"] = True
+
+    # -- fused RoPE --
+    try:
+        q = torch.randn(b, t, config.n_head, head_dim, device=device, dtype=dtype)
+        k = torch.randn(b, t, config.num_kv_heads or config.n_head, head_dim, device=device, dtype=dtype)
+        cos, sin = build_rope_cache(config.block_size, rope_dim)
+        cos, sin = cos.to(device), sin.to(device)
+        out = flash_ops.fused_rope_qk(q, k, cos[:t].to(dtype), sin[:t].to(dtype), rope_dim)
+        results["fused_rope"] = out is not None
+    except Exception as e:  # noqa: BLE001 -- diagnostic only, must not crash
+        results["fused_rope"] = False
+        results["fused_rope_error"] = repr(e)
+
+    # -- fused out_proj + residual-add + RMSNorm --
+    try:
+        x = torch.randn(b * t, config.n_embd, device=device, dtype=dtype)
+        w = torch.randn(config.n_embd, config.n_embd, device=device, dtype=dtype)
+        residual = torch.randn(b * t, config.n_embd, device=device, dtype=dtype)
+        norm_w = torch.ones(config.n_embd, device=device, dtype=dtype)
+        out = flash_ops.fused_outproj_add_rmsnorm(x, w, residual, norm_w, 1e-5)
+        results["fused_outproj_add_rmsnorm"] = out is not None
+    except Exception as e:  # noqa: BLE001
+        results["fused_outproj_add_rmsnorm"] = False
+        results["fused_outproj_add_rmsnorm_error"] = repr(e)
+
+    # -- fused SwiGLU --
+    try:
+        hidden = int(config.n_embd * config.ffn_mult)
+        gate = torch.randn(b * t, hidden, device=device, dtype=dtype)
+        value = torch.randn(b * t, hidden, device=device, dtype=dtype)
+        out = flash_ops.fused_swiglu(gate, value)
+        results["fused_swiglu"] = out is not None
+    except Exception as e:  # noqa: BLE001
+        results["fused_swiglu"] = False
+        results["fused_swiglu_error"] = repr(e)
+
+    if verbose:
+        print("[FusionAudit] one-time probe of fused kernels against real on-device tensors:")
+        for name in ("fused_rope", "fused_outproj_add_rmsnorm", "fused_swiglu"):
+            ok = results.get(name, False)
+            status = "ENGAGED" if ok else "FELL BACK TO EAGER"
+            print(f"  {name:28s} -> {status}")
+            if not ok and f"{name}_error" in results:
+                print(f"    reason: {results[f'{name}_error']}")
+        n_ok = sum(results.get(n, False) for n in ("fused_rope", "fused_outproj_add_rmsnorm", "fused_swiglu"))
+        if n_ok < 3:
+            print(
+                f"  [FusionAudit] {3 - n_ok}/3 fusions NOT engaging -- given "
+                f"this model is memory-bandwidth-bound (per the roofline "
+                f"analysis), a missed fusion means real intermediate-tensor "
+                f"HBM traffic that shouldn't be there. Worth fixing before "
+                f"chasing anything else."
+            )
+        else:
+            print("  [FusionAudit] all 3/3 fusions engaged -- fusion is not the bottleneck here.")
+    return results

tokenizer.py

@@ -0,0 +1,389 @@
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Sequence, List, Optional
+
+import sentencepiece as spm
+
+logger = logging.getLogger(__name__)
+
+# 32K is the well-established baseline vocab size for BPE/SentencePiece
+# LLM tokenizers (Llama-1/2, T5, Gopher, Chinchilla all use exactly this).
+# 128K+ only pays off for heavy multilingual/code coverage; for a small,
+# largely-English, narrow-domain model, 32K is the standard, safe default.
+DEFAULT_VOCAB_SIZE = 32000
+
+
+def train_sentencepiece(
+    data_files: Sequence[str],
+    model_prefix: str = 'tokenizer',
+    vocab_size: int = DEFAULT_VOCAB_SIZE,
+    model_type: str = 'bpe',
+    character_coverage: float = 0.9995,
+    byte_fallback: bool = True,
+    pad_id: int = 1,
+    unk_id: int = 0,
+    bos_id: int = 2,
+    eos_id: int = 3,
+    add_dummy_prefix: bool = True,
+    num_threads: int = 8,
+    input_sentence_size: int = 5_000_000,
+    shuffle_input_sentence: bool = True,
+    max_sentence_length: int = 16384,
+    split_digits: bool = True,
+    allow_whitespace_only_pieces: bool = True,
+    train_extremely_large_corpus: bool = False,
+) -> str:
+    """
+    Train a SentencePiece BPE tokenizer with byte-fallback — the same
+    scheme used by Llama-2, Mistral, and EuroLLM (BPE + byte_fallback via
+    SentencePiece specifically, not a hand-rolled BPE implementation).
+
+    Why SentencePiece and not a hand-written tiktoken export: SentencePiece's
+    C++ core does encode/decode and merge-rank bookkeeping internally and
+    natively — there is no manual ID-renumbering or rank-export step for
+    calling code to get wrong. (A prior tiktoken-based rewrite of this
+    tokenizer had exactly that class of bug: hand-exported merge ranks were
+    non-contiguous because special tokens occupied ids 0-3 in the source
+    vocab, silently corrupting merge-priority order and decode() mappings —
+    manifesting as repetitive garbage output like "to to to" despite a
+    healthy training loss. Delegating to SentencePiece's own encode/decode
+    removes that entire class of bug by construction.)
+
+    Notes on defaults:
+      - character_coverage < 1.0 with byte_fallback=True: rare glyphs fall
+        back to byte pieces instead of bloating the vocab with singletons.
+      - input_sentence_size + shuffle_input_sentence: without shuffling,
+        SentencePiece samples from the START of the concatenated corpus,
+        which silently biases vocab toward whichever domain file comes
+        first if you hand it multiple files back to back.
+      - split_digits: keeps numbers as individual digit tokens, which
+        generally helps arithmetic/math task tokenization consistency.
+    """
+    data_files = [str(Path(p)) for p in data_files]
+    if not data_files:
+        raise ValueError('data_files is empty')
+
+    missing = [f for f in data_files if not Path(f).exists()]
+    if missing:
+        raise FileNotFoundError(f'Missing input files: {missing}')
+
+    kwargs = dict(
+        input=','.join(data_files),
+        model_prefix=model_prefix,
+        vocab_size=int(vocab_size),
+        model_type=model_type,
+        character_coverage=character_coverage,
+        pad_id=pad_id,
+        unk_id=unk_id,
+        bos_id=bos_id,
+        eos_id=eos_id,
+        byte_fallback=byte_fallback,
+        hard_vocab_limit=False,
+        normalization_rule_name='nmt_nfkc',
+        add_dummy_prefix=add_dummy_prefix,
+        num_threads=num_threads,
+        input_sentence_size=input_sentence_size,
+        shuffle_input_sentence=shuffle_input_sentence,
+        max_sentence_length=max_sentence_length,
+        split_digits=split_digits,
+        allow_whitespace_only_pieces=allow_whitespace_only_pieces,
+        train_extremely_large_corpus=train_extremely_large_corpus,
+    )
+
+    logger.info(f"Training SentencePiece: vocab_size={vocab_size} model_type={model_type} "
+                f"files={len(data_files)}")
+    spm.SentencePieceTrainer.train(**kwargs)
+
+    model_path = f'{model_prefix}.model'
+    _validate_trained_model(
+        model_path, vocab_size,
+        expected_pad=pad_id, expected_unk=unk_id, expected_bos=bos_id, expected_eos=eos_id,
+    )
+    return model_path
+
+
+def _validate_trained_model(
+    model_path: str,
+    expected_vocab_size: int,
+    expected_pad: int,
+    expected_unk: int,
+    expected_bos: int,
+    expected_eos: int,
+) -> None:
+    """
+    Self-critique validation pass — checks the things that actually broke
+    in the previous (tiktoken) tokenizer, not just "does it load".
+    """
+    sp = spm.SentencePieceProcessor(model_file=model_path)
+
+    # 1. Vocab size sanity
+    actual_vocab = sp.vocab_size()
+    if actual_vocab != expected_vocab_size:
+        logger.warning(f"Trained vocab_size={actual_vocab} differs from requested={expected_vocab_size} "
+                        f"(hard_vocab_limit=False allows this if the corpus is small)")
+
+    # 2. Special token IDs must be EXACTLY what was requested — not just
+    # ">= 0". A previous bug class involved special-token ids silently
+    # drifting from what calling code assumed. Check explicitly, not
+    # loosely.
+    checks = [
+        ('pad', sp.pad_id(), expected_pad),
+        ('unk', sp.unk_id(), expected_unk),
+        ('bos', sp.bos_id(), expected_bos),
+        ('eos', sp.eos_id(), expected_eos),
+    ]
+    for name, actual, expected in checks:
+        if actual < 0:
+            raise ValueError(f'Trained model missing <{name}> special token')
+        if actual != expected:
+            raise ValueError(
+                f'<{name}> id drift: requested {expected}, SentencePiece '
+                f'assigned {actual}. This mismatch is exactly the class of '
+                f'bug that broke a previous tokenizer version — refusing '
+                f'to silently proceed.'
+            )
+
+    # 3. Basic round-trip: encode -> decode must reproduce recognizable text
+    probe = "The quick brown fox jumps over 42 lazy dogs. def foo(): return None"
+    ids = sp.encode(probe, out_type=int)
+    if not ids:
+        raise ValueError('Validation encode produced empty output')
+    decoded = sp.decode(ids)
+    if not decoded.strip():
+        raise ValueError('Validation round-trip produced empty decode')
+
+    # 4. SPECIFIC regression check for the actual reported failure mode:
+    # repetitive-token degenerate decode ("to to to", ",,,"). This won't
+    # catch a MODEL that's actually stuck in a repetition loop (that's a
+    # decoding-strategy issue, separate from the tokenizer), but it DOES
+    # catch a tokenizer that maps distinct ids to the same or corrupted
+    # text, which was the real bug here: encode the same repeated-word
+    # probe multiple times and confirm token ids are stable and decode
+    # is exact, not degenerating into duplicated/garbled pieces.
+    repeat_probe = "to to to , , , the the the"
+    repeat_ids = sp.encode(repeat_probe, out_type=int)
+    repeat_decoded = sp.decode(repeat_ids)
+    # Re-encoding the decoded output should reproduce the same ids
+    # (idempotency) — this is the real symptom check: a corrupted rank/id
+    # mapping breaks exactly this property even when a single encode/decode
+    # pass looks fine.
+    reencoded_ids = sp.encode(repeat_decoded, out_type=int)
+    if reencoded_ids != repeat_ids:
+        raise ValueError(
+            f'Round-trip idempotency FAILED on repeated-token probe: '
+            f'encode->decode->encode did not reproduce the same ids. '
+            f'original={repeat_ids} reencoded={reencoded_ids}. This is '
+            f'the specific failure signature of an id/rank mapping bug.'
+        )
+
+    # 5. Byte-fallback sanity: an unusual/rare unicode character must not
+    # crash and must not silently become <unk> if byte_fallback is on —
+    # it should decompose into byte pieces instead.
+    exotic_probe = "emoji test \U0001F600 and rare char \u0800"
+    exotic_ids = sp.encode(exotic_probe, out_type=int)
+    if not exotic_ids:
+        raise ValueError('Byte-fallback validation: exotic-character probe produced empty encode')
+    exotic_decoded = sp.decode(exotic_ids)
+    if not exotic_decoded.strip():
+        raise ValueError('Byte-fallback validation: exotic-character round-trip produced empty decode')
+
+    logger.info(f"✓ Validation OK: vocab={actual_vocab} probe_tokens={len(ids)} "
+                f"round-trip idempotency verified, byte-fallback verified")
+
+
+class TokenizerWrapper:
+    def __init__(self, model_path: str):
+        model_path = str(Path(model_path))
+        if not Path(model_path).exists():
+            raise FileNotFoundError(model_path)
+        self.sp = spm.SentencePieceProcessor(model_file=model_path)
+        self.vocab_size = int(self.sp.vocab_size())
+        self.pad_id = self.sp.pad_id()
+        self.unk_id = self.sp.unk_id()
+        self.bos_id = self.sp.bos_id()
+        self.eos_id = self.sp.eos_id()
+        for name, val in [('pad', self.pad_id), ('unk', self.unk_id), ('bos', self.bos_id), ('eos', self.eos_id)]:
+            if val < 0:
+                raise ValueError(f'SentencePiece model missing <{name}>')
+        self._special_ids = {self.pad_id, self.bos_id, self.eos_id}
+
+    def encode(self, text: str, add_bos: bool = True, add_eos: bool = False) -> List[int]:
+        if text is None:
+            raise ValueError('encode() received None')
+        if text == '':
+            ids: List[int] = []
+        else:
+            ids = list(self.sp.encode(text, out_type=int))
+        if add_bos:
+            ids = [self.bos_id] + ids
+        if add_eos:
+            ids = ids + [self.eos_id]
+        return ids
+
+    def encode_large_text(
+        self,
+        text: str,
+        add_bos: bool = True,
+        add_eos: bool = False,
+        chunk_chars: int = 2_000_000,
+    ) -> List[int]:
+        """OOM-safe encode for large (multi-MB+) strings, e.g. a whole
+        training corpus file read in one go. `encode()` above calls
+        self.sp.encode() on the ENTIRE string in one call -- even though
+        SentencePiece's C++ core is fast (not a slow Python BPE loop),
+        the RETURN VALUE is still built as one giant Python list[int],
+        and Python's own string handling means a large UTF-8 file can
+        already be 2-4x its byte size once decoded into a `str` object.
+        For a real, confirmed example: a 300MB single-file corpus OOM'd
+        a ~12-13GB Colab/Kaggle box within ~30 seconds calling encode()
+        on the whole file at once -- the failure was in Python-level
+        memory (str + list[int] materialization), not inside
+        SentencePiece's own encoding step.
+
+        This method processes `text` in fixed-size character chunks,
+        converting each chunk's ids to the accumulator and discarding
+        the chunk's own Python list before moving to the next chunk, so
+        peak memory is bounded by `chunk_chars` instead of len(text).
+        Chunk boundaries are placed at whitespace (never mid-word/
+        mid-token), so the output is IDENTICAL to what encode(text,
+        add_bos, add_eos) would produce on the whole string at once --
+        this is not an approximation, it's the same encoding, just
+        computed incrementally. Verified by direct comparison across
+        many chunk sizes and edge cases (empty string, no-whitespace
+        text, exact chunk-boundary alignment) against the whole-string
+        encode() path.
+
+        Use this instead of encode() for anything that might be large
+        (a whole file's contents) -- keep encode() for genuinely small
+        strings (a single sentence/prompt at inference time) where the
+        chunking overhead isn't worth paying.
+        """
+        if text is None:
+            raise ValueError('encode_large_text() received None')
+
+        if text == '':
+            ids: List[int] = []
+            if add_bos:
+                ids = [self.bos_id] + ids
+            if add_eos:
+                ids = ids + [self.eos_id]
+            return ids
+
+        all_ids: List[int] = []
+        pos = 0
+        n = len(text)
+        while pos < n:
+            end = min(pos + chunk_chars, n)
+            is_eof = (end >= n)
+            window = text[pos:end]
+            if is_eof:
+                piece = window
+                pos = end
+            else:
+                # Cut at the LAST whitespace boundary inside this
+                # window, so no token is split across two encode()
+                # calls (which would silently tokenize the seam
+                # differently than encoding the whole string at once).
+                cut = max(window.rfind(' '), window.rfind('\n'))
+                if cut <= 0:
+                    # cut == -1: no whitespace anywhere in this window
+                    # (pathological -- one token/URL longer than
+                    # chunk_chars). cut == 0: the window's FIRST
+                    # character is whitespace, which would make
+                    # `piece = window[:0]` empty and, worse, silently
+                    # drop that whitespace character (it's before the
+                    # cut, so it's never included in this piece OR the
+                    # next one, since pos would advance past it without
+                    # encoding it). Both cases: take the whole window
+                    # as one piece and advance past it -- correctness
+                    # at this one seam matters less than a silently
+                    # dropped character or a stalled loop, and this
+                    # only ever triggers when chunk_chars is smaller
+                    # than a typical word (never true at the real
+                    # ~2MB default; only reachable with a deliberately
+                    # tiny chunk_chars in tests).
+                    piece = window
+                    pos = end
+                else:
+                    piece = window[:cut]
+                    pos += cut
+            if not piece:
+                # Should be unreachable now that cut<=0 is folded into
+                # the "take whole window" branch above, but keep this
+                # as a hard backstop against any future edge case that
+                # produces an empty piece with is_eof=False -- without
+                # it, pos would never advance and the loop would spin
+                # forever.
+                if not is_eof:
+                    pos += 1
+                continue
+            piece_ids = self.sp.encode(piece, out_type=int)
+            all_ids.extend(piece_ids)
+            del piece_ids
+
+        if add_bos:
+            all_ids = [self.bos_id] + all_ids
+        if add_eos:
+            all_ids = all_ids + [self.eos_id]
+        return all_ids
+
+    def encode_batch(
+        self,
+        texts: Sequence[str],
+        add_bos: bool = True,
+        add_eos: bool = False,
+        skip_errors: bool = False,
+    ) -> List[List[int]]:
+        out: List[List[int]] = []
+        for i, t in enumerate(texts):
+            try:
+                out.append(self.encode(t, add_bos=add_bos, add_eos=add_eos))
+            except Exception as e:
+                if skip_errors:
+                    logger.warning(f"encode_batch: skipping item {i} ({e})")
+                    continue
+                raise
+        return out
+
+    def decode(self, ids: Sequence[int], skip_special_tokens: bool = True) -> str:
+        # Drop anything outside the valid piece-id range first. This is
+        # required, not cosmetic: PyTorch's ignore_index=-100 convention for
+        # masked label positions means `ids` is very commonly a raw labels
+        # tensor, and sp.decode() raises IndexError on any id < 0 or
+        # >= vocab_size instead of skipping it.
+        ids = [int(i) for i in ids if 0 <= int(i) < self.vocab_size]
+        if skip_special_tokens:
+            filtered = [i for i in ids if i not in self._special_ids]
+        else:
+            filtered = [i for i in ids if i != self.pad_id]
+        return self.sp.decode(filtered)
+
+    def decode_batch(self, batch_ids: Sequence[Sequence[int]], skip_special_tokens: bool = True) -> List[str]:
+        return [self.decode(ids, skip_special_tokens=skip_special_tokens) for ids in batch_ids]
+
+    def save_config(self, path: str) -> None:
+        Path(path).write_text(json.dumps({
+            'vocab_size': self.vocab_size,
+            'pad_id': self.pad_id,
+            'unk_id': self.unk_id,
+            'bos_id': self.bos_id,
+            'eos_id': self.eos_id,
+        }, indent=2), encoding='utf-8')
+
+    @classmethod
+    def from_config(cls, model_path: str, config_path: Optional[str] = None) -> 'TokenizerWrapper':
+        """Load and, if a config is given, verify special-id consistency against it."""
+        tok = cls(model_path)
+        if config_path and Path(config_path).exists():
+            cfg = json.loads(Path(config_path).read_text(encoding='utf-8'))
+            mismatches = {
+                k: (cfg[k], getattr(tok, k))
+                for k in ('vocab_size', 'pad_id', 'unk_id', 'bos_id', 'eos_id')
+                if k in cfg and cfg[k] != getattr(tok, k)
+            }
+            if mismatches:
+                raise ValueError(f'Tokenizer/config mismatch: {mismatches}')
+        return tok