Monody-Exp: monody-pass-7 EMA weights

README.md

@@ -0,0 +1,50 @@
+---
+license: apache-2.0
+language:
+  - en
+library_name: pytorch
+pipeline_tag: text-generation
+datasets:
+  - HuggingFaceFW/fineweb-edu
+tags:
+  - proof-of-concept
+  - base-model
+  - causal-lm
+  - gated-deltanet-2
+  - parallax-attention
+---
+
+# Monody-Exp
+
+**A proof of concept, published for comparison only.** Monody-Exp is a
+353M-parameter autoregressive language model, the baseline twin of the
+masked-diffusion model [Rondo-Exp](https://huggingface.co/KerzgrAI/Rondo-Exp),
+trained side by side with it to compare how each uses the same data over
+repeated passes. It is a small base model: it continues text, follows no
+instructions, and is not meant for real use.
+
+Both models trained on the same 88M FineWeb-Edu tokens, repeated 7 times (617M tokens), with 2,048-token windows and identical optimizer settings. These are the EMA weights after pass 7; held-out loss is on 101M unseen tokens.
+
+| | Rondo-Exp | Monody-Exp |
+|---|---|---|
+| Objective | masked diffusion | next token |
+| Token mixers | 12 bidirectional GDN-2, 4 bidirectional attention | 12 GDN-2, 4 Parallax attention |
+| Held-out loss | ≤ 3.810 (ELBO bound) | 3.632 (exact NLL) |
+| Held-out minus training loss | 0.13 | 1.59 |
+| Zero-shot accuracy, mean of 6 tasks | 44.8% | 45.6% |
+
+The six tasks are LAMBADA, ARC-Easy, PIQA, SciQ, HellaSwag and WinoGrande (1,000 examples each). Rondo's loss is an upper bound on its NLL, so it is not directly comparable to Monody's.
+
+## Run
+
+Linux with an NVIDIA GPU:
+
+```bash
+hf download KerzgrAI/Monody-Exp --local-dir Monody-Exp
+pip install -r Monody-Exp/requirements.txt
+python Monody-Exp/inference.py --prompt "The water cycle is"
+```
+
+The first run installs flash-linear-attention if it is missing. Options:
+`--max-new-tokens`, `--temperature` (0 is greedy), `--top-p`, `--top-k` and
+`--seed`.

config.json

@@ -0,0 +1,80 @@
+{
+  "allow_negative_eigenvalues": false,
+  "architecture": "Senryu2ForCausalLM",
+  "attention_dropout": 0.0,
+  "attention_head_dim": 128,
+  "attention_layer_type": "parallax",
+  "bos_token_id": 0,
+  "eos_token_id": 1,
+  "full_attention_interval": 4,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 2816,
+  "kda_lower_bound": -5.0,
+  "kda_safe_gate": true,
+  "layer_types": [
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "parallax",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "parallax",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "parallax",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "parallax"
+  ],
+  "linear_conv_kernel_dim": 4,
+  "linear_expand_v": 1.0,
+  "linear_full_rank_gate": true,
+  "linear_head_dim": 128,
+  "linear_layer_type": "gdn2",
+  "linear_num_heads": 8,
+  "linear_num_value_heads": 8,
+  "mask_token_id": null,
+  "max_position_embeddings": 8192,
+  "mla_kv_lora_rank": 256,
+  "mla_q_lora_rank": 512,
+  "mla_qk_nope_head_dim": 128,
+  "mla_v_head_dim": 128,
+  "mlp_activation": "situ_glu",
+  "model_type": "senryu2",
+  "mtp_adapter_rank": 128,
+  "mtp_loss_weight": 0.0,
+  "mtp_num_heads": 0,
+  "ngram_conv_kernel": 4,
+  "ngram_embedding_dim": 128,
+  "ngram_heads_per_order": 4,
+  "ngram_layer_index": 1,
+  "ngram_orders": [],
+  "ngram_table_rows": 2097152,
+  "num_attention_heads": 8,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 4,
+  "pad_token_id": 2,
+  "partial_rotary_factor": 0.5,
+  "residual_branches": 4,
+  "residual_gate_rank": 128,
+  "residual_mode": "gated_residual",
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000.0,
+  "shared_layer_indices": [],
+  "situ_gate_cap": 4.0,
+  "situ_up_cap": 25.0,
+  "tie_word_embeddings": true,
+  "training_sequence_length": 2048,
+  "unk_token_id": 3,
+  "use_block_attn_res": false,
+  "vocab_size": 65536,
+  "model_name": "Monody-Exp",
+  "developer": "KerzgrAI",
+  "torch_dtype": "bfloat16",
+  "weights": "ema",
+  "stage": "monody-pass-7"
+}

evaluation.json

@@ -0,0 +1,160 @@
+{
+  "elapsed_seconds": 1349.062306778971,
+  "evaluated_at_unix": 1793127069.5127065,
+  "held_out_tokens": 101191373,
+  "model": "monody",
+  "objective": "next-token NLL",
+  "optimizer_step": 2352,
+  "pass": 7,
+  "stage": "monody-pass-7",
+  "training_tokens": 616562688,
+  "unique_training_tokens": 88080384,
+  "weights": {
+    "ema": {
+      "downstream": {
+        "arc_easy": {
+          "accuracy": 0.475,
+          "accuracy_norm": 0.433,
+          "chance": 0.25043333333333334,
+          "elapsed_seconds": 7.513357853982598,
+          "examples": 1000,
+          "gold_bits_per_byte": 1.2626438911998918
+        },
+        "hellaswag": {
+          "accuracy": 0.264,
+          "accuracy_norm": 0.279,
+          "chance": 0.25,
+          "elapsed_seconds": 8.59502274915576,
+          "examples": 1000,
+          "gold_bits_per_byte": 1.1487319894695152
+        },
+        "lambada": {
+          "accuracy": 0.217,
+          "accuracy_norm": null,
+          "chance": 0.0,
+          "elapsed_seconds": 3.444128228817135,
+          "examples": 1000,
+          "gold_bits_per_byte": 1.2708033027176069,
+          "perplexity": 355.8159137348601
+        },
+        "mean_accuracy": 0.45649999999999996,
+        "mean_gold_bits_per_byte": 1.3049644023078721,
+        "piqa": {
+          "accuracy": 0.564,
+          "accuracy_norm": 0.564,
+          "chance": 0.5,
+          "elapsed_seconds": 5.2499302718788385,
+          "examples": 1000,
+          "gold_bits_per_byte": 1.3786440568762732
+        },
+        "sciq": {
+          "accuracy": 0.726,
+          "accuracy_norm": 0.636,
+          "chance": 0.25,
+          "elapsed_seconds": 8.13590192887932,
+          "examples": 1000,
+          "gold_bits_per_byte": 1.408172402357161
+        },
+        "winogrande": {
+          "accuracy": 0.493,
+          "accuracy_norm": null,
+          "chance": 0.5,
+          "elapsed_seconds": 4.800854941364378,
+          "examples": 1000,
+          "gold_bits_per_byte": 1.3607907712267844
+        }
+      },
+      "generalization_gap": 1.5875091574716764,
+      "held_out": {
+        "elapsed_seconds": 534.5517211458646,
+        "loss": 3.6317696267661397,
+        "loss_95ci": [
+          3.6271775943594977,
+          3.636572309692387
+        ],
+        "objective": "next-token NLL",
+        "perplexity": 37.77961331207202,
+        "tokens": 101191373,
+        "tokens_per_second": 189301.3697216918,
+        "windows": 49410
+      },
+      "memorization": {
+        "held_out": {
+          "continuation_tokens": 32,
+          "exact_match": 0.0,
+          "prefix_tokens": 64,
+          "token_accuracy": 0.0364990234375,
+          "windows": 256
+        },
+        "train": {
+          "continuation_tokens": 32,
+          "exact_match": 0.00390625,
+          "prefix_tokens": 64,
+          "token_accuracy": 0.0614013671875,
+          "windows": 256
+        }
+      },
+      "memorization_gap": 0.00390625,
+      "train_probe": {
+        "elapsed_seconds": 54.62738667288795,
+        "loss": 2.0442604692944633,
+        "loss_95ci": [
+          2.036346034331813,
+          2.052173914243236
+        ],
+        "objective": "next-token NLL",
+        "perplexity": 7.72344470088251,
+        "tokens": 10000384,
+        "tokens_per_second": 183065.39281995854,
+        "windows": 4883
+      }
+    },
+    "normal": {
+      "downstream": null,
+      "generalization_gap": 1.5069047089964025,
+      "held_out": {
+        "elapsed_seconds": 637.9913326767273,
+        "loss": 3.740345572211005,
+        "loss_95ci": [
+          3.7358146551118874,
+          3.745179232388174
+        ],
+        "objective": "next-token NLL",
+        "perplexity": 42.11254057450087,
+        "tokens": 101191373,
+        "tokens_per_second": 158609.32244243208,
+        "windows": 49410
+      },
+      "memorization": {
+        "held_out": {
+          "continuation_tokens": 32,
+          "exact_match": 0.0,
+          "prefix_tokens": 64,
+          "token_accuracy": 0.0335693359375,
+          "windows": 256
+        },
+        "train": {
+          "continuation_tokens": 32,
+          "exact_match": 0.00390625,
+          "prefix_tokens": 64,
+          "token_accuracy": 0.0570068359375,
+          "windows": 256
+        }
+      },
+      "memorization_gap": 0.00390625,
+      "train_probe": {
+        "elapsed_seconds": 67.42786940094084,
+        "loss": 2.2334408632146023,
+        "loss_95ci": [
+          2.223836386290301,
+          2.2425792741355606
+        ],
+        "objective": "next-token NLL",
+        "perplexity": 9.331920762841568,
+        "tokens": 10000384,
+        "tokens_per_second": 148312.32380390863,
+        "windows": 4883
+      }
+    }
+  }
+}

inference.py

@@ -0,0 +1,302 @@
+#!/usr/bin/env python3
+"""Text continuation for KerzgrAI/Rondo-Exp and KerzgrAI/Monody-Exp.
+
+Both are proof-of-concept base models: they continue text and follow no
+instructions. Rondo-Exp is masked diffusion: the continuation starts as masks
+and is revealed in random order, with up to --drafts tokens drafted and checked
+per forward (exact self-speculative decoding). Monody-Exp is its
+autoregressive twin and samples one token per forward with a cache.
+
+Run from a downloaded repo, which this script reads its model from:
+  python inference.py --prompt "The water cycle is"
+  python inference.py --prompt "Photosynthesis" --max-new-tokens 128 --temperature 0.8
+Or anywhere, downloading the model first:
+  python inference.py --repo KerzgrAI/Monody-Exp --prompt "The water cycle is"
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import subprocess
+import sys
+import time
+import warnings
+from pathlib import Path
+from types import ModuleType
+
+import torch
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+from tokenizers import Tokenizer
+
+REPOS = ("KerzgrAI/Rondo-Exp", "KerzgrAI/Monody-Exp")
+# The models trained on this flash-linear-attention commit (GDN-2 and Parallax attention).
+FLA_COMMIT = "cbb0a72efb55c18ca0ef4f298298317573ad2cb3"
+FLA_REPO = "https://github.com/fla-org/flash-linear-attention.git"
+FLA_MODULES = (
+    "fla.ops.gdn2",
+    "fla.layers.gdn2",
+    "fla.modules.rotary",
+    "fla.ops.parallax.parallel",
+)
+TINY_GDN_FILES = (
+    "tiny_gdn/__init__.py",
+    "tiny_gdn/config.py",
+    "tiny_gdn/detokenize.py",
+    "tiny_gdn/fugue.py",
+    "tiny_gdn/fugue_sampling.py",
+    "tiny_gdn/fused_ops.py",
+    "tiny_gdn/gated_residual_ops.py",
+    "tiny_gdn/haiku_layers.py",
+    "tiny_gdn/kdn/__init__.py",
+    "tiny_gdn/kdn/kernels.py",
+    "tiny_gdn/kdn/layer.py",
+    "tiny_gdn/kdn/recurrence.py",
+    "tiny_gdn/model.py",
+    "tiny_gdn/ngram.py",
+    "tiny_gdn/nn_common.py",
+    "tiny_gdn/parallax_ops.py",
+    "tiny_gdn/senryu2.py",
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Rondo-Exp / Monody-Exp text continuation"
+    )
+    parser.add_argument("--prompt", required=True, help="Text to continue")
+    parser.add_argument(
+        "--repo",
+        default=None,
+        help=f"Hub repo or local folder to load ({' or '.join(REPOS)}); default: this script's folder",
+    )
+    parser.add_argument("--max-new-tokens", type=int, default=256)
+    parser.add_argument("--temperature", type=float, default=1.0, help="0 is greedy")
+    parser.add_argument("--top-p", type=float, default=0.9)
+    parser.add_argument("--top-k", type=int, default=0, help="0 keeps every token")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument(
+        "--drafts",
+        type=int,
+        default=8,
+        help="Rondo-Exp: tokens drafted per forward; 0 reveals one token per forward",
+    )
+    args = parser.parse_args()
+    if args.max_new_tokens <= 0:
+        parser.error("--max-new-tokens must be positive")
+    if (
+        args.temperature < 0
+        or not 0 < args.top_p <= 1
+        or args.top_k < 0
+        or args.drafts < 0
+    ):
+        parser.error(
+            "--temperature >= 0, 0 < --top-p <= 1, --top-k >= 0 and --drafts >= 0"
+        )
+    return args
+
+
+def model_directory(repo: str | None) -> Path:
+    here = Path(__file__).resolve().parent
+    if repo is None:
+        if (here / "config.json").is_file():
+            return here
+        raise SystemExit(
+            f"No model next to {Path(__file__).name}; pass --repo {REPOS[0]} or --repo {REPOS[1]}"
+        )
+    if Path(repo).is_dir():
+        return Path(repo).resolve()
+    return Path(snapshot_download(repo))
+
+
+def ensure_fla() -> None:
+    try:
+        for name in FLA_MODULES:
+            importlib.import_module(name)
+        return
+    except Exception as error:  # noqa: BLE001
+        print(
+            f"Installing flash-linear-attention {FLA_COMMIT[:7]} ({error}) …",
+            file=sys.stderr,
+            flush=True,
+        )
+    pip = [sys.executable, "-m", "pip", "install"]
+    subprocess.check_call([*pip, "einops", "transformers>=4.45.0"])
+    subprocess.check_call(
+        [*pip, "--no-deps", "--force-reinstall", f"git+{FLA_REPO}@{FLA_COMMIT}"]
+    )
+    for name in [
+        name for name in sys.modules if name == "fla" or name.startswith("fla.")
+    ]:
+        del sys.modules[name]
+    importlib.invalidate_caches()
+    for name in FLA_MODULES:
+        importlib.import_module(name)
+
+
+def load_model(
+    directory: Path, device: torch.device
+) -> tuple[torch.nn.Module, ModuleType | None]:
+    """The model in ``directory`` on ``device`` in bfloat16, and fugue_sampling when it is Rondo."""
+
+    missing = [name for name in TINY_GDN_FILES if not (directory / name).is_file()]
+    if missing:
+        raise SystemExit(f"{directory} lacks the bundled code: {', '.join(missing)}")
+    sys.path.insert(0, str(directory))
+    config = importlib.import_module("tiny_gdn.config").TinyGDNConfig.from_json(
+        directory / "config.json"
+    )
+    if config.model_type == "fugue":
+        model_class = importlib.import_module("tiny_gdn.fugue").FugueForMaskedDiffusion
+        sampling = importlib.import_module("tiny_gdn.fugue_sampling")
+    else:
+        model_class = importlib.import_module("tiny_gdn.senryu2").Senryu2ForCausalLM
+        sampling = None
+    with torch.device(device):
+        model = model_class(config)
+    model = model.to(dtype=torch.bfloat16)
+    model.load_state_dict(
+        load_file(directory / "model.safetensors", device="cpu"), strict=True
+    )
+    return model.eval(), sampling
+
+
+def sample(
+    logits: torch.Tensor, args: argparse.Namespace, generator: torch.Generator
+) -> int:
+    logits = logits.float()
+    if args.temperature == 0:
+        return int(logits.argmax())
+    logits = logits / args.temperature
+    if 0 < args.top_k < logits.numel():
+        logits = logits.masked_fill(
+            logits < torch.topk(logits, args.top_k).values[-1], -torch.inf
+        )
+    if args.top_p < 1:
+        ordered, order = torch.sort(logits, descending=True)
+        drop = torch.softmax(ordered, dim=-1).cumsum(-1) > args.top_p
+        drop[1:] = drop[:-1].clone()
+        drop[0] = False
+        logits = logits.scatter(0, order, ordered.masked_fill(drop, -torch.inf))
+    return int(torch.multinomial(torch.softmax(logits, dim=-1), 1, generator=generator))
+
+
+@torch.inference_mode()
+def continue_autoregressive(model, prompt: list[int], args, generator, emit) -> int:
+    eos = int(model.config.eos_token_id)
+    room = model.config.max_position_embeddings - len(prompt)
+    output = model(
+        torch.tensor([prompt], device=generator.device),
+        return_logits=True,
+        logits_to_keep=1,
+        use_cache=True,
+    )
+    produced = 0
+    for _ in range(min(args.max_new_tokens, room)):
+        token = sample(output.logits[0, -1], args, generator)
+        if token == eos:
+            break
+        emit(token)
+        produced += 1
+        output = model(
+            torch.tensor([[token]], device=generator.device),
+            return_logits=True,
+            logits_to_keep=1,
+            past_key_values=output.past_key_values,
+            use_cache=True,
+        )
+    return produced
+
+
+@torch.inference_mode()
+def continue_diffusion(
+    model, sampling: ModuleType, prompt: list[int], args, generator, emit
+) -> int:
+    eos = int(model.config.eos_token_id)
+    mask = model.mask_token_id
+    settings = sampling.DiffusionSampling(
+        new_tokens=args.max_new_tokens,
+        block_length=None,
+        threshold=None,
+        order="random",
+        steps_per_block=None if args.drafts else args.max_new_tokens,
+        draft_tokens=args.drafts,
+        temperature=args.temperature,
+        top_k=args.top_k,
+        top_p=args.top_p,
+    )
+    shown = 0
+    forwards = 0
+
+    def show(canvas: torch.Tensor) -> bool:
+        nonlocal shown
+        committed = sampling.committed_prefix(canvas, mask).tolist()
+        for token in committed[shown:]:
+            if token == eos:
+                return True
+            emit(token)
+            shown += 1
+        return False
+
+    def on_step(state) -> bool:
+        nonlocal forwards
+        forwards = state.step
+        return show(state.tokens[0])
+
+    canvas = sampling.diffusion_generate(
+        model,
+        torch.tensor([prompt], device=generator.device),
+        settings,
+        generator=generator,
+        eos_token_id=eos,
+        on_step=on_step,
+    )
+    show(canvas[0])
+    print(f"\n[{forwards} forwards]", file=sys.stderr, end="", flush=True)
+    return shown
+
+
+def main() -> None:
+    args = parse_args()
+    warnings.filterwarnings("ignore", message=r"tl\.make_block_ptr is deprecated")
+    if not torch.cuda.is_available():
+        raise SystemExit(
+            "An NVIDIA GPU is required: the GDN-2 kernels run in Triton on CUDA."
+        )
+    device = torch.device("cuda")
+    directory = model_directory(args.repo)
+    ensure_fla()
+    model, sampling = load_model(directory, device)
+    tokenizer = Tokenizer.from_file(str(directory / "tokenizer.json"))
+    decoder = importlib.import_module("tiny_gdn.detokenize").IncrementalUtf8Decoder(
+        tokenizer
+    )
+    prompt = tokenizer.encode(args.prompt, add_special_tokens=False).ids
+    if not prompt:
+        raise SystemExit("The prompt produced no tokens")
+    room = model.config.max_position_embeddings - args.max_new_tokens
+    if room < 1:
+        raise SystemExit(
+            f"--max-new-tokens leaves no room in {model.config.max_position_embeddings} positions"
+        )
+    prompt = prompt[-room:]
+    generator = torch.Generator(device=device).manual_seed(args.seed)
+
+    def emit(token: int) -> None:
+        print(decoder.push(token), end="", flush=True)
+
+    print(args.prompt, end="", flush=True)
+    started = time.perf_counter()
+    if sampling is None:
+        produced = continue_autoregressive(model, prompt, args, generator, emit)
+    else:
+        produced = continue_diffusion(model, sampling, prompt, args, generator, emit)
+    shown = len(decoder.text)
+    print(decoder.finalize()[shown:], end="")
+    elapsed = time.perf_counter() - started
+    print(f"\n[{produced} tokens in {elapsed:.1f}s]", file=sys.stderr, flush=True)
+
+
+if __name__ == "__main__":
+    main()

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1f9a29ebdc0bfb39c4fa0cdffdee0e6754420f821896d0d39824095f2a28f8b
+size 706364032

requirements.txt

@@ -0,0 +1,11 @@
+torch>=2.7.0
+triton>=3.3
+safetensors>=0.4.0
+tokenizers>=0.20.0
+huggingface_hub>=0.26.0
+transformers>=4.45.0
+einops>=0.8.0
+numpy>=1.26.0
+
+# flash-linear-attention (commit cbb0a72, with GDN-2 and Parallax attention)
+# is installed by inference.py on first run when it is missing.

special_tokens_map.json

@@ -0,0 +1,63 @@
+{
+  "bos_token": "<|begin_of_text|>",
+  "eos_token": "<|end_of_text|>",
+  "pad_token": "<|padding|>",
+  "unk_token": "<|unknown|>",
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|tool_call|>",
+    "<|tool_response|>",
+    "<think>",
+    "</think>",
+    "<|fim_prefix|>",
+    "<|fim_middle|>",
+    "<|fim_suffix|>",
+    "<|fim_pad|>",
+    "<|no_think|>",
+    "<|think|>",
+    "<|mask|>",
+    "<|reserved_009|>",
+    "<|reserved_010|>",
+    "<|reserved_011|>",
+    "<|reserved_012|>",
+    "<|reserved_013|>",
+    "<|reserved_014|>",
+    "<|reserved_015|>",
+    "<|reserved_016|>",
+    "<|reserved_017|>",
+    "<|reserved_018|>",
+    "<|reserved_019|>",
+    "<|reserved_020|>",
+    "<|reserved_021|>",
+    "<|reserved_022|>",
+    "<|reserved_023|>",
+    "<|reserved_024|>",
+    "<|reserved_025|>",
+    "<|reserved_026|>",
+    "<|reserved_027|>",
+    "<|reserved_028|>",
+    "<|reserved_029|>",
+    "<|reserved_030|>",
+    "<|reserved_031|>",
+    "<|reserved_032|>",
+    "<|reserved_033|>",
+    "<|reserved_034|>",
+    "<|reserved_035|>",
+    "<|reserved_036|>",
+    "<|reserved_037|>",
+    "<|reserved_038|>",
+    "<|reserved_039|>",
+    "<|reserved_040|>",
+    "<|reserved_041|>",
+    "<|reserved_042|>",
+    "<|reserved_043|>",
+    "<|reserved_044|>",
+    "<|reserved_045|>",
+    "<|reserved_046|>",
+    "<|reserved_047|>",
+    "<|reserved_048|>",
+    "<|reserved_049|>"
+  ],
+  "mask_token": "<|mask|>"
+}

stage.json

@@ -0,0 +1,11 @@
+{
+  "checkpoint": "checkpoint-00002352",
+  "completed_at_unix": 1793125720.1193864,
+  "model": "monody",
+  "optimizer_step": 2352,
+  "pass": 7,
+  "skipped_optimizer_steps": 0,
+  "stage": "monody-pass-7",
+  "training_tokens": 616562688,
+  "unique_training_tokens": 88080384
+}

tiny_gdn/__init__.py

@@ -0,0 +1,20 @@
+from tiny_gdn.config import (
+    TinyGDNConfig,
+    haiku_config,
+    haiku_kdn_config,
+    prompttercet_config,
+    senryu2_1b_config,
+    senryu_config,
+)
+from tiny_gdn.model import TinyGDNForCausalLM, TinyGDNOutput
+
+__all__ = [
+    "TinyGDNConfig",
+    "TinyGDNForCausalLM",
+    "TinyGDNOutput",
+    "haiku_config",
+    "haiku_kdn_config",
+    "prompttercet_config",
+    "senryu2_1b_config",
+    "senryu_config",
+]

tiny_gdn/config.py

@@ -0,0 +1,488 @@
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, fields
+from pathlib import Path
+from typing import Any
+
+LINEAR_LAYER_TYPES = frozenset({"gdn2", "kda", "diag_kdn", "bigdn2"})
+ATTENTION_LAYER_TYPES = frozenset({"full_attention", "gated_mla", "parallax", "bidirectional"})
+# Gated-residual token-mixer pairs: Senryu-2's causal one, and Fugue's bidirectional one.
+GATED_RESIDUAL_MIXERS = frozenset({("gdn2", "parallax"), ("bigdn2", "bidirectional")})
+FUGUE_MODEL_TYPE = "fugue"
+RESIDUAL_MODES = frozenset({"pre_norm", "gated_residual"})
+DIAG_KDN_HEAD_DIMS = (64, 128)
+
+
+@dataclass(frozen=True)
+class TinyGDNConfig:
+    architecture: str = "TinyGDNForCausalLM"
+    model_type: str = "tiny_gdn"
+
+    vocab_size: int = 49_152
+    # Deep-thin sizing is deliberate: controlled sub-billion studies find
+    # depth materially more valuable than width around the 125M-150M scale.
+    hidden_size: int = 512
+    intermediate_size: int = 1_472
+    num_hidden_layers: int = 32
+
+    num_attention_heads: int = 4
+    num_key_value_heads: int = 1
+    attention_head_dim: int = 128
+    full_attention_interval: int = 4
+    attention_dropout: float = 0.0
+    partial_rotary_factor: float = 0.5
+    rope_theta: float = 1_000_000.0
+
+    linear_num_heads: int = 4
+    linear_num_value_heads: int = 4
+    linear_head_dim: int = 128
+    linear_expand_v: float = 1.0
+    linear_conv_kernel_dim: int = 4
+    allow_negative_eigenvalues: bool = False
+    linear_layer_type: str = "gdn2"
+    attention_layer_type: str = "full_attention"
+    mlp_activation: str = "swiglu"
+    use_block_attn_res: bool = False
+    kda_safe_gate: bool = True
+    kda_lower_bound: float = -5.0
+    mla_q_lora_rank: int | None = 512
+    mla_kv_lora_rank: int = 256
+    mla_qk_nope_head_dim: int = 128
+    mla_v_head_dim: int = 128
+    situ_gate_cap: float = 4.0
+    situ_up_cap: float = 25.0
+
+    # "gated_residual" replaces pre-norm with Qwen's multi-branch gated
+    # residual stream: each sublayer reads a gated mix of `residual_branches`
+    # normalized branches and writes its output back with per-branch scales.
+    residual_mode: str = "pre_norm"
+    residual_branches: int = 4
+    residual_gate_rank: int = 128
+    # GDN-2 output gate as one full-rank projection (Kimi K3) instead of fla's
+    # low-rank pair through head_v_dim.
+    linear_full_rank_gate: bool = False
+    # Hashed n-gram memory (Engram): `ngram_heads_per_order` tables per order,
+    # each with a distinct prime row count near `ngram_table_rows`, injected
+    # before block `ngram_layer_index`. An empty `ngram_orders` disables it.
+    ngram_orders: tuple[int, ...] = ()
+    ngram_heads_per_order: int = 4
+    ngram_table_rows: int = 2_097_152
+    ngram_embedding_dim: int = 128
+    ngram_layer_index: int = 1
+    ngram_conv_kernel: int = 4
+
+    max_position_embeddings: int = 32_768
+    training_sequence_length: int = 2_048
+    rms_norm_eps: float = 1e-6
+    initializer_range: float = 0.02
+    tie_word_embeddings: bool = True
+    shared_layer_indices: tuple[int, ...] = ()
+
+    # MTP is an opt-in ablation at this scale; static MTP is not assumed to
+    # improve a 150M model without a controlled pilot.
+    mtp_num_heads: int = 0
+    mtp_adapter_rank: int = 128
+    mtp_loss_weight: float = 0.0
+
+    bos_token_id: int = 0
+    eos_token_id: int = 1
+    pad_token_id: int = 2
+    unk_token_id: int = 3
+    # Masked-diffusion (Fugue) models replace noised tokens with this id.
+    mask_token_id: int | None = None
+    # Fugue only: 0 denoises whole sequences bidirectionally (Rondo); G > 0
+    # generates G-token blocks left to right over a causal clean stream
+    # (Fugue-1, docs/architecture.md "Fugue-1").
+    fugue_block_size: int = 0
+
+    def __post_init__(self) -> None:
+        for name in ("shared_layer_indices", "ngram_orders"):
+            object.__setattr__(self, name, tuple(getattr(self, name)))
+        if self.vocab_size <= 0 or self.vocab_size > 65_536:
+            raise ValueError("vocab_size must fit the uint16 token dataset")
+        if self.hidden_size != self.num_attention_heads * self.attention_head_dim:
+            raise ValueError("hidden_size must equal num_attention_heads * attention_head_dim")
+        if self.hidden_size != self.linear_num_heads * self.linear_head_dim:
+            raise ValueError("hidden_size must equal linear_num_heads * linear_head_dim")
+        if self.linear_num_value_heads < self.linear_num_heads:
+            raise ValueError("linear_num_value_heads must be at least linear_num_heads")
+        if self.linear_num_value_heads % self.linear_num_heads != 0:
+            raise ValueError("linear_num_value_heads must be divisible by linear_num_heads")
+        if self.num_attention_heads % self.num_key_value_heads != 0:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        if self.num_hidden_layers % self.full_attention_interval != 0:
+            raise ValueError("num_hidden_layers must be divisible by full_attention_interval")
+        if not 0.0 < self.partial_rotary_factor <= 1.0:
+            raise ValueError("partial_rotary_factor must be in (0, 1]")
+        rotary_dim = int(self.attention_head_dim * self.partial_rotary_factor)
+        if rotary_dim <= 0 or rotary_dim % 2:
+            raise ValueError("The partial rotary dimension must be positive and even")
+        if self.training_sequence_length > self.max_position_embeddings:
+            raise ValueError("training_sequence_length exceeds max_position_embeddings")
+        if len(set(self.shared_layer_indices)) != len(self.shared_layer_indices):
+            raise ValueError("shared_layer_indices must be unique")
+        if any(
+            index < 0 or index >= self.num_hidden_layers
+            for index in self.shared_layer_indices
+        ):
+            raise ValueError("shared_layer_indices contains an invalid layer")
+        if self.mtp_num_heads < 0:
+            raise ValueError("mtp_num_heads cannot be negative")
+        if self.mtp_num_heads and self.mtp_adapter_rank <= 0:
+            raise ValueError("mtp_adapter_rank must be positive when MTP is enabled")
+        if not 0.0 <= self.mtp_loss_weight <= 1.0:
+            raise ValueError("mtp_loss_weight must be between zero and one")
+        for token_id in (
+            self.bos_token_id,
+            self.eos_token_id,
+            self.pad_token_id,
+            self.unk_token_id,
+        ):
+            if not 0 <= token_id < self.vocab_size:
+                raise ValueError(f"Special token ID {token_id} is outside the vocabulary")
+        if self.linear_layer_type not in LINEAR_LAYER_TYPES:
+            raise ValueError(
+                f"linear_layer_type must be one of {sorted(LINEAR_LAYER_TYPES)}"
+            )
+        if self.linear_layer_type == "diag_kdn":
+            if self.linear_head_dim not in DIAG_KDN_HEAD_DIMS:
+                raise ValueError(
+                    f"diag_kdn needs linear_head_dim in {DIAG_KDN_HEAD_DIMS}"
+                )
+            if self.linear_num_value_heads != self.linear_num_heads:
+                raise ValueError("diag_kdn needs linear_num_value_heads == linear_num_heads")
+            if self.linear_expand_v != 1.0:
+                raise ValueError("diag_kdn needs linear_expand_v == 1.0")
+            if self.allow_negative_eigenvalues:
+                raise ValueError("diag_kdn derives its gain from the covariance; negative eigenvalues are unsupported")
+        if self.attention_layer_type not in ATTENTION_LAYER_TYPES:
+            raise ValueError(
+                f"attention_layer_type must be one of {sorted(ATTENTION_LAYER_TYPES)}"
+            )
+        if self.residual_mode not in RESIDUAL_MODES:
+            raise ValueError(f"residual_mode must be one of {sorted(RESIDUAL_MODES)}")
+        gated = self.residual_mode == "gated_residual"
+        if gated:
+            if self.residual_branches < 2:
+                raise ValueError("gated_residual needs at least two residual_branches")
+            if self.residual_gate_rank <= 0:
+                raise ValueError("residual_gate_rank must be positive")
+            if self.use_block_attn_res:
+                raise ValueError("gated_residual replaces block attention residuals")
+            if self.mtp_num_heads:
+                raise ValueError("gated_residual models do not support MTP heads")
+            if self.shared_layer_indices:
+                raise ValueError("gated_residual models do not support shared layers")
+            if (self.linear_layer_type, self.attention_layer_type) not in GATED_RESIDUAL_MIXERS:
+                raise ValueError(
+                    "gated_residual models pair gdn2 with parallax attention, or bigdn2 with bidirectional"
+                )
+        elif (
+            self.attention_layer_type in {"parallax", "bidirectional"}
+            or self.linear_layer_type == "bigdn2"
+            or self.ngram_orders
+        ):
+            raise ValueError(
+                "parallax and bidirectional mixers and n-gram memory need residual_mode 'gated_residual'"
+            )
+        fugue = self.model_type == FUGUE_MODEL_TYPE
+        if fugue != (self.linear_layer_type == "bigdn2"):
+            raise ValueError("Fugue models, and only Fugue models, use bidirectional bigdn2 mixers")
+        if fugue:
+            if self.mask_token_id is None or not 0 <= self.mask_token_id < self.vocab_size:
+                raise ValueError("Fugue models need a mask_token_id inside the vocabulary")
+            if self.ngram_orders:
+                raise ValueError("Fugue models have no n-gram memory")
+            if self.fugue_block_size < 0:
+                raise ValueError("fugue_block_size cannot be negative")
+            if self.fugue_block_size and self.training_sequence_length % self.fugue_block_size:
+                raise ValueError("training_sequence_length must be a multiple of fugue_block_size")
+        elif self.fugue_block_size:
+            raise ValueError("fugue_block_size applies to Fugue models only")
+        if self.linear_full_rank_gate and self.linear_layer_type not in {"gdn2", "bigdn2"}:
+            raise ValueError("linear_full_rank_gate applies to gdn2 and bigdn2 layers")
+        if self.ngram_orders:
+            if any(order < 2 for order in self.ngram_orders):
+                raise ValueError("ngram_orders must be at least 2")
+            if tuple(sorted(set(self.ngram_orders))) != tuple(self.ngram_orders):
+                raise ValueError("ngram_orders must be unique and ascending")
+            if self.ngram_heads_per_order <= 0 or self.ngram_embedding_dim <= 0:
+                raise ValueError("n-gram heads and embedding width must be positive")
+            if self.ngram_table_rows <= self.vocab_size:
+                raise ValueError("ngram_table_rows must exceed the vocabulary")
+            if not 0 <= self.ngram_layer_index < self.num_hidden_layers:
+                raise ValueError("ngram_layer_index must name a block")
+            if self.ngram_conv_kernel <= 0:
+                raise ValueError("ngram_conv_kernel must be positive")
+        if self.mlp_activation not in {"swiglu", "situ_glu"}:
+            raise ValueError("mlp_activation must be 'swiglu' or 'situ_glu'")
+        if self.mla_kv_lora_rank <= 0:
+            raise ValueError("mla_kv_lora_rank must be positive")
+        if self.mla_q_lora_rank is not None and self.mla_q_lora_rank <= 0:
+            raise ValueError("mla_q_lora_rank must be positive or None")
+        if self.mla_qk_nope_head_dim <= 0 or self.mla_v_head_dim <= 0:
+            raise ValueError("MLA head dimensions must be positive")
+        if self.kda_lower_bound >= 0.0:
+            raise ValueError("kda_lower_bound must be negative log-decay")
+        if self.situ_gate_cap <= 0.0 or self.situ_up_cap <= 0.0:
+            raise ValueError("SiTU caps must be positive")
+
+    @property
+    def layer_types(self) -> tuple[str, ...]:
+        return tuple(
+            self.attention_layer_type
+            if (index + 1) % self.full_attention_interval == 0
+            else self.linear_layer_type
+            for index in range(self.num_hidden_layers)
+        )
+
+    @property
+    def rotary_dim(self) -> int:
+        return int(self.attention_head_dim * self.partial_rotary_factor)
+
+    @property
+    def effective_num_layers(self) -> int:
+        return self.num_hidden_layers + len(self.shared_layer_indices)
+
+    @property
+    def ngram_table_count(self) -> int:
+        return len(self.ngram_orders) * self.ngram_heads_per_order
+
+    @property
+    def ngram_memory_dim(self) -> int:
+        return self.ngram_table_count * self.ngram_embedding_dim
+
+    def to_dict(self) -> dict[str, Any]:
+        payload = asdict(self)
+        payload["layer_types"] = list(self.layer_types)
+        return payload
+
+    def save_json(self, path: Path) -> None:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(
+            json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+
+    @classmethod
+    def from_json(cls, path: Path) -> TinyGDNConfig:
+        payload = json.loads(path.read_text(encoding="utf-8"))
+        payload.pop("layer_types", None)
+        for name in ("shared_layer_indices", "ngram_orders"):
+            if name in payload:
+                payload[name] = tuple(payload[name])
+        allowed = {item.name for item in fields(cls)}
+        return cls(**{key: value for key, value in payload.items() if key in allowed})
+
+
+def haiku_config(**overrides: Any) -> TinyGDNConfig:
+    """650M Haiku: 3×KDA + 1×Gated-MLA, AttnRes, SiTU-GLU, MTP."""
+
+    payload = {
+        "architecture": "HaikuForCausalLM",
+        "model_type": "haiku",
+        "vocab_size": 65_536,
+        "hidden_size": 1024,
+        "intermediate_size": 3840,
+        "num_hidden_layers": 36,
+        "num_attention_heads": 8,
+        "num_key_value_heads": 8,
+        "attention_head_dim": 128,
+        "full_attention_interval": 4,
+        "linear_num_heads": 8,
+        "linear_num_value_heads": 8,
+        "linear_head_dim": 128,
+        "linear_layer_type": "kda",
+        "attention_layer_type": "gated_mla",
+        "mlp_activation": "situ_glu",
+        "use_block_attn_res": True,
+        "mtp_num_heads": 2,
+        "mtp_adapter_rank": 128,
+        "mtp_loss_weight": 0.2,
+        "max_position_embeddings": 32_768,
+        "training_sequence_length": 2048,
+    }
+    payload.update(overrides)
+    return TinyGDNConfig(**payload)
+
+
+def haiku_kdn_config(**overrides: Any) -> TinyGDNConfig:
+    """Haiku with Diagonal KDN in place of KDA in the 27 linear layers.
+
+    Opt-in: at equal tokens it trailed KDA in Haiku pilots and trains ~23%
+    slower at full depth (docs/architecture.md, "Haiku v2").
+    """
+
+    return haiku_config(**{"linear_layer_type": "diag_kdn", **overrides})
+
+
+def senryu_config(**overrides: Any) -> TinyGDNConfig:
+    """800M Senryu: 3×Gated DeltaNet-2 + 1×Gated-MLA over 44 layers, AttnRes, SiTU-GLU.
+
+    Haiku's width, heads, tokenizer, and residual design with GDN-2's
+    channel-wise erase and write gates in the 33 linear layers. A rank-512
+    query LoRA costs as many parameters as a full 1024×1024 query, so queries
+    are full rank, and the KV latent doubles to 512. The feed-forward narrows
+    to 3392 to pay for the depth. Next-token only: MTP has no pilot win at
+    this scale (docs/architecture.md, "Senryu").
+    """
+
+    return haiku_config(
+        **{
+            "architecture": "SenryuForCausalLM",
+            "model_type": "senryu",
+            "intermediate_size": 3392,
+            "num_hidden_layers": 44,
+            "linear_layer_type": "gdn2",
+            "mla_q_lora_rank": None,
+            "mla_kv_lora_rank": 512,
+            "mtp_num_heads": 0,
+            "mtp_loss_weight": 0.0,
+            **overrides,
+        }
+    )
+
+
+def senryu2_1b_config(**overrides: Any) -> TinyGDNConfig:
+    """Senryu-2-1B: a 1.09B Senryu-2 backbone plus a 2.2B host-side n-gram memory.
+
+    28 layers of 3×Gated DeltaNet-2 (full-rank output gate) + 1×Parallax
+    attention, width 1536, SiTU-GLU 3584, a 4-branch gated residual stream,
+    and orders-2/3 hashed n-gram tables injected before block 1
+    (docs/architecture.md, "Senryu-2-1B").
+    """
+
+    payload: dict[str, Any] = {
+        "architecture": "Senryu2ForCausalLM",
+        "model_type": "senryu2",
+        "vocab_size": 65_536,
+        "hidden_size": 1536,
+        "intermediate_size": 3584,
+        "num_hidden_layers": 28,
+        "num_attention_heads": 12,
+        "num_key_value_heads": 4,
+        "attention_head_dim": 128,
+        "full_attention_interval": 4,
+        "partial_rotary_factor": 0.5,
+        "rope_theta": 1_000_000.0,
+        "linear_num_heads": 12,
+        "linear_num_value_heads": 12,
+        "linear_head_dim": 128,
+        "linear_layer_type": "gdn2",
+        "attention_layer_type": "parallax",
+        "mlp_activation": "situ_glu",
+        "residual_mode": "gated_residual",
+        "residual_branches": 4,
+        "residual_gate_rank": 128,
+        "linear_full_rank_gate": True,
+        "ngram_orders": (2, 3),
+        "ngram_heads_per_order": 4,
+        "ngram_table_rows": 2_097_152,
+        "ngram_embedding_dim": 128,
+        "ngram_layer_index": 1,
+        "ngram_conv_kernel": 4,
+        "max_position_embeddings": 32_768,
+        "training_sequence_length": 32_768,
+        "mtp_num_heads": 0,
+        "mtp_loss_weight": 0.0,
+    }
+    payload.update(overrides)
+    return TinyGDNConfig(**payload)
+
+
+def monody_config(**overrides: Any) -> TinyGDNConfig:
+    """Monody: the autoregressive control in the Rondo/Monody data-efficiency experiment.
+
+    Senryu-2's backbone at width 1024 with no n-gram memory: 16 layers of
+    3×Gated DeltaNet-2 + 1×Parallax attention, SiTU-GLU 2816, 8 query and 4 KV
+    heads, trained on 2,048-token windows (docs/architecture.md, "Fugue").
+    """
+
+    return senryu2_1b_config(
+        **{
+            "hidden_size": 1024,
+            "intermediate_size": 2816,
+            "num_hidden_layers": 16,
+            "num_attention_heads": 8,
+            "num_key_value_heads": 4,
+            "linear_num_heads": 8,
+            "linear_num_value_heads": 8,
+            "ngram_orders": (),
+            "max_position_embeddings": 8192,
+            "training_sequence_length": 2048,
+            **overrides,
+        }
+    )
+
+
+def rondo_config(**overrides: Any) -> TinyGDNConfig:
+    """Rondo: the first Fugue model, a masked-diffusion LM on Monody's backbone.
+
+    Monody's widths, depth and heads with every token mixer bidirectional: 12
+    bidirectional Gated DeltaNet-2 layers and 4 bidirectional softmax attention
+    layers, denoising `<|mask|>` (tokenizer v2, id 22). The backward scans'
+    own convolutions and decay projections leave it within 0.3% of Monody's
+    parameter count (docs/architecture.md, "Fugue").
+    """
+
+    return monody_config(
+        **{
+            "architecture": "FugueForMaskedDiffusion",
+            "model_type": FUGUE_MODEL_TYPE,
+            "linear_layer_type": "bigdn2",
+            "attention_layer_type": "bidirectional",
+            "mask_token_id": 22,
+            **overrides,
+        }
+    )
+
+
+def fugue1_1b_config(**overrides: Any) -> TinyGDNConfig:
+    """Fugue-1-1B: a 1.1B block-diffusion LM on Senryu-2-1B's backbone.
+
+    Senryu-2-1B's widths, depth and heads without n-gram memory, as a Fugue
+    model that writes 32-token blocks left to right: a causal clean stream
+    (next-token loss, and the exact cache every block reads) and a noisy stream
+    denoising each block bidirectionally against that cache, trained on
+    4,096-token windows (docs/architecture.md, "Fugue-1").
+    """
+
+    return senryu2_1b_config(
+        **{
+            "architecture": "FugueForMaskedDiffusion",
+            "model_type": FUGUE_MODEL_TYPE,
+            "linear_layer_type": "bigdn2",
+            "attention_layer_type": "bidirectional",
+            "mask_token_id": 22,
+            "fugue_block_size": 32,
+            "ngram_orders": (),
+            "max_position_embeddings": 32_768,
+            "training_sequence_length": 4096,
+            **overrides,
+        }
+    )
+
+
+def prompttercet_config(**overrides: Any) -> TinyGDNConfig:
+    """PromptTercet: a 501M Senryu-2 prompter, which writes the user's side of a chat.
+
+    Senryu-2-1B's layout at width 1024: 24 layers of 3×Gated DeltaNet-2 +
+    1×Parallax attention, SiTU-GLU 2816, 8 query and 4 KV heads, and n-gram
+    tables of 1,048,576 rows, 1.09B host-side parameters, half of Senryu-2-1B's
+    (docs/architecture.md, "PromptTercet").
+    """
+
+    return senryu2_1b_config(
+        **{
+            "hidden_size": 1024,
+            "intermediate_size": 2816,
+            "num_hidden_layers": 24,
+            "num_attention_heads": 8,
+            "num_key_value_heads": 4,
+            "linear_num_heads": 8,
+            "linear_num_value_heads": 8,
+            "ngram_table_rows": 1_048_576,
+            **overrides,
+        }
+    )

tiny_gdn/detokenize.py

@@ -0,0 +1,131 @@
+"""Byte-accurate incremental decode for GPT-2-style byte-level BPE.
+
+`tokenizer.decode(ids)` UTF-8-replaces incomplete sequences. Streaming that
+as a string prefix then breaks on the next byte: decode([20492]) is a
+replacement character and decode([20492, 294]) is `` ≈``, so the
+accumulated text is no longer a prefix and the UI re-emits the whole
+think block.
+
+This decoder concatenates raw BPE bytes and only emits complete UTF-8
+characters. The completing byte of ``≈`` yields `` ≈``; nothing is rewound.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Any
+
+
+def gpt2_bytes_to_unicode() -> dict[int, str]:
+    raw = (
+        list(range(ord("!"), ord("~") + 1))
+        + list(range(ord("¡"), ord("¬") + 1))
+        + list(range(ord("®"), ord("ÿ") + 1))
+    )
+    mapped = raw[:]
+    extra = 0
+    for byte in range(256):
+        if byte not in raw:
+            raw.append(byte)
+            mapped.append(256 + extra)
+            extra += 1
+    return dict(zip(raw, [chr(code) for code in mapped], strict=True))
+
+
+def unicode_to_gpt2_bytes() -> dict[str, int]:
+    return {char: byte for byte, char in gpt2_bytes_to_unicode().items()}
+
+
+def utf8_char_length(lead: int) -> int | None:
+    if lead < 0x80:
+        return 1
+    if 0xC2 <= lead <= 0xDF:
+        return 2
+    if 0xE0 <= lead <= 0xEF:
+        return 3
+    if 0xF0 <= lead <= 0xF4:
+        return 4
+    return None
+
+
+def split_complete_utf8(buffer: bytes) -> tuple[str, bytes]:
+    """Return (decoded prefix, leftover incomplete or invalid-lead bytes)."""
+    index = 0
+    end = len(buffer)
+    while index < end:
+        size = utf8_char_length(buffer[index])
+        if size is None:
+            break
+        if index + size > end:
+            break
+        chunk = buffer[index : index + size]
+        try:
+            chunk.decode("utf-8")
+        except UnicodeDecodeError:
+            break
+        index += size
+    if index == 0:
+        return "", buffer
+    return buffer[:index].decode("utf-8"), buffer[index:]
+
+
+def piece_to_bytes(piece: str, unicode_to_byte: Mapping[str, int]) -> bytes:
+    raw = bytearray()
+    for char in piece:
+        mapped = unicode_to_byte.get(char)
+        if mapped is None:
+            raw.extend(char.encode("utf-8"))
+        else:
+            raw.append(mapped)
+    return bytes(raw)
+
+
+class IncrementalUtf8Decoder:
+    """Push token ids; read only complete Unicode as it becomes available."""
+
+    def __init__(
+        self,
+        tokenizer: Any,
+        *,
+        skip_special_ids: Iterable[int] | None = None,
+    ) -> None:
+        vocab = tokenizer.get_vocab()
+        self._id_to_piece = {tid: piece for piece, tid in vocab.items()}
+        self._unicode_to_byte = unicode_to_gpt2_bytes()
+        self._skip_special_ids = set(skip_special_ids or ())
+        self._pending = b""
+        self._text = ""
+
+    @property
+    def text(self) -> str:
+        return self._text
+
+    def reset(self) -> None:
+        self._pending = b""
+        self._text = ""
+
+    def token_bytes(self, token_id: int) -> bytes:
+        piece = self._id_to_piece.get(token_id, "")
+        return piece_to_bytes(piece, self._unicode_to_byte)
+
+    def push(self, token_id: int) -> str:
+        if token_id in self._skip_special_ids:
+            return ""
+        self._pending += self.token_bytes(token_id)
+        complete, self._pending = split_complete_utf8(self._pending)
+        self._text += complete
+        return complete
+
+    def push_many(self, token_ids: Sequence[int]) -> str:
+        delta = []
+        for token_id in token_ids:
+            piece = self.push(token_id)
+            if piece:
+                delta.append(piece)
+        return "".join(delta)
+
+    def finalize(self) -> str:
+        if self._pending:
+            self._text += self._pending.decode("utf-8", errors="replace")
+            self._pending = b""
+        return self._text

tiny_gdn/fugue.py

@@ -0,0 +1,1012 @@
+"""Fugue: masked-diffusion language models on a bidirectional Gated DeltaNet-2 hybrid.
+
+Fugue keeps Senryu-2's block (4-branch Gated Residual stream, SiTU-GLU,
+full-rank sigmoid output gates) and makes every token mixer bidirectional:
+
+- Bidirectional GDN-2 reads each segment left to right and right to left over
+  one set of q/k/v, erase/write and output projections. Each direction has its
+  own short convolutions and low-rank decay (ViG). The backward scan runs on
+  the segment reversed in place, so both scans share ``cu_seqlens``, and it
+  reads its state before the current token (Hydra), so each token's own key
+  and value enter its output once.
+- Bidirectional attention is grouped-query softmax attention with per-head
+  QK-norm, partial RoPE and a sigmoid output gate.
+
+Masked diffusion (MDLM, LLaDA): each window draws a mask ratio t and every
+token becomes ``<|mask|>`` with probability t. The model predicts the
+originals; ``<|mask|>`` itself always gets probability zero. There is no time
+input; the mask count carries it. The trainer picks the per-token weight: 1/t
+for the ELBO, or 1 (unweighted masked cross-entropy). Evaluation always uses
+the ELBO (docs/architecture.md, "Fugue").
+
+Decoding runs equal-length canvases as a batch rather than packed, with no
+host sync, so a CUDA graph can capture each call: ``canvas_logits`` scores
+picked positions. It can also cache everything outside the block being
+denoised (Fast-dLLM's DualCache): ``refresh_block`` runs the whole canvas once
+and keeps, per layer, each scan's state and conv window where it enters the
+block and the attention keys and values; ``block_logits`` then runs only the
+block's tokens against them. Text outside the block keeps its features from
+the refresh.
+
+Fugue-1 (``fugue_block_size`` > 0) writes fixed-size blocks left to right
+instead, over a causal clean stream whose cache is exact; ``tiny_gdn.fugue_block``
+holds that mode's training forward and decoding, and the whole-sequence
+methods here refuse it.
+"""
+
+from __future__ import annotations
+
+import itertools
+import math
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from fla.modules import FusedRMSNormGated, ShortConvolution
+from fla.modules.rotary import rotary_embedding
+from fla.ops.gdn2 import chunk_gdn2
+from fla.ops.gdn2.fused_recurrent import fused_recurrent_gdn2_fwd
+from safetensors.torch import save_model
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from tiny_gdn.config import FUGUE_MODEL_TYPE, TinyGDNConfig
+from tiny_gdn.fused_ops import weighted_linear_cross_entropy
+from tiny_gdn.haiku_layers import SiTUGLU
+from tiny_gdn.model import TinyGDNOutput
+from tiny_gdn.nn_common import RMSNorm
+from tiny_gdn.senryu2 import (
+    GatedResidual,
+    RotaryTables,
+    mark_adamw_only,
+    mark_gain,
+    mark_muon_row_blocks,
+)
+
+IGNORE_INDEX = -100
+# Smallest mask ratio drawn in training; the 1/t weight stays below 1,000.
+MIN_MASK_RATIO = 1e-3
+
+
+# --------------------------------------------------------------------------- layout
+
+
+@dataclass(frozen=True)
+class SegmentLayout:
+    """Segments packed back to back in one row, and index maps within each segment."""
+
+    cu_seqlens: torch.Tensor
+    # Each position's mirror inside its segment; applying it twice is the identity.
+    reverse_index: torch.Tensor
+    # The next and previous position in the same segment; segment ends map to themselves.
+    next_index: torch.Tensor
+    previous_index: torch.Tensor
+    segment_start: torch.Tensor
+
+
+@dataclass(frozen=True)
+class ScanEdge:
+    """One scan direction of a GDN-2 layer where it enters the block: its state and q/k/v conv windows."""
+
+    state: torch.Tensor
+    convs: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+
+
+@dataclass(frozen=True)
+class GatedDeltaNetEdges:
+    forward: ScanEdge
+    backward: ScanEdge
+
+
+@dataclass(frozen=True)
+class AttentionContext:
+    """Rotated keys and values of the whole canvas; each block step overwrites the block's rows."""
+
+    keys: torch.Tensor
+    values: torch.Tensor
+
+
+LayerCache = GatedDeltaNetEdges | AttentionContext
+
+
+@dataclass
+class BlockCapture:
+    """What a decoding forward keeps for later block-only steps.
+
+    Per scan direction, in scan order, ``lead`` marks the tokens the scan reads
+    before it reaches the block (the prefix forwards, the suffix backwards) and
+    ``lead_count`` counts them. Both are tensors, so one CUDA graph serves every block.
+    """
+
+    lead: tuple[torch.Tensor, torch.Tensor]
+    lead_count: tuple[torch.Tensor, torch.Tensor]
+    edges: dict[tuple[int, int], ScanEdge] = field(default_factory=dict)
+    attention: dict[int, AttentionContext] = field(default_factory=dict)
+
+    @classmethod
+    def around(cls, positions: torch.Tensor, length: int) -> BlockCapture:
+        """The capture for the block at consecutive ``positions`` of ``length``-token rows."""
+
+        scan = torch.arange(length, device=positions.device)
+        before = positions[0]
+        after = length - 1 - positions[-1]
+        return cls(lead=(scan < before, scan < after), lead_count=(before, after))
+
+
+@dataclass(frozen=True)
+class BlockCache:
+    """Everything outside the block ``[start, stop)`` as the last full forward saw it.
+
+    ``offsets`` (the block start per row) and ``positions`` are tensors, so a
+    CUDA graph captured against one cache serves every block of the canvas
+    once ``copy_`` has loaded that block's cache into it.
+    """
+
+    layers: tuple[LayerCache, ...]
+    rotary: tuple[torch.Tensor, torch.Tensor]
+    offsets: torch.Tensor
+    positions: torch.Tensor
+
+    @property
+    def rows(self) -> int:
+        return self.offsets.shape[0]
+
+    @property
+    def width(self) -> int:
+        return self.positions.shape[0]
+
+    def tensors(self) -> Iterator[torch.Tensor]:
+        for layer in self.layers:
+            if isinstance(layer, AttentionContext):
+                yield from (layer.keys, layer.values)
+            else:
+                for edge in (layer.forward, layer.backward):
+                    yield edge.state
+                    yield from edge.convs
+        yield from self.rotary
+        yield from (self.offsets, self.positions)
+
+    def clone(self) -> BlockCache:
+        def edge(source: ScanEdge) -> ScanEdge:
+            q, k, v = (conv.clone() for conv in source.convs)
+            return ScanEdge(state=source.state.clone(), convs=(q, k, v))
+
+        return BlockCache(
+            layers=tuple(
+                AttentionContext(keys=layer.keys.clone(), values=layer.values.clone())
+                if isinstance(layer, AttentionContext)
+                else GatedDeltaNetEdges(forward=edge(layer.forward), backward=edge(layer.backward))
+                for layer in self.layers
+            ),
+            rotary=(self.rotary[0].clone(), self.rotary[1].clone()),
+            offsets=self.offsets.clone(),
+            positions=self.positions.clone(),
+        )
+
+    def copy_(self, source: BlockCache) -> None:
+        for target, value in zip(self.tensors(), source.tensors(), strict=True):
+            target.copy_(value)
+
+
+@dataclass
+class FugueContext:
+    """Per-forward layout shared by every block.
+
+    Training and evaluation pack all rows into one ``[1, length]`` sequence
+    split by ``layout.cu_seqlens``. Decoding runs ``[rows, length]`` canvases
+    that are each one whole segment as a batch (``packed=False``), so no kernel
+    needs the boundaries on the host.
+    """
+
+    length: int
+    layout: SegmentLayout
+    # Set when every segment has this length, so attention runs as one batched call.
+    segment_length: int | None
+    rotary: tuple[torch.Tensor, torch.Tensor]
+    capture: BlockCapture | None = None
+    rows: int = 1
+    packed: bool = True
+
+    @property
+    def cu_seqlens(self) -> torch.Tensor | None:
+        return self.layout.cu_seqlens if self.packed else None
+
+
+def whole_row_layout(length: int, device: torch.device) -> SegmentLayout:
+    """The layout of rows that are each one segment of ``length``, built without a host sync."""
+
+    positions = torch.arange(length, device=device)
+    return SegmentLayout(
+        cu_seqlens=torch.arange(2, device=device, dtype=torch.int32) * length,
+        reverse_index=positions.flip(0),
+        next_index=(positions + 1).clamp(max=length - 1),
+        previous_index=(positions - 1).clamp(min=0),
+        segment_start=positions.eq(0),
+    )
+
+
+def segment_layout(lengths: torch.Tensor) -> SegmentLayout:
+    lengths = lengths.to(torch.long)
+    cu_seqlens = F.pad(lengths.cumsum(0), (1, 0))
+    total = int(cu_seqlens[-1])
+    positions = torch.arange(total, device=lengths.device)
+    segment = torch.repeat_interleave(torch.arange(lengths.numel(), device=lengths.device), lengths)
+    offset = positions - cu_seqlens[segment]
+    segment_start = offset.eq(0)
+    segment_end = offset.eq(lengths[segment] - 1)
+    return SegmentLayout(
+        cu_seqlens=cu_seqlens.to(torch.int32),
+        reverse_index=cu_seqlens[segment] + cu_seqlens[segment + 1] - 1 - positions,
+        next_index=torch.where(segment_end, positions, positions + 1),
+        previous_index=torch.where(segment_start, positions, positions - 1),
+        segment_start=segment_start,
+    )
+
+
+# ------------------------------------------------------------------- token mixers
+
+
+class BidirectionalGatedDeltaNet2(nn.Module):
+    """GDN-2 in both directions over shared projections, averaged before the output gate.
+
+    S_t = (I - k_t (b_t ⊙ k_t)ᵀ) Diag(exp g_t) S_{t-1} + k_t (w_t ⊙ v_t)ᵀ runs
+    left to right and right to left; o_t = ½ q_tᵀ (S→_t + S←_{t+1}). Index 0 of
+    each per-direction module or parameter is the forward scan, index 1 the
+    backward one.
+    """
+
+    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
+        super().__init__()
+        hidden = config.hidden_size
+        self.layer_index = layer_index
+        self.num_heads = config.linear_num_heads
+        self.num_v_heads = config.linear_num_value_heads
+        self.head_k_dim = config.linear_head_dim
+        self.head_v_dim = int(config.linear_head_dim * config.linear_expand_v)
+        self.key_dim = self.num_heads * self.head_k_dim
+        self.value_dim = self.num_v_heads * self.head_v_dim
+        self.allow_neg_eigval = config.allow_negative_eigenvalues
+        kernel = config.linear_conv_kernel_dim
+
+        self.q_proj = nn.Linear(hidden, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden, self.value_dim, bias=False)
+        self.q_conv1d = nn.ModuleList(
+            ShortConvolution(self.key_dim, kernel, bias=False, activation="silu") for _ in range(2)
+        )
+        self.k_conv1d = nn.ModuleList(
+            ShortConvolution(self.key_dim, kernel, bias=False, activation="silu") for _ in range(2)
+        )
+        self.v_conv1d = nn.ModuleList(
+            ShortConvolution(self.value_dim, kernel, bias=False, activation="silu") for _ in range(2)
+        )
+        self.f_proj = nn.ModuleList(
+            nn.Sequential(
+                nn.Linear(hidden, self.head_v_dim, bias=False),
+                nn.Linear(self.head_v_dim, self.key_dim, bias=False),
+            )
+            for _ in range(2)
+        )
+        self.b_proj = nn.Linear(hidden, self.key_dim, bias=False)
+        self.w_proj = nn.Linear(hidden, self.value_dim, bias=False)
+        self.A_log = nn.Parameter(torch.empty(2, self.num_heads, dtype=torch.float32).uniform_(1, 16).log())
+        dt = torch.exp(
+            torch.rand(2, self.key_dim, dtype=torch.float32) * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
+        ).clamp(min=1e-4)
+        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
+        if config.linear_full_rank_gate:
+            self.g_proj: nn.Module = nn.Linear(hidden, self.value_dim, bias=True)
+            mark_adamw_only(self.g_proj.weight)
+        else:
+            self.g_proj = nn.Sequential(
+                nn.Linear(hidden, self.head_v_dim, bias=False),
+                nn.Linear(self.head_v_dim, self.value_dim, bias=True),
+            )
+            mark_adamw_only(*(layer.weight for layer in self.g_proj))
+        self.o_norm = FusedRMSNormGated(self.head_v_dim, activation="sigmoid", eps=config.rms_norm_eps)
+        self.o_proj = nn.Linear(self.value_dim, hidden, bias=False)
+
+        mark_gain(self.A_log, self.dt_bias)
+        mark_adamw_only(*(layer.weight for decay in self.f_proj for layer in decay))
+        for projection in (self.q_proj, self.k_proj, self.b_proj):
+            mark_muon_row_blocks(projection.weight, self.num_heads)
+        for projection in (self.v_proj, self.w_proj):
+            mark_muon_row_blocks(projection.weight, self.num_v_heads)
+
+    def convolutions(self, direction: int) -> tuple[ShortConvolution, ShortConvolution, ShortConvolution]:
+        return self.q_conv1d[direction], self.k_conv1d[direction], self.v_conv1d[direction]
+
+    def heads(
+        self,
+        direction: int,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        f: torch.Tensor,
+        b: torch.Tensor,
+        w: torch.Tensor,
+    ) -> tuple[torch.Tensor, ...]:
+        """Per-head ``(q, k, v, g, b, w)`` for the scan kernels, decay applied and key heads repeated."""
+
+        g = F.softplus(f + self.dt_bias[direction])
+        q, k, g, b = (rearrange(tensor, "... (h d) -> ... h d", d=self.head_k_dim) for tensor in (q, k, g, b))
+        v, w = (rearrange(tensor, "... (h d) -> ... h d", d=self.head_v_dim) for tensor in (v, w))
+        g = -self.A_log[direction].float().exp().unsqueeze(-1) * g
+        if self.num_v_heads > self.num_heads:
+            groups = self.num_v_heads // self.num_heads
+            q, k, g, b = (repeat(tensor, "... h d -> ... (h g) d", g=groups) for tensor in (q, k, g, b))
+        return q, k, v, g, b, w
+
+    def scan(
+        self,
+        direction: int,
+        hidden_states: torch.Tensor,
+        projected: tuple[torch.Tensor, ...],
+        layout: SegmentLayout,
+        cu_seqlens: torch.Tensor | None,
+        capture: BlockCapture | None = None,
+    ) -> torch.Tensor:
+        """One direction over every segment of ``layout``; ``cu_seqlens`` None scans rows that are one segment each."""
+
+        f = self.f_proj[direction](hidden_states).float()
+        q, k, v, f, b, w = (*projected[:3], f, *projected[3:])
+        if direction == 1:
+            q, k, v, f, b, w = (tensor.index_select(1, layout.reverse_index) for tensor in (q, k, v, f, b, w))
+        unconvolved = (q, k, v)
+        q_conv, k_conv, v_conv = self.convolutions(direction)
+        q, _ = q_conv(x=q, cu_seqlens=cu_seqlens)
+        k, _ = k_conv(x=k, cu_seqlens=cu_seqlens)
+        v, _ = v_conv(x=v, cu_seqlens=cu_seqlens)
+        if direction == 1:
+            # Each query reads the state one step early; its output moves back a step below.
+            q = q.index_select(1, layout.next_index)
+        q, k, v, g, b, w = self.heads(direction, q, k, v, f, b, w)
+        if capture is not None:
+            self._capture_edge(direction, unconvolved, (k, v, g, b, w), capture)
+        out, _ = chunk_gdn2(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            b=b,
+            w=w,
+            initial_state=None,
+            output_final_state=False,
+            use_qk_l2norm_in_kernel=True,
+            cu_seqlens=cu_seqlens,
+        )
+        if direction == 0:
+            return out
+        out = out.index_select(1, layout.previous_index).masked_fill(layout.segment_start[None, :, None, None], 0.0)
+        return out.index_select(1, layout.reverse_index)
+
+    def _capture_edge(
+        self,
+        direction: int,
+        unconvolved: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+        scanned: tuple[torch.Tensor, ...],
+        capture: BlockCapture,
+    ) -> None:
+        """Keep this direction's state and conv windows after the tokens it reads before the block.
+
+        The scan reruns over the whole row with zero key and decay past the
+        lead, which leaves the state unchanged, so its shape is the same for
+        every block.
+        """
+
+        lead, count = capture.lead[direction], capture.lead_count[direction]
+        width = self.q_conv1d[direction].kernel_size[0]
+        # A conv window holds the last `width` inputs, oldest first, zeros before the row starts.
+        slots = count - width + torch.arange(width, device=lead.device)
+        present = slots.ge(0)[None, :, None]
+        q_window, k_window, v_window = (
+            (tensor.index_select(1, slots.clamp(min=0)) * present).transpose(1, 2).contiguous()
+            for tensor in unconvolved
+        )
+        k, v, g, b, w = scanned
+        keep = lead[None, :, None, None]
+        k = k * keep
+        _, state = chunk_gdn2(
+            q=k, k=k, v=v, g=g * keep, b=b, w=w, output_final_state=True, use_qk_l2norm_in_kernel=True
+        )
+        capture.edges[(self.layer_index, direction)] = ScanEdge(state=state, convs=(q_window, k_window, v_window))
+
+    def scan_block(
+        self,
+        direction: int,
+        hidden_states: torch.Tensor,
+        projected: tuple[torch.Tensor, ...],
+        edge: ScanEdge,
+    ) -> torch.Tensor:
+        """``scan`` over ``[rows, width]`` block tokens, entering from ``edge``."""
+
+        f = self.f_proj[direction](hidden_states).float()
+        q, k, v, f, b, w = (*projected[:3], f, *projected[3:])
+        if direction == 1:
+            q, k, v, f, b, w = (tensor.flip(1) for tensor in (q, k, v, f, b, w))
+        # The kernels update conv windows in place; the cache must stay as the refresh left it.
+        q_conv, k_conv, v_conv = self.convolutions(direction)
+        q, _ = q_conv(x=q, cache=edge.convs[0].clone())
+        k, _ = k_conv(x=k, cache=edge.convs[1].clone())
+        v, _ = v_conv(x=v, cache=edge.convs[2].clone())
+        q, k, v, g, b, w = self.heads(direction, q, k, v, f, b, w)
+        if direction == 1:
+            # A leading slot with zero key, value and decay passes the edge state
+            # through, so query i reads the state before token i (Hydra).
+            q = F.pad(q, (0, 0, 0, 0, 0, 1))
+            k, v, g, b, w = (F.pad(tensor, (0, 0, 0, 0, 1, 0)) for tensor in (k, v, g, b, w))
+        out, _ = fused_recurrent_gdn2_fwd(
+            q=q.contiguous(),
+            k=k.contiguous(),
+            v=v.contiguous(),
+            g=g.contiguous(),
+            b=b.contiguous(),
+            w=w.contiguous(),
+            initial_state=edge.state,
+            use_qk_l2norm_in_kernel=True,
+        )
+        if direction == 0:
+            return out
+        return out[:, :-1].flip(1)
+
+    def project(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        b = self.b_proj(hidden_states).sigmoid()
+        if self.allow_neg_eigval:
+            b = b * 2.0
+        return (
+            self.q_proj(hidden_states),
+            self.k_proj(hidden_states),
+            self.v_proj(hidden_states),
+            b,
+            self.w_proj(hidden_states).sigmoid(),
+        )
+
+    def gated_output(self, out: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
+        gate = rearrange(self.g_proj(hidden_states), "... (h d) -> ... h d", d=self.head_v_dim)
+        out = self.o_norm(out, gate)
+        return self.o_proj(rearrange(out, "b t h d -> b t (h d)"))
+
+    def forward(self, hidden_states: torch.Tensor, context: FugueContext) -> torch.Tensor:
+        projected = self.project(hidden_states)
+        out = 0.5 * (
+            self.scan(0, hidden_states, projected, context.layout, context.cu_seqlens, context.capture)
+            + self.scan(1, hidden_states, projected, context.layout, context.cu_seqlens, context.capture)
+        )
+        return self.gated_output(out, hidden_states)
+
+    def forward_block(self, hidden_states: torch.Tensor, cache: LayerCache, block: BlockCache) -> torch.Tensor:
+        """``forward`` for ``[rows, width]`` block tokens between the cached edges."""
+
+        if not isinstance(cache, GatedDeltaNetEdges):
+            raise TypeError(f"layer {self.layer_index} needs GDN-2 edges, got {type(cache).__name__}")
+        projected = self.project(hidden_states)
+        out = 0.5 * (
+            self.scan_block(0, hidden_states, projected, cache.forward)
+            + self.scan_block(1, hidden_states, projected, cache.backward)
+        )
+        return self.gated_output(out, hidden_states)
+
+
+class BidirectionalAttention(nn.Module):
+    """Grouped-query softmax attention over each whole segment.
+
+    Per-head zero-centered RMSNorm on q and k, partial RoPE, and a full-rank
+    sigmoid output gate, as in Senryu-2's Parallax layers without the causal mask.
+    """
+
+    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
+        super().__init__()
+        self.layer_index = layer_index
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.attention_head_dim
+        self.scale = self.head_dim**-0.5
+        query_width = self.num_heads * self.head_dim
+        kv_width = self.num_kv_heads * self.head_dim
+        hidden = config.hidden_size
+        self.q_proj = nn.Linear(hidden, query_width, bias=False)
+        self.k_proj = nn.Linear(hidden, kv_width, bias=False)
+        self.v_proj = nn.Linear(hidden, kv_width, bias=False)
+        self.g_proj = nn.Linear(hidden, query_width, bias=True)
+        self.o_proj = nn.Linear(query_width, hidden, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        mark_muon_row_blocks(self.q_proj.weight, self.num_heads)
+        for projection in (self.k_proj, self.v_proj):
+            mark_muon_row_blocks(projection.weight, self.num_kv_heads)
+        mark_adamw_only(self.g_proj.weight)
+
+    def attend(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+        """``[n, T, H, D]`` queries over ``[n, T, H_kv, D]`` keys/values, no mask."""
+
+        out = F.scaled_dot_product_attention(
+            q.transpose(1, 2),
+            k.transpose(1, 2),
+            v.transpose(1, 2),
+            scale=self.scale,
+            enable_gqa=self.num_kv_heads != self.num_heads,
+        )
+        return out.transpose(1, 2)
+
+    def forward(self, hidden_states: torch.Tensor, context: FugueContext) -> torch.Tensor:
+        rows, length = hidden_states.shape[:2]
+        q = self.q_norm(self.q_proj(hidden_states).view(rows, length, self.num_heads, self.head_dim))
+        k = self.k_norm(self.k_proj(hidden_states).view(rows, length, self.num_kv_heads, self.head_dim))
+        v = self.v_proj(hidden_states).view(rows, length, self.num_kv_heads, self.head_dim)
+        cosine, sine = context.rotary
+        q, k = (
+            rotary_embedding(tensor, cosine, sine, seqlen_offsets=0, cu_seqlens=context.cu_seqlens)
+            for tensor in (q, k)
+        )
+        if context.capture is not None:
+            context.capture.attention[self.layer_index] = AttentionContext(keys=k, values=v)
+        if context.segment_length is not None:
+            segments = rows * length // context.segment_length
+            out = self.attend(
+                q.view(segments, context.segment_length, self.num_heads, self.head_dim),
+                k.view(segments, context.segment_length, self.num_kv_heads, self.head_dim),
+                v.view(segments, context.segment_length, self.num_kv_heads, self.head_dim),
+            ).reshape(rows, length, self.num_heads, self.head_dim)
+        else:
+            bounds = context.layout.cu_seqlens.tolist()
+            out = torch.cat(
+                [
+                    self.attend(q[:, start:stop], k[:, start:stop], v[:, start:stop])
+                    for start, stop in itertools.pairwise(bounds)
+                ],
+                dim=1,
+            )
+        gate = torch.sigmoid(self.g_proj(hidden_states)).view(rows, length, self.num_heads, self.head_dim)
+        return self.o_proj((out * gate).reshape(rows, length, self.num_heads * self.head_dim))
+
+    def forward_block(self, hidden_states: torch.Tensor, cache: LayerCache, block: BlockCache) -> torch.Tensor:
+        """``[rows, width]`` block queries over the cached canvas, with the block's own keys and values."""
+
+        if not isinstance(cache, AttentionContext):
+            raise TypeError(f"layer {self.layer_index} needs attention context, got {type(cache).__name__}")
+        rows, width = hidden_states.shape[:2]
+        q = self.q_norm(self.q_proj(hidden_states).view(rows, width, self.num_heads, self.head_dim))
+        k = self.k_norm(self.k_proj(hidden_states).view(rows, width, self.num_kv_heads, self.head_dim))
+        v = self.v_proj(hidden_states).view(rows, width, self.num_kv_heads, self.head_dim)
+        cosine, sine = block.rotary
+        q, k = (rotary_embedding(tensor, cosine, sine, seqlen_offsets=block.offsets) for tensor in (q, k))
+        cache.keys.index_copy_(1, block.positions, k)
+        cache.values.index_copy_(1, block.positions, v)
+        out = self.attend(q, cache.keys, cache.values)
+        gate = torch.sigmoid(self.g_proj(hidden_states)).view(rows, width, self.num_heads, self.head_dim)
+        return self.o_proj((out * gate).reshape(rows, width, self.num_heads * self.head_dim))
+
+
+# ------------------------------------------------------------------------- blocks
+
+
+class FugueBlock(nn.Module):
+    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
+        super().__init__()
+        self.layer_index = layer_index
+        self.layer_type = config.layer_types[layer_index]
+        residual = (config.residual_branches, config.hidden_size, config.residual_gate_rank, config.rms_norm_eps)
+        self.mixer_residual = GatedResidual(*residual, write=True)
+        self.mlp_residual = GatedResidual(*residual, write=True)
+        if self.layer_type == "bigdn2":
+            self.token_mixer: nn.Module = BidirectionalGatedDeltaNet2(config, layer_index)
+        elif self.layer_type == "bidirectional":
+            self.token_mixer = BidirectionalAttention(config, layer_index)
+        else:
+            raise ValueError(f"Fugue has no {self.layer_type} layers")
+        self.mlp = SiTUGLU(config)
+        mark_muon_row_blocks(self.mlp.gate_up_proj.weight, 2)
+
+    def mix(self, residual: torch.Tensor, mixer: Callable[[torch.Tensor], torch.Tensor]) -> torch.Tensor:
+        hidden_states, scales = self.mixer_residual.read(residual)
+        residual = GatedResidual.write(residual, scales, mixer(hidden_states).reshape(hidden_states.shape))
+        hidden_states, scales = self.mlp_residual.read(residual)
+        return GatedResidual.write(residual, scales, self.mlp(hidden_states))
+
+    def forward(self, residual: torch.Tensor, context: FugueContext) -> torch.Tensor:
+        return self.mix(
+            residual, lambda hidden: self.token_mixer(hidden.view(context.rows, context.length, -1), context)
+        )
+
+    def forward_block(self, residual: torch.Tensor, cache: LayerCache, block: BlockCache) -> torch.Tensor:
+        return self.mix(
+            residual,
+            lambda hidden: self.token_mixer.forward_block(hidden.view(block.rows, block.width, -1), cache, block),
+        )
+
+
+# -------------------------------------------------------------------------- noise
+
+
+def stratified_mask_ratios(
+    count: int,
+    *,
+    generator: torch.Generator,
+    device: torch.device | str,
+    minimum: float = MIN_MASK_RATIO,
+) -> torch.Tensor:
+    """``count`` low-discrepancy mask ratios in ``[minimum, 1]``: one uniform offset, evenly spaced."""
+
+    offset = torch.rand((), generator=generator, device=device)
+    strata = (offset + torch.arange(count, device=device, dtype=torch.float32) / count) % 1.0
+    order = torch.randperm(count, generator=generator, device=device)
+    return minimum + (1.0 - minimum) * strata[order]
+
+
+LOSS_WEIGHTINGS = ("elbo", "unweighted")
+
+
+def mask_tokens(
+    tokens: torch.Tensor,
+    valid: torch.Tensor,
+    ratios: torch.Tensor,
+    *,
+    mask_token_id: int,
+    generator: torch.Generator,
+    weighting: str = "elbo",
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Forward noise: ``(noisy ids, labels with -100 where kept, per-token weights)`` for ``[B, T]`` tokens.
+
+    Weights are 1/t for the ELBO and 1 for unweighted masked cross-entropy.
+    """
+
+    if weighting not in LOSS_WEIGHTINGS:
+        raise ValueError(f"weighting must be one of: {', '.join(LOSS_WEIGHTINGS)}")
+    draws = torch.rand(tokens.shape, generator=generator, device=tokens.device)
+    masked = (draws < ratios[:, None]) & valid
+    noisy = torch.where(masked, torch.full_like(tokens, mask_token_id), tokens)
+    labels = torch.where(masked, tokens, torch.full_like(tokens, IGNORE_INDEX))
+    row_weights = 1.0 / ratios if weighting == "elbo" else torch.ones_like(ratios)
+    return noisy, labels, row_weights[:, None].expand_as(tokens)
+
+
+def mask_exact_counts(
+    tokens: torch.Tensor,
+    valid: torch.Tensor,
+    counts: torch.Tensor,
+    *,
+    mask_token_id: int,
+    generator: torch.Generator,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Mask exactly ``counts[b]`` valid positions of each row, uniformly: ``(noisy ids, labels)``.
+
+    With ``l ~ U{1..L}`` masked and weight L/l, this is LLaDA's lower-variance
+    ELBO estimator; per token it is the mean CE over the masked positions.
+    """
+
+    draws = torch.rand(tokens.shape, generator=generator, device=tokens.device).masked_fill(~valid, 2.0)
+    ranks = draws.argsort(dim=1).argsort(dim=1)
+    masked = ranks < counts[:, None]
+    noisy = torch.where(masked, torch.full_like(tokens, mask_token_id), tokens)
+    labels = torch.where(masked, tokens, torch.full_like(tokens, IGNORE_INDEX))
+    return noisy, labels
+
+
+# -------------------------------------------------------------------------- model
+
+
+class FugueForMaskedDiffusion(nn.Module):
+    """Bidirectional GDN-2 hybrid that predicts the tokens hidden behind ``<|mask|>``."""
+
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        if config.model_type != FUGUE_MODEL_TYPE or config.mask_token_id is None:
+            raise ValueError("FugueForMaskedDiffusion needs a Fugue config with a mask_token_id")
+        self.config = config
+        self.mask_token_id = config.mask_token_id
+        # Fugue-1's block size; 0 for whole-sequence denoising (Rondo).
+        self.block_size = config.fugue_block_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(FugueBlock(config, index) for index in range(config.num_hidden_layers))
+        self.final_residual = GatedResidual(
+            config.residual_branches,
+            config.hidden_size,
+            config.residual_gate_rank,
+            config.rms_norm_eps,
+            write=False,
+        )
+        self.rotary = RotaryTables(config.rotary_dim, config.rope_theta)
+        self.checkpointed_layers: frozenset[int] = frozenset()
+        self.apply(self._initialize_module)
+        self._initialize_residual_projections()
+
+    # ---------------------------------------------------------------- setup
+    def _initialize_module(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+
+    def _initialize_residual_projections(self) -> None:
+        residual_std = self.config.initializer_range / math.sqrt(2 * self.config.num_hidden_layers)
+        for layer in self.layers:
+            nn.init.normal_(layer.token_mixer.o_proj.weight, mean=0.0, std=residual_std)
+            nn.init.normal_(layer.mlp.down_proj.weight, mean=0.0, std=residual_std)
+
+    def enable_gradient_checkpointing(self, enabled: bool = True, *, layers: int | None = None) -> None:
+        """Recompute block activations in backward; ``layers`` picks GDN-2 blocks first."""
+
+        if not enabled:
+            if layers is not None:
+                raise ValueError("layers needs gradient checkpointing enabled")
+            self.checkpointed_layers = frozenset()
+            return
+        count = len(self.layers) if layers is None else layers
+        if not 1 <= count <= len(self.layers):
+            raise ValueError(f"layers must be between 1 and {len(self.layers)}")
+        cheapest_replay_first = sorted(
+            range(len(self.layers)),
+            key=lambda index: self.layers[index].layer_type != "bigdn2",
+        )
+        self.checkpointed_layers = frozenset(cheapest_replay_first[:count])
+
+    def describe_gradient_checkpointing(self) -> str:
+        if not self.checkpointed_layers:
+            return "off"
+        layer_types = sorted({self.layers[index].layer_type for index in self.checkpointed_layers})
+        return f"{len(self.checkpointed_layers)}/{len(self.layers)} blocks ({', '.join(layer_types)})"
+
+    def project_to_vocabulary(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Logits over the vocabulary; ``<|mask|>`` is never a prediction (logit -inf)."""
+
+        logits = F.linear(hidden_states, self.embed_tokens.weight)
+        logits[..., self.mask_token_id] = float("-inf")
+        return logits
+
+    # -------------------------------------------------------------- forward
+    def _require_whole_sequence(self) -> None:
+        if self.block_size:
+            raise ValueError(
+                "this Fugue model denoises blocks over a causal clean stream "
+                f"(fugue_block_size={self.block_size}); train and decode it with tiny_gdn.fugue_block"
+            )
+
+    def _run_block(self, index: int, residual: torch.Tensor, context: FugueContext) -> torch.Tensor:
+        layer = self.layers[index]
+        if self.training and index in self.checkpointed_layers:
+            return checkpoint(layer, residual, context, use_reentrant=False)
+        return layer(residual, context)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        *,
+        labels: torch.Tensor | None = None,
+        label_weights: torch.Tensor | None = None,
+        loss_normalizer: float | None = None,
+        z_loss_coefficient: float = 0.0,
+        return_logits: bool = True,
+        logits_mask: torch.Tensor | None = None,
+        return_hidden_states: bool = False,
+    ) -> TinyGDNOutput:
+        """Score masked positions of ``[B, T]`` rows; each row is one bidirectional segment.
+
+        ``labels`` holds the original token at scored positions and -100
+        elsewhere; ``loss = Σ w·CE / loss_normalizer`` (default: the number of
+        scored positions). ``logits_mask`` limits logits to its true positions
+        and returns them flattened ``[selected, vocab]`` in row-major order.
+        """
+
+        self._require_whole_sequence()
+        if input_ids.ndim != 2:
+            raise ValueError("input_ids must have shape [batch, sequence]")
+        batch, length = input_ids.shape
+        if length > self.config.max_position_embeddings:
+            raise ValueError("Input exceeds max_position_embeddings")
+        device = input_ids.device
+        if attention_mask is None or bool(attention_mask.all()):
+            index = None
+            lengths = torch.full((batch,), length, dtype=torch.long, device=device)
+            segment_length: int | None = length
+        else:
+            if attention_mask.shape != input_ids.shape:
+                raise ValueError("attention_mask must match input_ids")
+            valid = attention_mask.reshape(-1).bool()
+            index = valid.nonzero().flatten()
+            lengths = attention_mask.bool().sum(dim=1)
+            if bool((lengths == 0).any()):
+                raise ValueError("every row needs at least one unmasked token")
+            segment_length = None
+
+        def pack(tensor: torch.Tensor) -> torch.Tensor:
+            flat = tensor.reshape(batch * length, *tensor.shape[2:])
+            return flat if index is None else flat.index_select(0, index)
+
+        tokens = pack(input_ids)
+        embeddings = self.embed_tokens(tokens)
+        context = FugueContext(
+            length=tokens.numel(),
+            layout=segment_layout(lengths),
+            segment_length=segment_length,
+            rotary=self.rotary.get(tokens.numel(), device, embeddings.dtype),
+        )
+        residual = embeddings.unsqueeze(1).expand(-1, self.config.residual_branches, -1)
+        for layer_index in range(len(self.layers)):
+            residual = self._run_block(layer_index, residual, context)
+        hidden_states, _ = self.final_residual.read(residual)
+
+        total_loss = main_loss = z_loss = None
+        if labels is not None:
+            if labels.shape != input_ids.shape:
+                raise ValueError("labels must match input_ids")
+            targets = pack(labels)
+            scored = targets.ne(IGNORE_INDEX).nonzero().flatten()
+            weights = (
+                torch.ones(scored.numel(), dtype=torch.float32, device=device)
+                if label_weights is None
+                else pack(label_weights).index_select(0, scored).float()
+            )
+            normalizer = float(scored.numel()) if loss_normalizer is None else float(loss_normalizer)
+            if scored.numel() == 0:
+                total_loss = hidden_states.sum() * 0.0
+                main_loss = total_loss.detach()
+                z_loss = total_loss.detach()
+            else:
+                objective, losses, log_partitions = weighted_linear_cross_entropy(
+                    hidden_states.index_select(0, scored),
+                    self.embed_tokens.weight,
+                    targets.index_select(0, scored),
+                    weights,
+                    z_loss_coefficient=z_loss_coefficient,
+                    excluded_token_id=self.mask_token_id,
+                )
+                total_loss = objective / max(normalizer, 1.0)
+                main_loss = (weights * losses).sum() / max(normalizer, 1.0)
+                z_loss = log_partitions.square().mean()
+
+        logits = None
+        if return_logits:
+            if logits_mask is not None:
+                selected = pack(logits_mask).nonzero().flatten()
+                logits = self.project_to_vocabulary(hidden_states.index_select(0, selected))
+            elif index is None:
+                logits = self.project_to_vocabulary(hidden_states).view(batch, length, -1)
+            else:
+                full = hidden_states.new_zeros(batch * length, hidden_states.shape[-1])
+                full.index_copy_(0, index, hidden_states)
+                logits = self.project_to_vocabulary(full).view(batch, length, -1)
+        return TinyGDNOutput(
+            loss=total_loss,
+            logits=logits,
+            main_loss=main_loss,
+            mtp_loss=None,
+            z_loss=z_loss if z_loss_coefficient > 0.0 else None,
+            hidden_states=hidden_states if return_hidden_states else None,
+        )
+
+    # ------------------------------------------------------------- decoding
+    def _decode(self, canvas: torch.Tensor, capture: BlockCapture | None) -> torch.Tensor:
+        """Final hidden states ``[rows, T, hidden]`` of unpadded canvases, run as a batch without host syncs."""
+
+        self._require_whole_sequence()
+        rows, length = canvas.shape
+        if length > self.config.max_position_embeddings:
+            raise ValueError("Input exceeds max_position_embeddings")
+        embeddings = self.embed_tokens(canvas.reshape(-1))
+        context = FugueContext(
+            length=length,
+            layout=whole_row_layout(length, canvas.device),
+            segment_length=length,
+            rotary=self.rotary.get(length, canvas.device, embeddings.dtype),
+            capture=capture,
+            rows=rows,
+            packed=False,
+        )
+        residual = embeddings.unsqueeze(1).expand(-1, self.config.residual_branches, -1)
+        for layer in self.layers:
+            residual = layer(residual, context)
+        hidden_states, _ = self.final_residual.read(residual)
+        return hidden_states.view(rows, length, -1)
+
+    def canvas_logits(self, canvas: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        """Logits ``[rows, picks, vocab]`` at ``positions`` ``[rows, picks]`` of ``[rows, T]`` canvases."""
+
+        hidden_states = self._decode(canvas, None)
+        picked = hidden_states.gather(1, positions.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1]))
+        return self.project_to_vocabulary(picked)
+
+    def refresh_block(self, canvas: torch.Tensor, positions: torch.Tensor) -> tuple[torch.Tensor, BlockCache]:
+        """Full forward of ``[rows, T]`` canvases: logits ``[rows, width, vocab]`` for the block and its cache.
+
+        ``positions`` holds the block's consecutive canvas positions, after at
+        least one token. The cache keeps ``positions`` itself, so a graph that
+        refreshes from a static tensor caches a block that later calls can move.
+        """
+
+        rows, length = canvas.shape
+        capture = BlockCapture.around(positions, length)
+        hidden_states = self._decode(canvas, capture)
+        logits = self.project_to_vocabulary(hidden_states.index_select(1, positions))
+        layers: list[LayerCache] = []
+        for index, layer in enumerate(self.layers):
+            if layer.layer_type == "bigdn2":
+                layers.append(GatedDeltaNetEdges(forward=capture.edges[(index, 0)], backward=capture.edges[(index, 1)]))
+            else:
+                layers.append(capture.attention[index])
+        cache = BlockCache(
+            layers=tuple(layers),
+            rotary=self.rotary.get(length, canvas.device, self.embed_tokens.weight.dtype),
+            offsets=positions[:1].to(torch.int32).repeat(rows),
+            positions=positions,
+        )
+        return logits, cache
+
+    def block_logits(self, block_ids: torch.Tensor, cache: BlockCache) -> torch.Tensor:
+        """Logits ``[rows, width, vocab]`` for the block's current ``[rows, width]`` tokens against ``cache``."""
+
+        self._require_whole_sequence()
+        if block_ids.shape != (cache.rows, cache.width):
+            raise ValueError("block_ids must match the cached block")
+        embeddings = self.embed_tokens(block_ids.reshape(-1))
+        residual = embeddings.unsqueeze(1).expand(-1, self.config.residual_branches, -1)
+        for layer, layer_cache in zip(self.layers, cache.layers, strict=True):
+            residual = layer.forward_block(residual, layer_cache, cache)
+        hidden_states, _ = self.final_residual.read(residual)
+        return self.project_to_vocabulary(hidden_states).view(cache.rows, cache.width, -1)
+
+    # ------------------------------------------------------------ reporting
+    def parameter_report(self) -> dict[str, int]:
+        total = sum(parameter.numel() for parameter in self.parameters())
+        embeddings = self.embed_tokens.weight.numel()
+        return {
+            "deployable_core": total,
+            "training_total": total,
+            "embedding": embeddings,
+            "mtp_auxiliary": 0,
+            "non_embedding_core": total - embeddings,
+            "ngram_table": 0,
+        }
+
+    def save_checkpoint(self, output_dir: Path) -> None:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        self.config.save_json(output_dir / "config.json")
+        save_model(self, output_dir / "model.safetensors")
+
+    def extra_repr(self) -> str:
+        return f"core_parameters={self.parameter_report()['deployable_core']:,}"
+
+    def get_architecture_metadata(self) -> dict[str, Any]:
+        if self.block_size:
+            objective = [
+                (
+                    f"block diffusion over {self.block_size}-token blocks: a causal clean stream with a next-token "
+                    "loss, and a noisy stream denoising each block against the clean prefix (exact block cache)"
+                ),
+                "token editing: corrupted visible tokens are trained back to the originals",
+                "3:1 Gated DeltaNet-2 + attention hybrid, bidirectional inside the active block",
+                (
+                    "GDN-2 noisy reads: the clean forward state where the block starts, plus forward and "
+                    "backward scans inside the block (exclusive backward read, Hydra)"
+                ),
+            ]
+        else:
+            objective = [
+                "masked diffusion (MDLM/LLaDA, SUBS: mask never predicted, low-discrepancy mask ratios, no time input)",
+                "3:1 bidirectional Gated DeltaNet-2 + bidirectional attention hybrid",
+                (
+                    "bidirectional GDN-2: shared projections, per-direction short convs and low-rank decay, "
+                    "exclusive backward read (Hydra)"
+                ),
+            ]
+        return {
+            "architecture": self.config.architecture,
+            "layer_types": list(self.config.layer_types),
+            "effective_num_layers": self.config.effective_num_layers,
+            "shared_layer_indices": [],
+            "parameter_report": self.parameter_report(),
+            "features": [
+                *objective,
+                "4-branch Gated Residual stream (gated read, per-branch write scales)",
+                "full-rank sigmoid output gates",
+                "partial RoPE, per-head QK normalization",
+                "SiTU-GLU",
+                "tied input-output embeddings",
+            ],
+        }

tiny_gdn/fugue_sampling.py

@@ -0,0 +1,654 @@
+"""Generation for Fugue masked-diffusion models.
+
+The continuation starts as ``<|mask|>`` tokens after the prompt. Each model
+call scores masked positions and reveals some of them, one block at a time:
+the whole continuation (``block_length=None``) or LLaDA's semi-autoregressive
+blocks, left to right. A call reveals, in one of three ways:
+
+- ``random`` order with ``steps_per_block`` steps: an even share of the block
+  at uniformly random positions, the reverse of the i.i.d. masking Fugue
+  trains on (MDLM's sampler). Fewer steps reveal more tokens per call, each
+  sampled without seeing the others.
+- ``random`` order with ``draft_tokens``: exact self-speculative decoding. It
+  samples what random order revealing one token per call would, but each call
+  checks up to ``draft_tokens`` tokens that the previous call drafted, with
+  Leviathan et al.'s (2023) rejection rule along the random order. Draft i is
+  checked against row i of one batched forward, which holds drafts 0..i-1.
+- ``confidence`` order: each token whose sampled probability clears
+  ``threshold`` and at least the single most confident one (Fast-dLLM); with
+  ``threshold=None``, the most confident even share per step.
+
+Confidence order lets an undertrained model copy its likeliest phrase into
+loops (docs/architecture.md, "Fugue"); random order keeps it on the training
+distribution.
+
+With ``block_cache``, a call runs the whole canvas and caches everything
+outside the block (Fast-dLLM's DualCache, see
+``FugueForMaskedDiffusion.refresh_block``); the next calls run only the block
+until a row has revealed ``refresh_tokens`` more of it. Text outside the block
+keeps its features from the last full pass in between. Rondo's samples matched
+exact decoding within noise when refreshed every 8 tokens; refreshed only at
+block starts, their perplexity under Monody rose by a fifth.
+
+With ``cuda_graphs``, every call shape is captured once as a CUDA graph per
+model and replayed, since a launch-bound eager forward costs several times
+its GPU time.
+"""
+
+from __future__ import annotations
+
+import math
+import weakref
+from collections import OrderedDict
+from collections.abc import Callable, Hashable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+from tiny_gdn.fugue import BlockCache, FugueForMaskedDiffusion
+
+DEFAULT_BLOCK_LENGTH = 32
+DEFAULT_CONFIDENCE_THRESHOLD = 0.9
+DEFAULT_REFRESH_TOKENS = 8
+UNMASKING_ORDERS: tuple[str, ...] = ("random", "confidence")
+# Eager calls before capture, so Triton has compiled and tuned every kernel.
+GRAPH_WARMUP_STEPS = 2
+# Larger calls are bound by GPU work rather than launches and run eagerly.
+GRAPH_TOKEN_LIMIT = 32_768
+GRAPH_CACHE_SIZE = 4
+
+
+@dataclass(frozen=True)
+class DiffusionSampling:
+    new_tokens: int
+    # None decodes the whole continuation as one block.
+    block_length: int | None = DEFAULT_BLOCK_LENGTH
+    threshold: float | None = DEFAULT_CONFIDENCE_THRESHOLD
+    steps_per_block: int | None = None
+    order: str = "confidence"
+    temperature: float = 0.0
+    top_k: int = 0
+    top_p: float = 1.0
+    stop_at_eos: bool = True
+    block_cache: bool = False
+    refresh_tokens: int = DEFAULT_REFRESH_TOKENS
+    draft_tokens: int = 0
+    cuda_graphs: bool = False
+
+    def __post_init__(self) -> None:
+        if self.new_tokens <= 0:
+            raise ValueError("new_tokens must be positive")
+        if self.block_length is not None and self.block_length <= 0:
+            raise ValueError("block_length must be positive")
+        if self.order not in UNMASKING_ORDERS:
+            raise ValueError(f"order must be one of: {', '.join(UNMASKING_ORDERS)}")
+        if self.threshold is not None and not 0.0 < self.threshold <= 1.0:
+            raise ValueError("threshold must be in (0, 1]")
+        if self.threshold is not None and self.order != "confidence":
+            raise ValueError("a confidence threshold needs confidence order")
+        if self.steps_per_block is not None and self.steps_per_block <= 0:
+            raise ValueError("steps_per_block must be positive")
+        if self.draft_tokens < 0:
+            raise ValueError("draft_tokens cannot be negative")
+        if self.draft_tokens and (self.order != "random" or self.block_cache):
+            raise ValueError("drafts need random order without a block cache")
+        if self.threshold is None and self.steps_per_block is None and not self.draft_tokens:
+            raise ValueError("set a confidence threshold, steps_per_block or draft_tokens")
+        if self.temperature < 0.0:
+            raise ValueError("temperature cannot be negative")
+        if self.top_k < 0 or not 0.0 < self.top_p <= 1.0:
+            raise ValueError("top_k must be >= 0 and top_p in (0, 1]")
+        if self.refresh_tokens <= 0:
+            raise ValueError("refresh_tokens must be positive")
+
+    @property
+    def block_width(self) -> int:
+        return self.block_length or self.new_tokens
+
+
+@dataclass(frozen=True)
+class DenoisingStep:
+    """State after one model call: ``tokens`` is the whole canvas, masks included."""
+
+    step: int
+    block: int
+    tokens: torch.Tensor
+    revealed: int
+    remaining: int
+
+
+# ------------------------------------------------------------------ sampling
+
+
+def _filter_logits(logits: torch.Tensor, *, top_k: int, top_p: float) -> torch.Tensor:
+    if top_k > 0 and top_k < logits.shape[-1]:
+        kth = torch.topk(logits, top_k, dim=-1).values[..., -1:]
+        logits = logits.masked_fill(logits < kth, float("-inf"))
+    if top_p < 1.0:
+        sorted_logits, order = torch.sort(logits, dim=-1, descending=True)
+        cumulative = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+        drop = cumulative - sorted_logits.softmax(dim=-1) >= top_p
+        sorted_logits = sorted_logits.masked_fill(drop, float("-inf"))
+        logits = torch.full_like(logits, float("-inf")).scatter(-1, order, sorted_logits)
+    return logits
+
+
+def _distribution(logits: torch.Tensor, sampling: DiffusionSampling) -> torch.Tensor:
+    """What a token is sampled from, in float64 over the last dimension; one-hot at temperature 0.
+
+    Float32 truncates the tail and lowers the effective temperature of
+    masked-diffusion samplers (Zheng et al., 2024).
+    """
+
+    if sampling.temperature == 0.0:
+        return F.one_hot(logits.argmax(dim=-1), logits.shape[-1]).double()
+    filtered = _filter_logits(logits.double() / sampling.temperature, top_k=sampling.top_k, top_p=sampling.top_p)
+    return filtered.softmax(dim=-1)
+
+
+def _propose(
+    logits: torch.Tensor,
+    sampling: DiffusionSampling,
+    generator: torch.Generator | None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """``(tokens, confidence)`` per row of ``[N, vocab]`` logits."""
+
+    if sampling.temperature == 0.0:
+        confidence, tokens = logits.float().softmax(dim=-1).max(dim=-1)
+        return tokens, confidence
+    probabilities = _distribution(logits, sampling)
+    tokens = torch.multinomial(probabilities, 1, generator=generator).squeeze(-1)
+    confidence = probabilities.gather(-1, tokens[:, None]).squeeze(-1).float()
+    return tokens, confidence
+
+
+def _top_ranked(priority: torch.Tensor, masked: torch.Tensor, steps_left: int) -> torch.Tensor:
+    """An even share of each row's masked positions, highest ``priority`` first."""
+
+    take = torch.ceil(masked.sum(dim=1).float() / steps_left).long()
+    ranks = priority.masked_fill(~masked, -1.0).argsort(dim=1, descending=True).argsort(dim=1)
+    return masked & (ranks < take[:, None])
+
+
+def _reveal_confident(
+    sampling: DiffusionSampling, masked: torch.Tensor, confidence: torch.Tensor, steps_left: int
+) -> torch.Tensor:
+    """Confidence order over ``[rows, width]``; ``confidence`` is -1 off the masks."""
+
+    if sampling.threshold is None:
+        return _top_ranked(confidence, masked, steps_left)
+    reveal = masked & (confidence >= sampling.threshold)
+    best = confidence.argmax(dim=1)
+    rows = torch.arange(masked.shape[0], device=masked.device)
+    has_masks = masked.any(dim=1)
+    reveal[rows[has_masks], best[has_masks]] = True
+    return reveal
+
+
+def _random_order(
+    batch: int, prompt_length: int, sampling: DiffusionSampling, generator: torch.Generator | None, device: torch.device
+) -> torch.Tensor:
+    """Each row's continuation positions ``[batch, new_tokens]`` in a uniformly random order within each block."""
+
+    blocks = []
+    for start in range(0, sampling.new_tokens, sampling.block_width):
+        width = min(sampling.block_width, sampling.new_tokens - start)
+        draws = torch.rand(batch, width, generator=generator, device=device)
+        blocks.append(prompt_length + start + draws.argsort(dim=1))
+    return torch.cat(blocks, dim=1)
+
+
+def _first_eos(tokens: torch.Tensor, eos_token_id: int) -> torch.Tensor:
+    """Each row's first EOS index in ``[rows, T]``, or ``T`` for rows without one."""
+
+    length = tokens.shape[1]
+    positions = torch.arange(length, device=tokens.device)
+    return torch.where(tokens.eq(eos_token_id), positions, length).min(dim=1).values
+
+
+def _committed_eos(tokens: torch.Tensor, mask_token_id: int, eos_token_id: int) -> torch.Tensor:
+    """Rows of ``[rows, T]`` whose revealed prefix reaches an EOS."""
+
+    length = tokens.shape[1]
+    positions = torch.arange(length, device=tokens.device)
+    first_mask = torch.where(tokens.eq(mask_token_id), positions, length).min(dim=1).values
+    return _first_eos(tokens, eos_token_id) < first_mask
+
+
+def _cut_order(
+    order: torch.Tensor, start: torch.Tensor | int, stop: torch.Tensor | int, cutoff: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Move each row's pending positions ``order[start:stop]`` past ``cutoff`` behind the rest.
+
+    Returns the order, the rest keeping its relative order, and how many positions each row moved.
+    """
+
+    index = torch.arange(order.shape[1], device=order.device)
+    start = torch.as_tensor(start, device=order.device).reshape(-1, 1)
+    stop = torch.as_tensor(stop, device=order.device).reshape(-1, 1)
+    pending = (index >= start) & (index < stop)
+    dead = pending & (order > cutoff[:, None])
+    moved = dead.sum(dim=1)
+    if not bool(moved.any()):
+        return order, moved
+    key = torch.where(index < start, 0, torch.where(pending, 1 + dead.long(), 3))
+    return order.gather(1, torch.sort(key, dim=1, stable=True).indices), moved
+
+
+# -------------------------------------------------------------------- graphs
+
+
+class _Captured:
+    """A function of tensors, captured once as a CUDA graph over static copies of its inputs."""
+
+    def __init__(self, function: Callable[..., Any], inputs: tuple[torch.Tensor, ...], pool: Any) -> None:
+        self.inputs = tuple(tensor.clone() for tensor in inputs)
+        device = self.inputs[0].device
+        current = torch.cuda.current_stream(device)
+        stream = torch.cuda.Stream(device=device)
+        stream.wait_stream(current)
+        with torch.cuda.stream(stream):
+            for _ in range(GRAPH_WARMUP_STEPS):
+                function(*self.inputs)
+        current.wait_stream(stream)
+        self.graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.graph, pool=pool, capture_error_mode="thread_local"):
+            self.outputs = function(*self.inputs)
+
+    def __call__(self, *inputs: torch.Tensor) -> Any:
+        for static, tensor in zip(self.inputs, inputs, strict=True):
+            static.copy_(tensor)
+        self.graph.replay()
+        return self.outputs
+
+
+class _GraphCache:
+    """One model's captured decoding graphs by call shape, least recently used dropped first."""
+
+    def __init__(self) -> None:
+        self.graphs: OrderedDict[Hashable, Any] = OrderedDict()
+
+    def get(self, key: Hashable, build: Callable[[], Any]) -> Any:
+        graph = self.graphs.pop(key, None)
+        if graph is None:
+            while len(self.graphs) >= GRAPH_CACHE_SIZE:
+                self.graphs.popitem(last=False)
+            graph = build()
+        self.graphs[key] = graph
+        return graph
+
+
+# Graphs bake in the parameters' addresses, so they live and die with the model object.
+_GRAPH_CACHES: weakref.WeakKeyDictionary[FugueForMaskedDiffusion, _GraphCache] = weakref.WeakKeyDictionary()
+
+
+def _graphs(model: FugueForMaskedDiffusion) -> _GraphCache:
+    cache = _GRAPH_CACHES.get(model)
+    if cache is None:
+        cache = _GRAPH_CACHES[model] = _GraphCache()
+    return cache
+
+
+def _graphable(sampling: DiffusionSampling, canvas: torch.Tensor) -> bool:
+    return sampling.cuda_graphs and canvas.is_cuda and canvas.numel() <= GRAPH_TOKEN_LIMIT
+
+
+def _canvas_logits(
+    model: FugueForMaskedDiffusion, canvas: torch.Tensor, positions: torch.Tensor, sampling: DiffusionSampling
+) -> torch.Tensor:
+    """``model.canvas_logits``, replayed from a graph when ``sampling`` allows it."""
+
+    if not _graphable(sampling, canvas):
+        return model.canvas_logits(canvas, positions)
+    key = ("canvas", *canvas.shape, positions.shape[1])
+    graph = _graphs(model).get(key, lambda: _Captured(model.canvas_logits, (canvas, positions), None))
+    return graph(canvas, positions)
+
+
+class _BlockGraphs:
+    """``refresh_block`` and ``block_logits`` for one canvas and block shape: two graphs over one cache."""
+
+    def __init__(self, model: FugueForMaskedDiffusion, canvas: torch.Tensor, positions: torch.Tensor) -> None:
+        pool = torch.cuda.graph_pool_handle()
+        self.refresh = _Captured(model.refresh_block, (canvas, positions), pool)
+        # Fill the cache before the block step warms up on it.
+        self.refresh(canvas, positions)
+        cache: BlockCache = self.refresh.outputs[1]
+        self.cache = cache
+        block = canvas.index_select(1, positions)
+        self.block = _Captured(lambda ids: model.block_logits(ids, cache), (block,), pool)
+
+
+# ------------------------------------------------------------------- scorers
+
+
+class _CanvasScorer:
+    """Logits from a forward of the whole canvas at every call."""
+
+    def __init__(self, model: FugueForMaskedDiffusion, sampling: DiffusionSampling) -> None:
+        self.model = model
+        self.sampling = sampling
+        self.start = 0
+        self.block_mask: torch.Tensor | None = None
+
+    def begin(self, canvas: torch.Tensor, start: int, stop: int) -> None:
+        self.start = start
+        self.block_mask = None
+        if not _graphable(self.sampling, canvas):
+            self.block_mask = torch.zeros_like(canvas, dtype=torch.bool)
+            self.block_mask[:, start:stop] = True
+
+    def __call__(self, canvas: torch.Tensor, masked: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        """Logits ``[rows, picks, vocab]`` at ``positions`` ``[rows, picks]`` of the block."""
+
+        if self.block_mask is None:
+            return _canvas_logits(self.model, canvas, positions, self.sampling)
+        # Without graphs, the packed forward scores the block, as evaluations always have.
+        logits = self.model(canvas, return_logits=True, logits_mask=self.block_mask).logits
+        logits = logits.view(canvas.shape[0], -1, logits.shape[-1])
+        return logits.gather(1, (positions - self.start).unsqueeze(-1).expand(-1, -1, logits.shape[-1]))
+
+
+class _BlockCacheScorer:
+    """Logits from a full forward, then from the block alone against its cache until a refresh is due."""
+
+    def __init__(self, model: FugueForMaskedDiffusion, sampling: DiffusionSampling) -> None:
+        self.model = model
+        self.sampling = sampling
+        self.start = 0
+        self.positions: torch.Tensor | None = None
+        self.cache: BlockCache | None = None
+        self.masked_at_refresh: torch.Tensor | None = None
+        self.graphs: _BlockGraphs | None = None
+
+    def begin(self, canvas: torch.Tensor, start: int, stop: int) -> None:
+        self.start = start
+        self.positions = torch.arange(start, stop, device=canvas.device)
+        self.cache = None
+        self.graphs = None
+        if _graphable(self.sampling, canvas):
+            key = ("block", *canvas.shape, stop - start)
+            positions = self.positions
+            self.graphs = _graphs(self.model).get(key, lambda: _BlockGraphs(self.model, canvas, positions))
+
+    def __call__(self, canvas: torch.Tensor, masked: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        """Logits ``[rows, picks, vocab]`` at ``positions`` ``[rows, picks]`` of the block."""
+
+        if self.positions is None:
+            raise RuntimeError("begin a block before scoring it")
+        masks_left = masked.sum(dim=1)
+        if (
+            self.cache is not None
+            and self.masked_at_refresh is not None
+            and int((self.masked_at_refresh - masks_left).max()) >= self.sampling.refresh_tokens
+        ):
+            self.cache = None
+        if self.cache is None:
+            if self.graphs is None:
+                logits, self.cache = self.model.refresh_block(canvas, self.positions)
+            else:
+                logits, self.cache = self.graphs.refresh(canvas, self.positions)
+            self.masked_at_refresh = masks_left
+        else:
+            block = canvas.index_select(1, self.positions)
+            if self.graphs is None:
+                logits = self.model.block_logits(block, self.cache)
+            else:
+                logits = self.graphs.block(block)
+        offsets = (positions - self.start).unsqueeze(-1).expand(-1, -1, logits.shape[-1])
+        return logits.gather(1, offsets)
+
+
+# ------------------------------------------------------------------ decoders
+
+
+@torch.no_grad()
+def diffusion_generate(
+    model: FugueForMaskedDiffusion,
+    prompt_ids: torch.Tensor,
+    sampling: DiffusionSampling,
+    *,
+    generator: torch.Generator | None = None,
+    eos_token_id: int | None = None,
+    on_step: Callable[[DenoisingStep], bool | None] | None = None,
+) -> torch.Tensor:
+    """Fill ``new_tokens`` masks after each ``[B, P]`` prompt row; returns ``[B, new_tokens]``.
+
+    ``on_step`` sees every model call and may return True to cancel. With
+    ``stop_at_eos``, positions after a row's first revealed EOS leave its
+    reveal order and keep their masks; the row stops once everything before
+    that EOS is revealed, and the run stops when every row has.
+    """
+
+    if prompt_ids.ndim != 2:
+        raise ValueError("prompt_ids must have shape [batch, prompt]")
+    if prompt_ids.shape[1] == 0:
+        raise ValueError("the prompt needs at least one token")
+    mask_id = model.mask_token_id
+    batch, prompt_length = prompt_ids.shape
+    total = prompt_length + sampling.new_tokens
+    canvas = torch.cat(
+        [prompt_ids, prompt_ids.new_full((batch, sampling.new_tokens), mask_id)],
+        dim=1,
+    )
+    if sampling.draft_tokens:
+        return _speculative_generate(model, canvas, prompt_length, sampling, generator, eos_token_id, on_step)
+    stop_at_eos = sampling.stop_at_eos and eos_token_id is not None
+    finished = torch.zeros(batch, dtype=torch.bool, device=canvas.device)
+    cutoff = torch.full((batch,), total, device=canvas.device)
+    scorer = _BlockCacheScorer(model, sampling) if sampling.block_cache else _CanvasScorer(model, sampling)
+    order = _random_order(batch, prompt_length, sampling, generator, canvas.device) if sampling.order == "random" else None
+    step = 0
+    for block, start in enumerate(range(prompt_length, total, sampling.block_width)):
+        stop = min(start + sampling.block_width, total)
+        width = stop - start
+        block_steps = min(sampling.steps_per_block or width, width)
+        scorer.begin(canvas, start, stop)
+        everywhere = torch.arange(start, stop, device=canvas.device).expand(batch, -1)
+        # Random order picks the same number of positions every call, so one graph fits the block.
+        share = math.ceil(width / block_steps)
+        if order is not None:
+            block_order = order[:, start - prompt_length : stop - prompt_length]
+            # Slots past the block's end repeat positions revealed by the first call.
+            block_order = torch.cat([block_order, block_order[:, :share]], dim=1)
+        cursor = 0
+        for block_step in range(block_steps if sampling.threshold is None else width):
+            window = canvas[:, start:stop]
+            masked = window.eq(mask_id) & ~finished[:, None] & (everywhere < cutoff[:, None])
+            remaining = int(masked.sum())
+            if remaining == 0:
+                break
+            steps_left = block_steps - block_step
+            if order is not None:
+                take = math.ceil((width - cursor) / steps_left)
+                picks = block_order[:, cursor : cursor + share]
+                cursor += take
+                live = masked.gather(1, picks - start) & (torch.arange(share, device=canvas.device) < take)
+                logits = scorer(canvas, masked, picks)
+                tokens, _ = _propose(logits.flatten(0, 1), sampling, generator)
+                current = canvas.gather(1, picks)
+                canvas.scatter_(1, picks, torch.where(live, tokens.view(batch, share), current))
+                revealed = int(live.sum())
+            else:
+                logits = scorer(canvas, masked, everywhere)
+                proposals, confidence = _propose(logits.flatten(0, 1), sampling, generator)
+                proposals = proposals.view(batch, width)
+                confidence = confidence.view(batch, width).masked_fill(~masked, -1.0)
+                reveal = _reveal_confident(sampling, masked, confidence, steps_left)
+                canvas[:, start:stop] = torch.where(reveal, proposals, window)
+                revealed = int(reveal.sum())
+            step += 1
+            if stop_at_eos:
+                cutoff = prompt_length + _first_eos(canvas[:, prompt_length:], eos_token_id)
+                if order is not None:
+                    block_order, _ = _cut_order(block_order, cursor, width, cutoff)
+                finished |= _committed_eos(canvas[:, prompt_length:], mask_id, eos_token_id)
+            if on_step is not None:
+                cancelled = on_step(
+                    DenoisingStep(
+                        step=step,
+                        block=block,
+                        tokens=canvas[:, prompt_length:],
+                        revealed=revealed,
+                        remaining=remaining - revealed,
+                    )
+                )
+                if cancelled:
+                    return canvas[:, prompt_length:]
+            if stop_at_eos and bool(finished.all()):
+                return canvas[:, prompt_length:]
+        if stop_at_eos:
+            finished |= canvas[:, start:stop].eq(eos_token_id).any(dim=1)
+            if bool(finished.all()):
+                break
+    return canvas[:, prompt_length:]
+
+
+def _speculative_generate(
+    model: FugueForMaskedDiffusion,
+    canvas: torch.Tensor,
+    prompt_length: int,
+    sampling: DiffusionSampling,
+    generator: torch.Generator | None,
+    eos_token_id: int | None,
+    on_step: Callable[[DenoisingStep], bool | None] | None,
+) -> torch.Tensor:
+    """Random order, one token at a time in distribution, verifying up to ``draft_tokens`` drafts per call.
+
+    Each call scores ``drafts + 1`` copies of the canvas: copy i holds drafts
+    0..i-1, so its prediction at draft i's position is the conditional that
+    draft i is checked against, and its predictions at the next ``drafts``
+    positions become the next drafts if it ends up the last accepted copy.
+    """
+
+    batch, total = canvas.shape
+    device = canvas.device
+    mask_id = model.mask_token_id
+    drafts_per_call = sampling.draft_tokens
+    copies = drafts_per_call + 1
+    stop_at_eos = sampling.stop_at_eos and eos_token_id is not None
+    # Window slots past the end point at position 0, a prompt token that is never revealed.
+    order = F.pad(_random_order(batch, prompt_length, sampling, generator, device), (0, 2 * drafts_per_call + 1))
+    slots = torch.arange(2 * drafts_per_call + 1, device=device)
+    draft_slots = torch.arange(drafts_per_call, device=device)
+    # Copy i holds drafts j < i.
+    holds = torch.arange(copies, device=device)[:, None] > draft_slots[None, :]
+    rows = torch.arange(batch, device=device)
+    cursor = torch.zeros(batch, dtype=torch.long, device=device)
+    # order[:cursor] is revealed, order[cursor:limit] is masked and still to come.
+    limit = torch.full((batch,), sampling.new_tokens, dtype=torch.long, device=device)
+    drafts = torch.zeros(batch, drafts_per_call, dtype=torch.long, device=device)
+    draft_count = torch.zeros(batch, dtype=torch.long, device=device)
+    draft_probabilities: torch.Tensor | None = None
+    finished = torch.zeros(batch, dtype=torch.bool, device=device)
+    step = 0
+    while True:
+        remaining = (limit - cursor).masked_fill(finished, 0)
+        left = int(remaining.sum())
+        if left == 0:
+            break
+        draft_count = torch.minimum(draft_count, remaining)
+        window = order.gather(1, cursor[:, None] + slots)
+        live = draft_slots[None, :] < draft_count[:, None]
+        variants = canvas[:, None, :].repeat(1, copies, 1)
+        if drafts_per_call:
+            targets = window[:, None, :drafts_per_call].expand(-1, copies, -1)
+            placed = torch.where(
+                holds[None] & live[:, None, :], drafts[:, None, :].expand(-1, copies, -1), variants.gather(2, targets)
+            )
+            variants.scatter_(2, targets, placed)
+        picks = window.unfold(1, copies, 1)
+        logits = _canvas_logits(
+            model, variants.view(batch * copies, total), picks.reshape(batch * copies, copies), sampling
+        ).view(batch, copies, copies, -1)
+        verified = _distribution(logits[:, :, 0], sampling)
+        accepted = torch.zeros(batch, dtype=torch.long, device=device)
+        ends_at_eos = torch.zeros(batch, dtype=torch.bool, device=device)
+        if draft_probabilities is not None:
+            target = verified[:, :drafts_per_call].gather(2, drafts[..., None]).squeeze(-1)
+            proposal = draft_probabilities.gather(2, drafts[..., None]).squeeze(-1)
+            coins = torch.rand(batch, drafts_per_call, generator=generator, device=device, dtype=torch.float64)
+            accepted = ((coins * proposal < target) & live).long().cumprod(dim=1).sum(dim=1)
+            if stop_at_eos:
+                # An accepted EOS cuts the order the later drafts were drawn for, so the run ends with it.
+                accepted_eos = drafts.eq(eos_token_id) & (draft_slots[None, :] < accepted[:, None])
+                first_eos = torch.where(accepted_eos, draft_slots[None, :], drafts_per_call).min(dim=1).values
+                ends_at_eos = first_eos < accepted
+                accepted = torch.minimum(accepted, first_eos + 1)
+        following = verified[rows, accepted]
+        if draft_probabilities is not None:
+            # A rejected draft's position samples from what its target has beyond the draft.
+            rejected = accepted < draft_count
+            excess = (following - draft_probabilities[rows, accepted.clamp(max=drafts_per_call - 1)]).clamp(min=0.0)
+            mass = excess.sum(dim=-1, keepdim=True)
+            excess = torch.where(mass > 0.0, excess / mass.clamp(min=1e-300), following)
+            following = torch.where(rejected[:, None], excess, following)
+        next_token = torch.multinomial(following, 1, generator=generator).squeeze(-1)
+        has_next = (accepted < remaining) & ~ends_at_eos
+        if drafts_per_call:
+            kept = draft_slots[None, :] < accepted[:, None]
+            draft_positions = window[:, :drafts_per_call]
+            canvas.scatter_(1, draft_positions, torch.where(kept, drafts, canvas.gather(1, draft_positions)))
+        next_position = window.gather(1, accepted[:, None])
+        canvas.scatter_(1, next_position, torch.where(has_next[:, None], next_token[:, None], canvas.gather(1, next_position)))
+        revealed = accepted + has_next.long()
+        cursor += revealed
+        if drafts_per_call:
+            draft_probabilities = _distribution(logits[rows, accepted, 1:], sampling)
+            drafts = torch.multinomial(draft_probabilities.flatten(0, 1), 1, generator=generator).view(batch, -1)
+            draft_count = (limit - cursor).clamp(min=0, max=drafts_per_call)
+        if stop_at_eos:
+            cutoff = prompt_length + _first_eos(canvas[:, prompt_length:], eos_token_id)
+            order, moved = _cut_order(order, cursor, limit, cutoff)
+            limit -= moved
+            finished |= _committed_eos(canvas[:, prompt_length:], mask_id, eos_token_id)
+            # Drafts after an EOS or a cut were drawn for positions that no longer come next.
+            draft_count = draft_count.masked_fill(finished | ends_at_eos | moved.gt(0), 0)
+        step += 1
+        if on_step is not None:
+            count = int(revealed.sum())
+            cancelled = on_step(
+                DenoisingStep(
+                    step=step,
+                    block=min(int(cursor.min()), sampling.new_tokens - 1) // sampling.block_width,
+                    tokens=canvas[:, prompt_length:],
+                    revealed=count,
+                    remaining=left - count,
+                )
+            )
+            if cancelled:
+                break
+    return canvas[:, prompt_length:]
+
+
+@torch.no_grad()
+def denoise_canvas(model: FugueForMaskedDiffusion, canvas: torch.Tensor, *, steps: int) -> torch.Tensor:
+    """Greedily fill every ``<|mask|>`` of a ``[B, T]`` canvas in ``steps`` calls, most confident first."""
+
+    mask_id = model.mask_token_id
+    canvas = canvas.clone()
+    for step in range(steps):
+        masked = canvas.eq(mask_id)
+        remaining = masked.sum(dim=1)
+        if int(remaining.max()) == 0:
+            break
+        logits = model(canvas, return_logits=True, logits_mask=masked).logits
+        confidence, proposals = logits.float().softmax(dim=-1).max(dim=-1)
+        scores = torch.full(canvas.shape, -1.0, device=canvas.device)
+        scores[masked] = confidence
+        filled = torch.zeros_like(canvas)
+        filled[masked] = proposals
+        take = torch.ceil(remaining.float() / (steps - step)).long()
+        ranks = scores.argsort(dim=1, descending=True).argsort(dim=1)
+        reveal = masked & (ranks < take[:, None])
+        canvas = torch.where(reveal, filled, canvas)
+    return canvas
+
+
+def committed_prefix(tokens: torch.Tensor, mask_token_id: int) -> torch.Tensor:
+    """The revealed tokens of a ``[N]`` canvas up to its first remaining mask."""
+
+    masked = tokens.eq(mask_token_id).nonzero()
+    return tokens if masked.numel() == 0 else tokens[: int(masked[0])]

tiny_gdn/fused_ops.py

@@ -0,0 +1,548 @@
+"""Fused CUDA kernels for memory-bound blocks: SiTU-GLU, the LM loss, token log-probs."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+from triton.language.extra import libdevice
+
+from tiny_gdn.nn_common import (
+    linear_cross_entropy_reference,
+    linear_token_logprobs_reference,
+    situ_glu_reference,
+)
+
+BLOCK_SIZE = 1024
+VOCAB_BLOCK_SIZE = 4096
+LOGITS_CHUNK_BYTES = 1 << 30
+
+
+@triton.jit
+def _situ_glu_forward_kernel(
+    gate_up,
+    output,
+    total,
+    width,
+    gate_cap,
+    up_cap,
+    BLOCK: tl.constexpr,
+):
+    offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+    mask = offsets < total
+    row = offsets // width
+    gate_offsets = offsets + row * width
+    gate = tl.load(gate_up + gate_offsets, mask=mask, other=0.0).to(tl.float32)
+    up = tl.load(gate_up + gate_offsets + width, mask=mask, other=0.0).to(tl.float32)
+    gated = gate_cap * libdevice.tanh(gate / gate_cap) * tl.sigmoid(gate)
+    result = gated * up_cap * libdevice.tanh(up / up_cap)
+    tl.store(output + offsets, result.to(output.dtype.element_ty), mask=mask)
+
+
+@triton.jit
+def _situ_glu_backward_kernel(
+    gate_up,
+    grad_output,
+    grad_gate_up,
+    total,
+    width,
+    gate_cap,
+    up_cap,
+    BLOCK: tl.constexpr,
+):
+    offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+    mask = offsets < total
+    row = offsets // width
+    gate_offsets = offsets + row * width
+    gate = tl.load(gate_up + gate_offsets, mask=mask, other=0.0).to(tl.float32)
+    up = tl.load(gate_up + gate_offsets + width, mask=mask, other=0.0).to(tl.float32)
+    grad = tl.load(grad_output + offsets, mask=mask, other=0.0).to(tl.float32)
+    gate_tanh = libdevice.tanh(gate / gate_cap)
+    up_tanh = libdevice.tanh(up / up_cap)
+    sigmoid = tl.sigmoid(gate)
+    gated = gate_cap * gate_tanh * sigmoid
+    capped_up = up_cap * up_tanh
+    gate_slope = sigmoid * ((1.0 - gate_tanh * gate_tanh) + gate_cap * gate_tanh * (1.0 - sigmoid))
+    grad_gate = grad * capped_up * gate_slope
+    grad_up = grad * gated * (1.0 - up_tanh * up_tanh)
+    dtype = grad_gate_up.dtype.element_ty
+    tl.store(grad_gate_up + gate_offsets, grad_gate.to(dtype), mask=mask)
+    tl.store(grad_gate_up + gate_offsets + width, grad_up.to(dtype), mask=mask)
+
+
+def _launch_grid(total: int) -> tuple[int]:
+    return (triton.cdiv(total, BLOCK_SIZE),)
+
+
+class _SiTUGLUFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, gate_up: torch.Tensor, gate_cap: float, up_cap: float) -> torch.Tensor:
+        gate_up = gate_up.contiguous()
+        width = gate_up.shape[-1] // 2
+        output = gate_up.new_empty((*gate_up.shape[:-1], width))
+        total = output.numel()
+        if total:
+            _situ_glu_forward_kernel[_launch_grid(total)](
+                gate_up, output, total, width, gate_cap, up_cap, BLOCK=BLOCK_SIZE
+            )
+        ctx.save_for_backward(gate_up)
+        ctx.caps = (gate_cap, up_cap)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]:
+        (gate_up,) = ctx.saved_tensors
+        gate_cap, up_cap = ctx.caps
+        grad_output = grad_output.contiguous()
+        grad_gate_up = torch.empty_like(gate_up)
+        width = gate_up.shape[-1] // 2
+        total = grad_output.numel()
+        if total:
+            _situ_glu_backward_kernel[_launch_grid(total)](
+                gate_up,
+                grad_output,
+                grad_gate_up,
+                total,
+                width,
+                gate_cap,
+                up_cap,
+                BLOCK=BLOCK_SIZE,
+            )
+        return grad_gate_up, None, None
+
+
+def situ_glu(gate_up: torch.Tensor, gate_cap: float, up_cap: float) -> torch.Tensor:
+    """``softcap(gate) * sigmoid(gate) * softcap(up)`` over ``[..., gate | up]``.
+
+    The CUDA path evaluates in FP32 and rounds once; backward recomputes the
+    activations from ``gate_up`` instead of saving intermediates.
+    """
+
+    if gate_up.shape[-1] % 2:
+        raise ValueError("gate_up must have an even last dimension")
+    if not gate_up.is_cuda:
+        return situ_glu_reference(gate_up, gate_cap, up_cap)
+    return _SiTUGLUFunction.apply(gate_up, gate_cap, up_cap)
+
+
+@triton.jit
+def _cross_entropy_rows_kernel(
+    logits,
+    targets,
+    losses,
+    log_partitions,
+    valid_count,
+    row_stride,
+    vocab_size,
+    z_loss_coefficient,
+    ignore_index,
+    COMPUTE_GRAD: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """Per row: logsumexp and NLL; optionally overwrite the row with d(mean loss)/d(logits)."""
+
+    row = tl.program_id(0).to(tl.int64)
+    row_logits = logits + row * row_stride
+    target = tl.load(targets + row)
+    columns = tl.arange(0, BLOCK)
+    if target == ignore_index:
+        tl.store(losses + row, 0.0)
+        tl.store(log_partitions + row, 0.0)
+        if COMPUTE_GRAD:
+            for start in range(0, vocab_size, BLOCK):
+                offsets = start + columns
+                zeros = tl.zeros((BLOCK,), dtype=logits.dtype.element_ty)
+                tl.store(row_logits + offsets, zeros, mask=offsets < vocab_size)
+    else:
+        running_max = tl.full((), float("-inf"), tl.float32)
+        running_sum = tl.zeros((), dtype=tl.float32)
+        for start in range(0, vocab_size, BLOCK):
+            offsets = start + columns
+            values = tl.load(
+                row_logits + offsets, mask=offsets < vocab_size, other=float("-inf")
+            ).to(tl.float32)
+            new_max = tl.maximum(running_max, tl.max(values, axis=0))
+            running_sum = running_sum * tl.exp(running_max - new_max) + tl.sum(
+                tl.exp(values - new_max), axis=0
+            )
+            running_max = new_max
+        log_partition = running_max + tl.log(running_sum)
+        target_logit = tl.load(row_logits + target).to(tl.float32)
+        tl.store(losses + row, log_partition - target_logit)
+        tl.store(log_partitions + row, log_partition)
+        if COMPUTE_GRAD:
+            grad_scale = 1.0 / tl.load(valid_count).to(tl.float32)
+            softmax_scale = grad_scale * (1.0 + 2.0 * z_loss_coefficient * log_partition)
+            for start in range(0, vocab_size, BLOCK):
+                offsets = start + columns
+                in_vocab = offsets < vocab_size
+                values = tl.load(row_logits + offsets, mask=in_vocab, other=0.0).to(tl.float32)
+                grad = tl.exp(values - log_partition) * softmax_scale
+                grad = tl.where(offsets == target, grad - grad_scale, grad)
+                tl.store(row_logits + offsets, grad.to(logits.dtype.element_ty), mask=in_vocab)
+
+
+class _LinearCrossEntropyFunction(torch.autograd.Function):
+    """Chunked ``states @ weight.T`` + cross-entropy + z-loss with gradients precomputed.
+
+    Each logits chunk is consumed in place, so peak extra memory is one chunk
+    instead of ``[tokens, vocab]`` logits plus their softmax for backward.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        states: torch.Tensor,
+        weight: torch.Tensor,
+        targets: torch.Tensor,
+        z_loss_coefficient: float,
+        ignore_index: int,
+        compute_grad: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        states = states.contiguous()
+        targets = targets.contiguous()
+        rows, vocab_size = states.shape[0], weight.shape[0]
+        valid = targets.ne(ignore_index)
+        valid_count = valid.sum(dtype=torch.float32)
+        losses = states.new_empty(rows, dtype=torch.float32)
+        log_partitions = states.new_empty(rows, dtype=torch.float32)
+        grad_states = torch.empty_like(states) if compute_grad else None
+        grad_weight = (
+            torch.zeros(weight.shape, dtype=torch.float32, device=weight.device)
+            if compute_grad
+            else None
+        )
+        chunk_rows = max(1, LOGITS_CHUNK_BYTES // (vocab_size * states.element_size()))
+        for start in range(0, rows, chunk_rows):
+            end = min(start + chunk_rows, rows)
+            chunk = states[start:end]
+            logits = chunk @ weight.t()
+            _cross_entropy_rows_kernel[(end - start,)](
+                logits,
+                targets[start:end],
+                losses[start:end],
+                log_partitions[start:end],
+                valid_count,
+                logits.stride(0),
+                vocab_size,
+                z_loss_coefficient,
+                ignore_index,
+                COMPUTE_GRAD=compute_grad,
+                BLOCK=VOCAB_BLOCK_SIZE,
+                num_warps=8,
+            )
+            if compute_grad:
+                torch.mm(logits, weight, out=grad_states[start:end])
+                if chunk.dtype == torch.float32:
+                    grad_weight.addmm_(logits.t(), chunk)
+                else:
+                    torch.addmm(
+                        grad_weight, logits.t(), chunk, out_dtype=torch.float32, out=grad_weight
+                    )
+        cross_entropy = losses.sum() / valid_count
+        z_loss = log_partitions.square().sum() / valid_count
+        objective = cross_entropy + z_loss_coefficient * z_loss
+        if compute_grad:
+            ctx.save_for_backward(grad_states, grad_weight)
+        ctx.weight_dtype = weight.dtype
+        ctx.mark_non_differentiable(cross_entropy, z_loss)
+        return objective, cross_entropy, z_loss
+
+    @staticmethod
+    def backward(
+        ctx, grad_objective: torch.Tensor, *_: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, None, None, None, None]:
+        grad_states, grad_weight = ctx.saved_tensors
+        scale = grad_objective.to(torch.float32)
+        grad_states = (grad_states * scale).to(grad_states.dtype)
+        grad_weight = (grad_weight * scale).to(ctx.weight_dtype)
+        return grad_states, grad_weight, None, None, None, None
+
+
+class _WeightedLinearCrossEntropyFunction(torch.autograd.Function):
+    """``Σ_i w_i (CE_i + z log Z_i²)`` over chunked ``states @ weight.T``, gradients precomputed.
+
+    The row kernel runs with unit scale and each row's gradient is then
+    multiplied by its weight, so per-token weights cost one in-place multiply.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        states: torch.Tensor,
+        weight: torch.Tensor,
+        targets: torch.Tensor,
+        token_weights: torch.Tensor,
+        z_loss_coefficient: float,
+        compute_grad: bool,
+        excluded_token_id: int | None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        states = states.contiguous()
+        targets = targets.contiguous()
+        token_weights = token_weights.float().contiguous()
+        rows, vocab_size = states.shape[0], weight.shape[0]
+        unit = torch.ones((), dtype=torch.float32, device=states.device)
+        losses = states.new_empty(rows, dtype=torch.float32)
+        log_partitions = states.new_empty(rows, dtype=torch.float32)
+        grad_states = torch.empty_like(states) if compute_grad else None
+        grad_weight = (
+            torch.zeros(weight.shape, dtype=torch.float32, device=weight.device)
+            if compute_grad
+            else None
+        )
+        chunk_rows = max(1, LOGITS_CHUNK_BYTES // (vocab_size * states.element_size()))
+        for start in range(0, rows, chunk_rows):
+            end = min(start + chunk_rows, rows)
+            chunk = states[start:end]
+            logits = chunk @ weight.t()
+            if excluded_token_id is not None:
+                logits[:, excluded_token_id] = float("-inf")
+            _cross_entropy_rows_kernel[(end - start,)](
+                logits,
+                targets[start:end],
+                losses[start:end],
+                log_partitions[start:end],
+                unit,
+                logits.stride(0),
+                vocab_size,
+                z_loss_coefficient,
+                -100,
+                COMPUTE_GRAD=compute_grad,
+                BLOCK=VOCAB_BLOCK_SIZE,
+                num_warps=8,
+            )
+            if compute_grad:
+                logits.mul_(token_weights[start:end, None].to(logits.dtype))
+                torch.mm(logits, weight, out=grad_states[start:end])
+                if chunk.dtype == torch.float32:
+                    grad_weight.addmm_(logits.t(), chunk)
+                else:
+                    torch.addmm(
+                        grad_weight, logits.t(), chunk, out_dtype=torch.float32, out=grad_weight
+                    )
+        objective = (token_weights * (losses + z_loss_coefficient * log_partitions.square())).sum()
+        if compute_grad:
+            ctx.save_for_backward(grad_states, grad_weight)
+        ctx.weight_dtype = weight.dtype
+        ctx.mark_non_differentiable(losses, log_partitions)
+        return objective, losses, log_partitions
+
+    @staticmethod
+    def backward(
+        ctx, grad_objective: torch.Tensor, *_: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, None, None, None, None, None]:
+        grad_states, grad_weight = ctx.saved_tensors
+        scale = grad_objective.to(torch.float32)
+        grad_states = (grad_states * scale).to(grad_states.dtype)
+        grad_weight = (grad_weight * scale).to(ctx.weight_dtype)
+        return grad_states, grad_weight, None, None, None, None, None
+
+
+def weighted_linear_cross_entropy(
+    states: torch.Tensor,
+    weight: torch.Tensor,
+    targets: torch.Tensor,
+    token_weights: torch.Tensor,
+    *,
+    z_loss_coefficient: float = 0.0,
+    excluded_token_id: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """``(Σ_i w_i (CE_i + z log Z_i²), CE [rows], log Z [rows])`` without full logits.
+
+    Every row is scored; only the weighted sum carries gradients.
+    ``excluded_token_id`` gets probability zero (its logit is -inf).
+    """
+
+    if states.ndim != 2 or targets.shape != states.shape[:1] or token_weights.shape != targets.shape:
+        raise ValueError("states must be [tokens, hidden], targets and token_weights [tokens]")
+    if not states.is_cuda:
+        logits = F.linear(states, weight).float()
+        if excluded_token_id is not None:
+            logits[:, excluded_token_id] = float("-inf")
+        log_partitions = torch.logsumexp(logits, dim=-1)
+        losses = log_partitions - logits.gather(1, targets[:, None]).squeeze(1)
+        weights = token_weights.float()
+        objective = (weights * (losses + z_loss_coefficient * log_partitions.square())).sum()
+        return objective, losses.detach(), log_partitions.detach()
+    compute_grad = torch.is_grad_enabled() and (states.requires_grad or weight.requires_grad)
+    return _WeightedLinearCrossEntropyFunction.apply(
+        states, weight, targets, token_weights, z_loss_coefficient, compute_grad, excluded_token_id
+    )
+
+
+def linear_cross_entropy(
+    states: torch.Tensor,
+    weight: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    z_loss_coefficient: float = 0.0,
+    ignore_index: int = -100,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Mean ``cross_entropy(states @ weight.T, targets)`` without full logits.
+
+    Returns ``(objective, cross_entropy, z_loss)``. Only ``objective =
+    cross_entropy + z_loss_coefficient * z_loss`` carries gradients; ``z_loss``
+    is the mean squared log-partition over non-ignored targets.
+    """
+
+    if states.ndim != 2 or targets.shape != states.shape[:1]:
+        raise ValueError("states must be [tokens, hidden] and targets [tokens]")
+    if not states.is_cuda:
+        return linear_cross_entropy_reference(
+            states,
+            weight,
+            targets,
+            z_loss_coefficient=z_loss_coefficient,
+            ignore_index=ignore_index,
+        )
+    compute_grad = torch.is_grad_enabled() and (states.requires_grad or weight.requires_grad)
+    return _LinearCrossEntropyFunction.apply(
+        states, weight, targets, z_loss_coefficient, ignore_index, compute_grad
+    )
+
+
+@triton.jit
+def _token_logprob_grad_rows_kernel(
+    logits,
+    targets,
+    log_partitions,
+    row_grads,
+    row_stride,
+    vocab_size,
+    ignore_index,
+    BLOCK: tl.constexpr,
+):
+    """Overwrite each row with ``g * (onehot(target) - softmax)``, the gradient of ``g * log p``."""
+
+    row = tl.program_id(0).to(tl.int64)
+    row_logits = logits + row * row_stride
+    target = tl.load(targets + row)
+    columns = tl.arange(0, BLOCK)
+    if target == ignore_index:
+        for start in range(0, vocab_size, BLOCK):
+            offsets = start + columns
+            zeros = tl.zeros((BLOCK,), dtype=logits.dtype.element_ty)
+            tl.store(row_logits + offsets, zeros, mask=offsets < vocab_size)
+    else:
+        scale = tl.load(row_grads + row).to(tl.float32)
+        log_partition = tl.load(log_partitions + row)
+        for start in range(0, vocab_size, BLOCK):
+            offsets = start + columns
+            in_vocab = offsets < vocab_size
+            values = tl.load(row_logits + offsets, mask=in_vocab, other=0.0).to(tl.float32)
+            grad = -scale * tl.exp(values - log_partition)
+            grad = tl.where(offsets == target, grad + scale, grad)
+            tl.store(row_logits + offsets, grad.to(logits.dtype.element_ty), mask=in_vocab)
+
+
+class _LinearTokenLogprobFunction(torch.autograd.Function):
+    """Per-row ``log_softmax(states @ weight.T)[target]`` without keeping logits.
+
+    Preference losses weight every sequence differently, so the gradient
+    cannot be precomputed in forward like the mean LM loss. Forward keeps one
+    log-partition per row; backward recomputes each logits chunk and turns it
+    into the scaled softmax gradient in place.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        states: torch.Tensor,
+        weight: torch.Tensor,
+        targets: torch.Tensor,
+        ignore_index: int,
+    ) -> torch.Tensor:
+        states = states.contiguous()
+        targets = targets.contiguous()
+        rows, vocab_size = states.shape[0], weight.shape[0]
+        negative_logprobs = states.new_empty(rows, dtype=torch.float32)
+        log_partitions = states.new_empty(rows, dtype=torch.float32)
+        unused_count = states.new_ones((), dtype=torch.float32)
+        chunk_rows = max(1, LOGITS_CHUNK_BYTES // (vocab_size * states.element_size()))
+        for start in range(0, rows, chunk_rows):
+            end = min(start + chunk_rows, rows)
+            logits = states[start:end] @ weight.t()
+            _cross_entropy_rows_kernel[(end - start,)](
+                logits,
+                targets[start:end],
+                negative_logprobs[start:end],
+                log_partitions[start:end],
+                unused_count,
+                logits.stride(0),
+                vocab_size,
+                0.0,
+                ignore_index,
+                COMPUTE_GRAD=False,
+                BLOCK=VOCAB_BLOCK_SIZE,
+                num_warps=8,
+            )
+        ctx.save_for_backward(states, weight, targets, log_partitions)
+        ctx.ignore_index = ignore_index
+        ctx.chunk_rows = chunk_rows
+        return negative_logprobs.neg_()
+
+    @staticmethod
+    def backward(
+        ctx, grad_logprobs: torch.Tensor
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None, None, None]:
+        states, weight, targets, log_partitions = ctx.saved_tensors
+        row_grads = grad_logprobs.to(torch.float32).contiguous()
+        rows, vocab_size = states.shape[0], weight.shape[0]
+        grad_states = torch.empty_like(states) if ctx.needs_input_grad[0] else None
+        grad_weight = (
+            torch.zeros(weight.shape, dtype=torch.float32, device=weight.device)
+            if ctx.needs_input_grad[1]
+            else None
+        )
+        for start in range(0, rows, ctx.chunk_rows):
+            end = min(start + ctx.chunk_rows, rows)
+            chunk = states[start:end]
+            logits = chunk @ weight.t()
+            _token_logprob_grad_rows_kernel[(end - start,)](
+                logits,
+                targets[start:end],
+                log_partitions[start:end],
+                row_grads[start:end],
+                logits.stride(0),
+                vocab_size,
+                ctx.ignore_index,
+                BLOCK=VOCAB_BLOCK_SIZE,
+                num_warps=8,
+            )
+            if grad_states is not None:
+                torch.mm(logits, weight, out=grad_states[start:end])
+            if grad_weight is not None:
+                if chunk.dtype == torch.float32:
+                    grad_weight.addmm_(logits.t(), chunk)
+                else:
+                    torch.addmm(
+                        grad_weight, logits.t(), chunk, out_dtype=torch.float32, out=grad_weight
+                    )
+        return (
+            grad_states,
+            None if grad_weight is None else grad_weight.to(weight.dtype),
+            None,
+            None,
+        )
+
+
+def linear_token_logprobs(
+    states: torch.Tensor,
+    weight: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    ignore_index: int = -100,
+) -> torch.Tensor:
+    """fp32 ``log_softmax(states @ weight.T)[row, target]`` per row without full logits.
+
+    Rows whose target is ``ignore_index`` return 0 and receive no gradient.
+    """
+
+    if states.ndim != 2 or targets.shape != states.shape[:1]:
+        raise ValueError("states must be [tokens, hidden] and targets [tokens]")
+    if not states.is_cuda:
+        return linear_token_logprobs_reference(
+            states, weight, targets, ignore_index=ignore_index
+        )
+    return _LinearTokenLogprobFunction.apply(states, weight, targets, ignore_index)

tiny_gdn/gated_residual_ops.py

@@ -0,0 +1,343 @@
+"""Fused CUDA kernels for Senryu-2's gated residual read and write.
+
+The read normalizes every branch, gates it through a rank-``r`` bottleneck over
+all branches and averages the gated branches; the write adds one scaled update
+to every branch. The bottleneck GEMMs run on cuBLAS and every elementwise or
+reduction stage around them is a single Triton pass over the branches, which
+replaces dozens of memory-bound PyTorch ops per read on 32k-token rows.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+GEMM_ALIGNMENT = 8
+COLUMN_BLOCK = 512
+GRAD_ROWS_PER_PROGRAM = 32
+
+
+def _warps(block: int) -> int:
+    return 8 if block >= 2048 else 4
+
+
+@triton.jit
+def _read_normalize_kernel(
+    residual,
+    gain,
+    flat,
+    rstd,
+    hidden_size,
+    eps,
+    BRANCHES: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    vector = tl.program_id(0).to(tl.int64)
+    branch = vector % BRANCHES
+    columns = tl.arange(0, BLOCK)
+    mask = columns < hidden_size
+    offsets = vector * hidden_size + columns
+    values = tl.load(residual + offsets, mask=mask, other=0.0).to(tl.float32)
+    inverse_rms = tl.rsqrt(tl.sum(values * values, axis=0) / hidden_size + eps)
+    scale = tl.load(gain + branch * hidden_size + columns, mask=mask, other=0.0)
+    tl.store(rstd + vector, inverse_rms)
+    tl.store(flat + offsets, (values * inverse_rms * scale).to(flat.dtype.element_ty), mask=mask)
+
+
+@triton.jit
+def _read_combine_kernel(
+    residual,
+    rstd,
+    gain,
+    gate_logits,
+    out,
+    hidden_size,
+    BRANCHES: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    columns = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    mask = columns < hidden_size
+    total = tl.zeros([BLOCK], dtype=tl.float32)
+    for branch in tl.static_range(BRANCHES):
+        vector = row * BRANCHES + branch
+        offsets = vector * hidden_size + columns
+        values = tl.load(residual + offsets, mask=mask, other=0.0).to(tl.float32)
+        logits = tl.load(gate_logits + offsets, mask=mask, other=0.0).to(tl.float32)
+        scale = tl.load(gain + branch * hidden_size + columns, mask=mask, other=0.0)
+        total += tl.sigmoid(logits) * (values * tl.load(rstd + vector) * scale)
+    tl.store(out + row * hidden_size + columns, (total / BRANCHES).to(out.dtype.element_ty), mask=mask)
+
+
+@triton.jit
+def _read_backward_gates_kernel(
+    residual,
+    rstd,
+    gain,
+    gate_logits,
+    grad_out,
+    flat,
+    grad_gate_logits,
+    hidden_size,
+    BRANCHES: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    vector = tl.program_id(0).to(tl.int64)
+    row = vector // BRANCHES
+    branch = vector % BRANCHES
+    columns = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    mask = columns < hidden_size
+    offsets = vector * hidden_size + columns
+    values = tl.load(residual + offsets, mask=mask, other=0.0).to(tl.float32)
+    scale = tl.load(gain + branch * hidden_size + columns, mask=mask, other=0.0)
+    hat = values * tl.load(rstd + vector) * scale
+    gates = tl.sigmoid(tl.load(gate_logits + offsets, mask=mask, other=0.0).to(tl.float32))
+    grad = tl.load(grad_out + row * hidden_size + columns, mask=mask, other=0.0).to(tl.float32) / BRANCHES
+    tl.store(flat + offsets, hat.to(flat.dtype.element_ty), mask=mask)
+    grad_logits = grad * hat * gates * (1.0 - gates)
+    tl.store(grad_gate_logits + offsets, grad_logits.to(grad_gate_logits.dtype.element_ty), mask=mask)
+
+
+@triton.jit
+def _read_backward_residual_kernel(
+    residual,
+    rstd,
+    gain,
+    gate_logits,
+    grad_out,
+    grad_flat,
+    grad_residual,
+    grad_gain_partial,
+    rows,
+    hidden_size,
+    BRANCHES: tl.constexpr,
+    ROWS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    block = tl.program_id(0).to(tl.int64)
+    branch = tl.program_id(1)
+    columns = tl.arange(0, BLOCK)
+    mask = columns < hidden_size
+    scale = tl.load(gain + branch * hidden_size + columns, mask=mask, other=0.0)
+    grad_gain = tl.zeros([BLOCK], dtype=tl.float32)
+    for index in range(ROWS):
+        row = block * ROWS + index
+        valid = row < rows
+        row_mask = mask & valid
+        vector = row * BRANCHES + branch
+        offsets = vector * hidden_size + columns
+        inverse_rms = tl.load(rstd + vector, mask=valid, other=0.0)
+        normalized = tl.load(residual + offsets, mask=row_mask, other=0.0).to(tl.float32) * inverse_rms
+        gates = tl.sigmoid(tl.load(gate_logits + offsets, mask=row_mask, other=0.0).to(tl.float32))
+        grad = tl.load(grad_out + row * hidden_size + columns, mask=row_mask, other=0.0).to(tl.float32)
+        grad_hat = gates * (grad / BRANCHES)
+        grad_hat += tl.load(grad_flat + offsets, mask=row_mask, other=0.0).to(tl.float32)
+        grad_gain += grad_hat * normalized
+        grad_normalized = grad_hat * scale
+        mean = tl.sum(grad_normalized * normalized, axis=0) / hidden_size
+        result = inverse_rms * (grad_normalized - normalized * mean)
+        tl.store(grad_residual + offsets, result.to(grad_residual.dtype.element_ty), mask=row_mask)
+    tl.store(grad_gain_partial + (block * BRANCHES + branch) * hidden_size + columns, grad_gain, mask=mask)
+
+
+@triton.jit
+def _write_forward_kernel(
+    residual,
+    scales,
+    update,
+    out,
+    hidden_size,
+    BRANCHES: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    columns = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    mask = columns < hidden_size
+    values = tl.load(update + row * hidden_size + columns, mask=mask, other=0.0).to(tl.float32)
+    for branch in tl.static_range(BRANCHES):
+        vector = row * BRANCHES + branch
+        offsets = vector * hidden_size + columns
+        scale = tl.load(scales + vector).to(tl.float32)
+        written = tl.load(residual + offsets, mask=mask, other=0.0).to(tl.float32) + scale * values
+        tl.store(out + offsets, written.to(out.dtype.element_ty), mask=mask)
+
+
+@triton.jit
+def _write_backward_kernel(
+    grad,
+    scales,
+    update,
+    grad_scales,
+    grad_update,
+    hidden_size,
+    BRANCHES: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    columns = tl.arange(0, BLOCK)
+    mask = columns < hidden_size
+    values = tl.load(update + row * hidden_size + columns, mask=mask, other=0.0).to(tl.float32)
+    total = tl.zeros([BLOCK], dtype=tl.float32)
+    for branch in tl.static_range(BRANCHES):
+        vector = row * BRANCHES + branch
+        incoming = tl.load(grad + vector * hidden_size + columns, mask=mask, other=0.0).to(tl.float32)
+        total += incoming * tl.load(scales + vector).to(tl.float32)
+        tl.store(grad_scales + vector, tl.sum(incoming * values, axis=0).to(grad_scales.dtype.element_ty))
+    tl.store(grad_update + row * hidden_size + columns, total.to(grad_update.dtype.element_ty), mask=mask)
+
+
+def _bottleneck_projection(down: torch.Tensor, write: torch.Tensor | None) -> torch.Tensor:
+    """``[down; write]`` padded with zero rows to a cuBLAS-friendly height."""
+
+    rows = down.shape[0] + (0 if write is None else write.shape[0])
+    padded = triton.cdiv(rows, GEMM_ALIGNMENT) * GEMM_ALIGNMENT
+    if write is None and padded == rows:
+        return down
+    projection = down.new_zeros(padded, down.shape[1])
+    projection[: down.shape[0]] = down
+    if write is not None:
+        projection[down.shape[0] : rows] = write
+    return projection
+
+
+class _FusedGatedResidualRead(torch.autograd.Function):
+    """CUDA twin of ``senryu2._GatedResidualRead``: same math, rounding points and saved tensors."""
+
+    @staticmethod
+    def forward(ctx, residual, gamma, down, up, write, eps):
+        residual = residual.contiguous()
+        rows, branches, hidden_size = residual.shape
+        rank = down.shape[0]
+        gain = (1.0 + gamma.float()).contiguous()
+        flat = residual.new_empty(rows, branches * hidden_size, dtype=down.dtype)
+        rstd = residual.new_empty(rows, branches, dtype=torch.float32)
+        out = residual.new_empty(rows, hidden_size)
+        full = triton.next_power_of_2(hidden_size)
+        block = min(COLUMN_BLOCK, full)
+        _read_normalize_kernel[(rows * branches,)](
+            residual, gain, flat, rstd, hidden_size, eps, BRANCHES=branches, BLOCK=full, num_warps=_warps(full)
+        )
+        bottleneck = flat @ _bottleneck_projection(down, write).t()
+        hidden = bottleneck[:, :rank] / branches
+        gate_logits = F.silu(hidden.float()).to(down.dtype) @ up.t()
+        _read_combine_kernel[(rows, triton.cdiv(hidden_size, block))](
+            residual, rstd, gain, gate_logits, out, hidden_size, BRANCHES=branches, BLOCK=block
+        )
+        if write is None:
+            scales = residual.new_empty(rows, 0, dtype=torch.float32)
+        else:
+            scales = 2.0 * torch.sigmoid(bottleneck[:, rank : rank + branches].float() / branches)
+        ctx.save_for_backward(residual, gamma, down, up, write, rstd, hidden, scales)
+        scales_out = scales.to(residual.dtype, copy=True)
+        if write is None:
+            ctx.mark_non_differentiable(scales_out)
+        return out, scales_out
+
+    @staticmethod
+    def backward(ctx, grad_out, grad_scales):
+        residual, gamma, down, up, write, rstd, hidden, scales = ctx.saved_tensors
+        rows, branches, hidden_size = residual.shape
+        rank = down.shape[0]
+        dtype = down.dtype
+        gain = (1.0 + gamma.float()).contiguous()
+        grad_out = grad_out.contiguous()
+        full = triton.next_power_of_2(hidden_size)
+        block = min(COLUMN_BLOCK, full)
+
+        pre = hidden.float()
+        sigmoid_pre = torch.sigmoid(pre)
+        activated = F.silu(pre).to(dtype)
+        gate_logits = activated @ up.t()
+        flat = torch.empty_like(gate_logits)
+        grad_gate_logits = torch.empty_like(gate_logits)
+        _read_backward_gates_kernel[(rows * branches, triton.cdiv(hidden_size, block))](
+            residual, rstd, gain, gate_logits, grad_out, flat, grad_gate_logits, hidden_size,
+            BRANCHES=branches, BLOCK=block,
+        )
+        grad_up = grad_gate_logits.t() @ activated
+        grad_activated = (grad_gate_logits @ up).float()
+        del grad_gate_logits
+
+        projection = _bottleneck_projection(down, write)
+        grad_bottleneck = residual.new_zeros(rows, projection.shape[0], dtype=dtype)
+        grad_pre = grad_activated * sigmoid_pre * (1.0 + pre * (1.0 - sigmoid_pre))
+        grad_bottleneck[:, :rank] = (grad_pre / branches).to(dtype)
+        use_scales = write is not None and grad_scales is not None
+        if use_scales:
+            grad_scale_logits = grad_scales.float() * scales * (1.0 - scales / 2.0)
+            grad_bottleneck[:, rank : rank + branches] = (grad_scale_logits / branches).to(dtype)
+        grad_projection = grad_bottleneck.t() @ flat
+        grad_flat = grad_bottleneck @ projection
+        del flat
+
+        grad_residual = torch.empty_like(residual)
+        programs = triton.cdiv(rows, GRAD_ROWS_PER_PROGRAM)
+        grad_gain_partial = residual.new_empty(programs, branches, hidden_size, dtype=torch.float32)
+        _read_backward_residual_kernel[(programs, branches)](
+            residual, rstd, gain, gate_logits, grad_out, grad_flat, grad_residual, grad_gain_partial,
+            rows, hidden_size, BRANCHES=branches, ROWS=GRAD_ROWS_PER_PROGRAM, BLOCK=full, num_warps=_warps(full),
+        )
+        grad_write = None
+        if write is not None:
+            grad_write = (
+                grad_projection[rank : rank + branches].to(write.dtype) if use_scales else torch.zeros_like(write)
+            )
+        return (
+            grad_residual,
+            grad_gain_partial.sum(dim=0).to(gamma.dtype),
+            grad_projection[:rank].to(down.dtype),
+            grad_up.to(up.dtype),
+            grad_write,
+            None,
+        )
+
+
+class _FusedGatedResidualWrite(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, residual, scales, update):
+        residual = residual.contiguous()
+        scales = scales.contiguous()
+        update = update.contiguous()
+        rows, branches, hidden_size = residual.shape
+        out = torch.empty_like(residual)
+        block = min(COLUMN_BLOCK, triton.next_power_of_2(hidden_size))
+        _write_forward_kernel[(rows, triton.cdiv(hidden_size, block))](
+            residual, scales, update, out, hidden_size, BRANCHES=branches, BLOCK=block
+        )
+        ctx.save_for_backward(scales, update)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad):
+        scales, update = ctx.saved_tensors
+        grad = grad.contiguous()
+        rows, branches, hidden_size = grad.shape
+        grad_scales = torch.empty_like(scales)
+        grad_update = torch.empty_like(update)
+        full = triton.next_power_of_2(hidden_size)
+        _write_backward_kernel[(rows,)](
+            grad, scales, update, grad_scales, grad_update, hidden_size,
+            BRANCHES=branches, BLOCK=full, num_warps=_warps(full),
+        )
+        return grad, grad_scales, grad_update
+
+
+def fused_gated_residual_read(
+    residual: torch.Tensor,
+    gamma: torch.Tensor,
+    down: torch.Tensor,
+    up: torch.Tensor,
+    write: torch.Tensor | None,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """``(x [N, d], scales [N, n])`` from CUDA branches ``[N, n, d]``."""
+
+    return _FusedGatedResidualRead.apply(residual, gamma, down, up, write, eps)
+
+
+def fused_gated_residual_write(residual: torch.Tensor, scales: torch.Tensor, update: torch.Tensor) -> torch.Tensor:
+    """``residual + scales[..., None] * update[:, None]`` on CUDA, in FP32 and rounded once."""
+
+    return _FusedGatedResidualWrite.apply(residual, scales, update)

tiny_gdn/haiku_layers.py

@@ -0,0 +1,199 @@
+"""Haiku mixers: Gated MLA (NoPE), SiTU-GLU, and block Attention Residuals."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.attention import sdpa_kernel
+
+from tiny_gdn.config import TinyGDNConfig
+from tiny_gdn.nn_common import RMSNorm, attention_sdpa_backends, situ_glu_reference
+
+try:
+    from tiny_gdn.fused_ops import situ_glu
+except ImportError:
+    situ_glu = situ_glu_reference
+
+
+class GatedMultiheadLatentAttention(nn.Module):
+    """DeepSeek-style MLA with Kimi K3 NoPE and a full-rank output gate.
+
+    Queries and keys are content-only. Position is left to the KDA layers.
+    KV is compressed through a latent bottleneck, then expanded per head.
+    """
+
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.num_heads = config.num_attention_heads
+        self.qk_dim = config.mla_qk_nope_head_dim
+        self.v_dim = config.mla_v_head_dim
+        self.dropout = config.attention_dropout
+        q_out = self.num_heads * self.qk_dim
+        v_out = self.num_heads * self.v_dim
+        kv_out = self.num_heads * (self.qk_dim + self.v_dim)
+
+        if config.mla_q_lora_rank is None:
+            self.q_proj: nn.Module = nn.Linear(config.hidden_size, q_out, bias=False)
+        else:
+            self.q_proj = nn.Sequential(
+                nn.Linear(config.hidden_size, config.mla_q_lora_rank, bias=False),
+                RMSNorm(config.mla_q_lora_rank, config.rms_norm_eps),
+                nn.Linear(config.mla_q_lora_rank, q_out, bias=False),
+            )
+        self.kv_down = nn.Linear(config.hidden_size, config.mla_kv_lora_rank, bias=False)
+        self.kv_norm = RMSNorm(config.mla_kv_lora_rank, config.rms_norm_eps)
+        self.kv_up = nn.Linear(config.mla_kv_lora_rank, kv_out, bias=False)
+        self.q_norm = RMSNorm(self.qk_dim, config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.qk_dim, config.rms_norm_eps)
+        self.output_gate_proj = nn.Linear(config.hidden_size, v_out, bias=False)
+        self.o_proj = nn.Linear(v_out, config.hidden_size, bias=False)
+
+    def _attention_mask(
+        self,
+        attention_mask: torch.Tensor | None,
+        sequence_length: int,
+        device: torch.device,
+    ) -> torch.Tensor | None:
+        if attention_mask is None:
+            return None
+        if attention_mask.ndim != 2:
+            raise ValueError("attention_mask must have shape [batch, sequence]")
+        if attention_mask.shape[1] != sequence_length:
+            raise ValueError("attention_mask sequence length does not match input")
+        causal = torch.ones(
+            sequence_length,
+            sequence_length,
+            dtype=torch.bool,
+            device=device,
+        ).tril()
+        valid_keys = attention_mask[:, None, None, :].to(dtype=torch.bool, device=device)
+        return causal[None, None, :, :] & valid_keys
+
+    def _project_kv(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size, sequence_length, _ = hidden_states.shape
+        compressed = self.kv_norm(self.kv_down(hidden_states))
+        key_value = self.kv_up(compressed)
+        key, value = key_value.split(
+            [self.num_heads * self.qk_dim, self.num_heads * self.v_dim],
+            dim=-1,
+        )
+        key = self.k_norm(
+            key.view(batch_size, sequence_length, self.num_heads, self.qk_dim)
+        ).transpose(1, 2)
+        value = value.view(
+            batch_size, sequence_length, self.num_heads, self.v_dim
+        ).transpose(1, 2)
+        return key, value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+        batch_size, sequence_length, _ = hidden_states.shape
+        query = self.q_proj(hidden_states)
+        query = self.q_norm(
+            query.view(batch_size, sequence_length, self.num_heads, self.qk_dim)
+        ).transpose(1, 2)
+        key, value = self._project_kv(hidden_states)
+        if past_key_value is not None:
+            key = torch.cat([past_key_value[0], key], dim=2)
+            value = torch.cat([past_key_value[1], value], dim=2)
+        present = (key, value)
+        past_len = 0 if past_key_value is None else past_key_value[0].shape[2]
+        kv_len = key.shape[2]
+
+        if attention_mask is not None and past_key_value is None:
+            sdpa_mask = self._attention_mask(
+                attention_mask, sequence_length, hidden_states.device
+            )
+            is_causal = False
+        elif past_key_value is not None and sequence_length == 1:
+            sdpa_mask = (
+                None
+                if attention_mask is None
+                else attention_mask[:, None, None, :].to(
+                    dtype=torch.bool, device=hidden_states.device
+                )
+            )
+            is_causal = False
+        elif past_key_value is not None:
+            q_idx = torch.arange(
+                past_len, past_len + sequence_length, device=hidden_states.device
+            )[:, None]
+            k_idx = torch.arange(kv_len, device=hidden_states.device)[None, :]
+            sdpa_mask = (k_idx <= q_idx)[None, None, :, :]
+            is_causal = False
+        else:
+            sdpa_mask = None
+            is_causal = True
+
+        with sdpa_kernel(attention_sdpa_backends(query.device)):
+            attention_output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                attn_mask=sdpa_mask,
+                dropout_p=self.dropout if self.training else 0.0,
+                is_causal=is_causal,
+            )
+        attention_output = attention_output.transpose(1, 2).reshape(
+            batch_size, sequence_length, -1
+        )
+        attention_output = attention_output * torch.sigmoid(
+            self.output_gate_proj(hidden_states)
+        )
+        return self.o_proj(attention_output), present
+
+
+class SiTUGLU(nn.Module):
+    """Sigmoid-Tanh Unit GLU from Kimi K3.
+
+    Caps the Swish linear factor and the up-projection so the routed / deep
+    stack cannot explode, while matching SwiGLU near the origin.
+    """
+
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.gate_up_proj = nn.Linear(
+            config.hidden_size,
+            config.intermediate_size * 2,
+            bias=False,
+        )
+        self.down_proj = nn.Linear(
+            config.intermediate_size,
+            config.hidden_size,
+            bias=False,
+        )
+        self.gate_cap = config.situ_gate_cap
+        self.up_cap = config.situ_up_cap
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        gate_up = self.gate_up_proj(hidden_states)
+        return self.down_proj(situ_glu(gate_up, self.gate_cap, self.up_cap))
+
+
+class BlockAttentionResidual(nn.Module):
+    """Selective depth skip over block outputs (Kimi K3 AttnRes, block-level)."""
+
+    def __init__(self, hidden_size: int) -> None:
+        super().__init__()
+        self.query = nn.Parameter(torch.zeros(hidden_size))
+        self.scale = 1.0 / math.sqrt(hidden_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        memories: list[torch.Tensor],
+    ) -> torch.Tensor:
+        if not memories:
+            return hidden_states
+        stacked = torch.stack(memories, dim=2)
+        scores = torch.einsum("d,bsnd->bsn", self.query, stacked) * self.scale
+        weights = torch.softmax(scores, dim=-1)
+        mixed = torch.einsum("bsn,bsnd->bsd", weights, stacked)
+        return hidden_states + mixed

tiny_gdn/kdn/__init__.py

@@ -0,0 +1,5 @@
+"""Diagonal Kalman Delta Network (KDN) token mixer and recurrences."""
+
+from tiny_gdn.kdn.layer import SUPPORTED_HEAD_DIMS, DiagonalKalmanDeltaNet
+
+__all__ = ["SUPPORTED_HEAD_DIMS", "DiagonalKalmanDeltaNet"]

tiny_gdn/kdn/layer.py

@@ -0,0 +1,293 @@
+"""Diagonal Kalman Delta Network token mixer (arXiv:2609.07816)."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+from einops import rearrange
+from fla.layers.utils import get_layer_cache, update_layer_cache
+from fla.modules import FusedRMSNormGated, ShortConvolution
+from torch import nn
+
+from tiny_gdn.kdn import recurrence
+from tiny_gdn.kdn.kernels import (
+    _diag_kdn_frontend_precumsum_paired,
+    _diag_kdn_gain,
+    _diag_kdn_normalize,
+    chunk_kalman,
+)
+
+SUPPORTED_HEAD_DIMS = (64, 128)
+
+
+class DiagonalKalmanDeltaNet(nn.Module):
+    """Delta-rule memory whose write strength is a per-channel Kalman gain.
+
+    Each head keeps a diagonal posterior covariance over its key channels:
+    channels that have absorbed little evidence take large corrections and
+    well-determined ones small corrections, replacing KDA's scalar ``beta``.
+    Stateless CUDA BF16 calls run the fused Triton kernels. Cached prefill,
+    decode, padding masks, and other dtypes use the exact PyTorch recurrences.
+
+    The cache holds ``recurrent_state = (memory [B,H,K,K], precision [B,H,K])``
+    in float32 plus the three short-convolution states.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        head_dim: int,
+        num_heads: int,
+        *,
+        conv_size: int = 4,
+        norm_eps: float = 1e-5,
+        layer_idx: int | None = None,
+    ) -> None:
+        super().__init__()
+        if head_dim not in SUPPORTED_HEAD_DIMS:
+            raise ValueError(
+                f"Diagonal KDN head_dim must be one of {SUPPORTED_HEAD_DIMS}, got {head_dim}"
+            )
+        if num_heads <= 0 or hidden_size <= 0:
+            raise ValueError("hidden_size and num_heads must be positive")
+        if conv_size <= 0:
+            raise ValueError("conv_size must be positive")
+        self.hidden_size = hidden_size
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+        self.key_dim = num_heads * head_dim
+        self.layer_idx = layer_idx
+        self.scale = head_dim**-0.5
+        self.info_scale = float(head_dim)
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.q_conv1d, self.k_conv1d, self.v_conv1d = (
+            ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=False,
+                activation="silu",
+            )
+            for _ in range(3)
+        )
+
+        self.f_proj = nn.Sequential(
+            nn.Linear(hidden_size, head_dim, bias=False),
+            nn.Linear(head_dim, self.key_dim, bias=False),
+        )
+        self.A_log = nn.Parameter(
+            torch.log(torch.empty(num_heads, dtype=torch.float32).uniform_(1, 16))
+        )
+        self.A_log._no_weight_decay = True
+        dt = torch.exp(
+            torch.rand(self.key_dim, dtype=torch.float32)
+            * (math.log(0.1) - math.log(0.001))
+            + math.log(0.001)
+        ).clamp(min=1e-4)
+        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
+        self.dt_bias._no_weight_decay = True
+
+        self.omega_proj = nn.Sequential(
+            nn.Linear(hidden_size, head_dim, bias=False),
+            nn.Linear(head_dim, self.key_dim, bias=True),
+        )
+        self.r_proj = nn.Linear(hidden_size, num_heads, bias=True)
+        self.initial_precision_param = nn.Parameter(
+            torch.full(
+                (num_heads,),
+                recurrence.initial_precision_logit(),
+                dtype=torch.float32,
+            )
+        )
+        self.initial_precision_param._no_weight_decay = True
+
+        self.g_proj = nn.Sequential(
+            nn.Linear(hidden_size, head_dim, bias=False),
+            nn.Linear(head_dim, self.key_dim, bias=True),
+        )
+        self.o_norm = FusedRMSNormGated(head_dim, activation="sigmoid", eps=norm_eps)
+        self.o_proj = nn.Linear(self.key_dim, hidden_size, bias=False)
+        self.reset_noise_parameters()
+
+    def reset_noise_parameters(self) -> None:
+        """Isotropic process noise and a neutral observation-noise bias at init.
+
+        Blanket initializers overwrite these, so callers re-run this afterwards.
+        """
+
+        nn.init.zeros_(self.omega_proj[1].weight)
+        nn.init.zeros_(self.omega_proj[1].bias)
+        nn.init.zeros_(self.r_proj.bias)
+
+    def _fused_supported(self, hidden_states: torch.Tensor) -> bool:
+        return hidden_states.is_cuda and hidden_states.dtype == torch.bfloat16
+
+    def _split_heads(self, tensor: torch.Tensor) -> torch.Tensor:
+        return rearrange(tensor, "... (h d) -> ... h d", d=self.head_dim).contiguous()
+
+    def _fused(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        retention_raw: torch.Tensor,
+        omega_raw: torch.Tensor,
+        r_raw: torch.Tensor,
+    ) -> torch.Tensor:
+        q_memory, k_gain, k_memory = _diag_kdn_normalize(q, k)
+        g_cumsum, alpha, omega, r, initial_precision = (
+            _diag_kdn_frontend_precumsum_paired(
+                retention_raw,
+                omega_raw,
+                r_raw,
+                self.A_log.float().contiguous(),
+                self.dt_bias.float().view(self.num_heads, self.head_dim).contiguous(),
+                self.initial_precision_param.float().contiguous(),
+            )
+        )
+        kappa = _diag_kdn_gain(
+            k_gain, alpha, omega, r, initial_precision, omega_raw, r_raw
+        )
+        return chunk_kalman(q_memory, k_memory, kappa, v, g_cumsum, self.scale)
+
+    def _stateful(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        retention_raw: torch.Tensor,
+        omega_raw: torch.Tensor,
+        r_raw: torch.Tensor,
+        state: tuple[torch.Tensor, torch.Tensor] | None,
+        token_mask: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        log_alpha, omega, r, initial_precision = recurrence.kalman_noise(
+            retention_raw,
+            omega_raw,
+            r_raw,
+            self.A_log,
+            self.dt_bias,
+            self.initial_precision_param,
+        )
+        q, k = recurrence.normalize_query_key(q, k)
+        v = v.to(q.dtype)
+        if token_mask is not None:
+            k, v, log_alpha, omega = recurrence.mask_padding(
+                token_mask, k, v, log_alpha, omega
+            )
+        if state is None:
+            memory, precision = recurrence.initial_state(
+                initial_precision,
+                batch_size=q.shape[0],
+                key_dim=self.head_dim,
+                value_dim=self.head_dim,
+            )
+        else:
+            memory, precision = state
+        if q.shape[1] == 1:
+            output, memory, precision = recurrence.diag_kdn_step(
+                q[:, 0],
+                k[:, 0],
+                v[:, 0],
+                log_alpha[:, 0],
+                omega[:, 0],
+                r[:, 0],
+                memory,
+                precision,
+                scale=self.scale,
+                info_scale=self.info_scale,
+            )
+            output = output[:, None]
+        else:
+            output, memory, precision = recurrence.diag_kdn_chunked(
+                q,
+                k,
+                v,
+                log_alpha,
+                omega,
+                r,
+                memory,
+                precision,
+                scale=self.scale,
+                info_scale=self.info_scale,
+            )
+        return output, (memory, precision)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Any | None = None,
+        use_cache: bool = False,
+        output_attentions: bool = False,
+        **kwargs: Any,
+    ) -> tuple[torch.Tensor, None, Any | None]:
+        if hidden_states.ndim != 3 or hidden_states.shape[-1] != self.hidden_size:
+            raise ValueError(
+                f"hidden_states must have shape [B, T, {self.hidden_size}]"
+            )
+        if attention_mask is not None and attention_mask.ndim != 2:
+            raise ValueError("attention_mask must be a [batch, sequence] padding mask")
+        if output_attentions:
+            raise ValueError("Diagonal KDN has no attention weights to return")
+        length = hidden_states.shape[1]
+        last_state = get_layer_cache(self, past_key_values)
+        token_mask = (
+            None
+            if attention_mask is None
+            else attention_mask[:, -length:].to(device=hidden_states.device, dtype=torch.bool)
+        )
+        conv_mask = (
+            None if token_mask is None else token_mask.to(hidden_states.dtype)
+        )
+        conv_q, conv_k, conv_v = (
+            (None, None, None) if last_state is None else last_state["conv_state"]
+        )
+        q, conv_q = self.q_conv1d(
+            x=self.q_proj(hidden_states),
+            mask=conv_mask,
+            cache=conv_q,
+            output_final_state=use_cache,
+        )
+        k, conv_k = self.k_conv1d(
+            x=self.k_proj(hidden_states),
+            mask=conv_mask,
+            cache=conv_k,
+            output_final_state=use_cache,
+        )
+        v, conv_v = self.v_conv1d(
+            x=self.v_proj(hidden_states),
+            mask=conv_mask,
+            cache=conv_v,
+            output_final_state=use_cache,
+        )
+        q, k, v = (self._split_heads(tensor) for tensor in (q, k, v))
+        retention_raw = self._split_heads(self.f_proj(hidden_states))
+        omega_raw = self._split_heads(self.omega_proj(hidden_states))
+        r_raw = self.r_proj(hidden_states).contiguous()
+
+        recurrent_state = None if last_state is None else last_state["recurrent_state"]
+        stateless = recurrent_state is None and token_mask is None and not use_cache
+        if stateless and self._fused_supported(hidden_states):
+            output = self._fused(q, k, v, retention_raw, omega_raw, r_raw)
+        else:
+            output, recurrent_state = self._stateful(
+                q, k, v, retention_raw, omega_raw, r_raw, recurrent_state, token_mask
+            )
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_q, conv_k, conv_v),
+            offset=length,
+        )
+        output = self.o_norm(
+            output.to(v.dtype),
+            self._split_heads(self.g_proj(hidden_states)),
+        )
+        output = self.o_proj(rearrange(output, "b t h d -> b t (h d)"))
+        return output, None, past_key_values

tiny_gdn/kdn/recurrence.py

@@ -0,0 +1,311 @@
+"""Exact PyTorch Diagonal KDN recurrences that carry state across calls.
+
+The Triton kernels cover stateless fixed-length training. These functions
+take and return ``(memory, precision)`` so prefill, decode, and padded batches
+continue a sequence exactly; the sequential form is also the test oracle.
+
+Per head, with one covariance coordinate ``p`` per key channel::
+
+    p_hat_t = alpha_t**2 * p_{t-1} + omega_t
+    kappa_t = p_hat_t * k_t / (r_t + sum_i p_hat_ti * k_ti**2)
+    1 / p_t = 1 / p_hat_t + info_scale * k_t**2 / r_t
+    S_t     = (I - kappa_t k_t^T) diag(alpha_t) S_{t-1} + kappa_t v_t^T
+    o_t     = scale * S_t^T q_t
+
+Shapes: ``q, k, log_alpha, omega`` are ``[B, T, H, K]``, ``v`` is
+``[B, T, H, V]``, ``r`` is ``[B, T, H]``, ``memory`` is ``[B, H, K, V]`` and
+``precision`` (``1 / p``) is ``[B, H, K]``. All math runs in float32.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+OMEGA_MIN = 0.0
+R_MIN = 0.01
+INITIAL_PRECISION = 1.0
+INITIAL_PRECISION_FLOOR = 0.1
+LOG_MEMORY_ALPHA_FLOOR = math.log(1e-6)
+NORM_EPS = 1e-12
+DEFAULT_CHUNK_SIZE = 64
+
+
+def initial_precision_logit() -> float:
+    """Raw parameter value whose floored softplus equals ``INITIAL_PRECISION``."""
+
+    return math.log(math.expm1(INITIAL_PRECISION - INITIAL_PRECISION_FLOOR))
+
+
+def _work(tensor: torch.Tensor) -> torch.Tensor:
+    return tensor.to(torch.promote_types(tensor.dtype, torch.float32))
+
+
+def kalman_noise(
+    retention_raw: torch.Tensor,
+    omega_raw: torch.Tensor,
+    r_raw: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    initial_precision_param: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Return ``(log_alpha, omega, r, initial_precision)`` like the Triton frontend."""
+
+    heads, key_dim = retention_raw.shape[-2:]
+    rate = _work(A_log).exp().view(heads, 1)
+    log_alpha = -rate * F.softplus(
+        _work(retention_raw) + _work(dt_bias).view(heads, key_dim)
+    )
+    omega = OMEGA_MIN + F.softplus(_work(omega_raw))
+    r = R_MIN + F.softplus(_work(r_raw))
+    initial_precision = INITIAL_PRECISION_FLOOR + F.softplus(
+        _work(initial_precision_param)
+    )
+    return log_alpha, omega, r, initial_precision
+
+
+def normalize_query_key(
+    q: torch.Tensor, k: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return (
+        F.normalize(_work(q), dim=-1, eps=NORM_EPS),
+        F.normalize(_work(k), dim=-1, eps=NORM_EPS),
+    )
+
+
+def initial_state(
+    initial_precision: torch.Tensor,
+    *,
+    batch_size: int,
+    key_dim: int,
+    value_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Zero memory and the learned per-head prior precision for new sequences."""
+
+    heads = initial_precision.shape[0]
+    prior = _work(initial_precision)
+    memory = prior.new_zeros((batch_size, heads, key_dim, value_dim))
+    precision = prior.view(1, heads, 1).expand(batch_size, heads, key_dim).contiguous()
+    return memory, precision
+
+
+def mask_padding(
+    token_mask: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    log_alpha: torch.Tensor,
+    omega: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Turn padded positions into identity steps that leave the state untouched."""
+
+    keep = token_mask[:, :, None, None]
+    return (
+        torch.where(keep, k, 0.0),
+        torch.where(keep, v, 0.0),
+        torch.where(keep, log_alpha, 0.0),
+        torch.where(keep, omega, 0.0),
+    )
+
+
+def diag_kdn_step(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    log_alpha: torch.Tensor,
+    omega: torch.Tensor,
+    r: torch.Tensor,
+    memory: torch.Tensor,
+    precision: torch.Tensor,
+    *,
+    scale: float,
+    info_scale: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Advance one token: ``q, k, log_alpha, omega`` are ``[B, H, K]``."""
+
+    k_squared = k.square()
+    predicted = (2.0 * log_alpha).exp() / precision + omega
+    noise = r[..., None]
+    gain = predicted * k / (noise + (predicted * k_squared).sum(-1, keepdim=True))
+    precision = predicted.reciprocal() + info_scale * k_squared / noise
+    decayed = memory * log_alpha.clamp_min(LOG_MEMORY_ALPHA_FLOOR).exp()[..., None]
+    innovation = v - torch.einsum("bhk,bhkv->bhv", k, decayed)
+    memory = decayed + gain[..., None] * innovation[..., None, :]
+    output = scale * torch.einsum("bhk,bhkv->bhv", q, memory)
+    return output, memory, precision
+
+
+def diag_kdn_recurrent(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    log_alpha: torch.Tensor,
+    omega: torch.Tensor,
+    r: torch.Tensor,
+    memory: torch.Tensor,
+    precision: torch.Tensor,
+    *,
+    scale: float,
+    info_scale: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Token-by-token reference; returns ``(output [B,T,H,V], memory, precision)``."""
+
+    outputs = []
+    for token in range(q.shape[1]):
+        output, memory, precision = diag_kdn_step(
+            q[:, token],
+            k[:, token],
+            v[:, token],
+            log_alpha[:, token],
+            omega[:, token],
+            r[:, token],
+            memory,
+            precision,
+            scale=scale,
+            info_scale=info_scale,
+        )
+        outputs.append(output)
+    return torch.stack(outputs, dim=1), memory, precision
+
+
+def _compose(
+    later: tuple[torch.Tensor, ...], earlier: tuple[torch.Tensor, ...]
+) -> tuple[torch.Tensor, ...]:
+    """Product ``later @ earlier`` of 2x2 maps stored as (m00, m01, m10, m11)."""
+
+    a00, a01, a10, a11 = later
+    b00, b01, b10, b11 = earlier
+    product = (
+        a00 * b00 + a01 * b10,
+        a00 * b01 + a01 * b11,
+        a10 * b00 + a11 * b10,
+        a10 * b01 + a11 * b11,
+    )
+    # Mobius maps are invariant to scaling, so renormalize to stay in range.
+    magnitude = torch.stack(product).abs().amax(dim=0).clamp_min(
+        torch.finfo(product[0].dtype).tiny
+    )
+    return tuple(entry / magnitude for entry in product)
+
+
+def _inclusive_scan(maps: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
+    """Hillis-Steele prefix products ``M_t @ ... @ M_1`` along dim 1."""
+
+    length = maps[0].shape[1]
+    offset = 1
+    while offset < length:
+        combined = _compose(
+            tuple(entry[:, offset:] for entry in maps),
+            tuple(entry[:, :-offset] for entry in maps),
+        )
+        maps = tuple(
+            torch.cat((entry[:, :offset], update), dim=1)
+            for entry, update in zip(maps, combined)
+        )
+        offset *= 2
+    return maps
+
+
+def kalman_gain(
+    k: torch.Tensor,
+    log_alpha: torch.Tensor,
+    omega: torch.Tensor,
+    r: torch.Tensor,
+    precision: torch.Tensor,
+    *,
+    info_scale: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Vector gains for every token via a parallel Mobius scan, plus final precision."""
+
+    k_squared = k.square()
+    noise = r[..., None]
+    alpha_squared = (2.0 * log_alpha).exp()
+    information = info_scale * k_squared / noise
+    maps = _inclusive_scan(
+        (
+            alpha_squared,
+            omega,
+            information * alpha_squared,
+            1.0 + information * omega,
+        )
+    )
+    start = precision.reciprocal()[:, None]
+    covariance = (maps[0] * start + maps[1]) / (maps[2] * start + maps[3])
+    previous = torch.cat((start, covariance[:, :-1]), dim=1)
+    predicted = alpha_squared * previous + omega
+    gain = predicted * k / (noise + (predicted * k_squared).sum(-1, keepdim=True))
+    return gain, covariance[:, -1].reciprocal()
+
+
+def chunked_memory(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    log_alpha: torch.Tensor,
+    gain: torch.Tensor,
+    memory: torch.Tensor,
+    *,
+    scale: float,
+    chunk_size: int = DEFAULT_CHUNK_SIZE,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Two-key delta-rule memory in WY form, one triangular solve per chunk."""
+
+    if chunk_size <= 0:
+        raise ValueError("chunk_size must be positive")
+    log_alpha = log_alpha.clamp_min(LOG_MEMORY_ALPHA_FLOOR)
+    length = q.shape[1]
+    outputs = []
+    for start in range(0, length, chunk_size):
+        stop = min(start + chunk_size, length)
+        size = stop - start
+        q_chunk, k_chunk, v_chunk, gain_chunk = (
+            tensor[:, start:stop].transpose(1, 2) for tensor in (q, k, v, gain)
+        )
+        cumulative = log_alpha[:, start:stop].transpose(1, 2).cumsum(dim=2)
+        causal = torch.ones(size, size, dtype=torch.bool, device=q.device).tril()
+        differences = cumulative[:, :, :, None, :] - cumulative[:, :, None, :, :]
+        weights = (
+            torch.where(causal[:, :, None], differences, -math.inf).exp()
+            * gain_chunk[:, :, None, :, :]
+        )
+        interaction = torch.einsum("bhtk,bhtsk->bhts", k_chunk, weights)
+        readout = torch.einsum("bhtk,bhtsk->bhts", q_chunk, weights)
+        decay = cumulative.exp()
+        corrected = torch.linalg.solve_triangular(
+            interaction,
+            v_chunk - (decay * k_chunk) @ memory,
+            upper=False,
+            unitriangular=True,
+        )
+        outputs.append(
+            (scale * ((q_chunk * decay) @ memory + readout @ corrected)).transpose(1, 2)
+        )
+        to_end = (cumulative[:, :, -1:] - cumulative).exp() * gain_chunk
+        memory = memory * decay[:, :, -1, :, None] + to_end.transpose(-1, -2) @ corrected
+    return torch.cat(outputs, dim=1), memory
+
+
+def diag_kdn_chunked(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    log_alpha: torch.Tensor,
+    omega: torch.Tensor,
+    r: torch.Tensor,
+    memory: torch.Tensor,
+    precision: torch.Tensor,
+    *,
+    scale: float,
+    info_scale: float,
+    chunk_size: int = DEFAULT_CHUNK_SIZE,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Parallel form of :func:`diag_kdn_recurrent` with the same signature."""
+
+    gain, precision = kalman_gain(
+        k, log_alpha, omega, r, precision, info_scale=info_scale
+    )
+    output, memory = chunked_memory(
+        q, k, v, log_alpha, gain, memory, scale=scale, chunk_size=chunk_size
+    )
+    return output, memory, precision

tiny_gdn/model.py

@@ -0,0 +1,782 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from safetensors.torch import load_model, save_model
+from torch import nn
+from torch.nn.attention import sdpa_kernel
+from torch.utils.checkpoint import checkpoint
+
+from tiny_gdn.config import LINEAR_LAYER_TYPES, TinyGDNConfig
+from tiny_gdn.haiku_layers import (
+    BlockAttentionResidual,
+    GatedMultiheadLatentAttention,
+    SiTUGLU,
+)
+from tiny_gdn.nn_common import (
+    RMSNorm,
+    attention_sdpa_backends,
+    linear_cross_entropy_reference,
+)
+
+try:
+    from tiny_gdn.fused_ops import linear_cross_entropy
+except ImportError:
+    linear_cross_entropy = linear_cross_entropy_reference
+
+try:
+    # Import the module directly — `from fla.layers import GatedDeltaNet2`
+    # executes layers/__init__.py and eagerly loads every attention kernel.
+    from fla.layers.gdn2 import GatedDeltaNet2
+except ImportError as import_error:
+    GatedDeltaNet2 = None
+    FLA_IMPORT_ERROR: ImportError | None = import_error
+else:
+    FLA_IMPORT_ERROR = None
+
+try:
+    from fla.layers.kda import KimiDeltaAttention
+except ImportError as import_error:
+    KimiDeltaAttention = None
+    KDA_IMPORT_ERROR: ImportError | None = import_error
+else:
+    KDA_IMPORT_ERROR = None
+
+try:
+    from tiny_gdn.kdn import DiagonalKalmanDeltaNet
+except ImportError as import_error:
+    DiagonalKalmanDeltaNet = None
+    KDN_IMPORT_ERROR: ImportError | None = import_error
+else:
+    KDN_IMPORT_ERROR = None
+
+
+@dataclass
+class TinyGDNOutput:
+    loss: torch.Tensor | None
+    logits: torch.Tensor | None
+    main_loss: torch.Tensor | None
+    mtp_loss: torch.Tensor | None
+    z_loss: torch.Tensor | None
+    hidden_states: torch.Tensor | None = None
+    past_key_values: Any | None = None
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, rotary_dim: int, rope_theta: float) -> None:
+        super().__init__()
+        inverse_frequency = 1.0 / (
+            rope_theta
+            ** (
+                torch.arange(0, rotary_dim, 2, dtype=torch.float32)
+                / rotary_dim
+            )
+        )
+        self.rotary_dim = rotary_dim
+        self.register_buffer("inverse_frequency", inverse_frequency, persistent=False)
+
+    def forward(
+        self,
+        sequence_length: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        position_offset: int = 0,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        positions = torch.arange(
+            position_offset,
+            position_offset + sequence_length,
+            device=device,
+            dtype=torch.float32,
+        )
+        frequencies = torch.outer(positions, self.inverse_frequency.float())
+        embeddings = torch.cat((frequencies, frequencies), dim=-1)
+        return embeddings.cos().to(dtype=dtype), embeddings.sin().to(dtype=dtype)
+
+
+def rotate_half(hidden_states: torch.Tensor) -> torch.Tensor:
+    first, second = hidden_states.chunk(2, dim=-1)
+    return torch.cat((-second, first), dim=-1)
+
+
+def apply_rotary_embedding(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cosine: torch.Tensor,
+    sine: torch.Tensor,
+    rotary_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    cosine = cosine[None, None, :, :]
+    sine = sine[None, None, :, :]
+    query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:]
+    key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:]
+    query_rotary = query_rotary * cosine + rotate_half(query_rotary) * sine
+    key_rotary = key_rotary * cosine + rotate_half(key_rotary) * sine
+    return (
+        torch.cat((query_rotary, query_pass), dim=-1),
+        torch.cat((key_rotary, key_pass), dim=-1),
+    )
+
+
+class GatedGroupedQueryAttention(nn.Module):
+    """QK-normalized, partially rotary GQA with a learned sigmoid output gate."""
+
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.head_dim = config.attention_head_dim
+        self.rotary_dim = config.rotary_dim
+        self.dropout = config.attention_dropout
+
+        query_size = self.num_heads * self.head_dim
+        key_value_size = self.num_key_value_heads * self.head_dim
+        self.q_gate_proj = nn.Linear(config.hidden_size, query_size * 2, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, key_value_size, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, key_value_size, bias=False)
+        self.o_proj = nn.Linear(query_size, config.hidden_size, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        self.rotary = RotaryEmbedding(self.rotary_dim, config.rope_theta)
+
+    def _attention_mask(
+        self,
+        attention_mask: torch.Tensor | None,
+        sequence_length: int,
+        device: torch.device,
+    ) -> torch.Tensor | None:
+        if attention_mask is None:
+            return None
+        if attention_mask.ndim != 2:
+            raise ValueError("attention_mask must have shape [batch, sequence]")
+        if attention_mask.shape[1] != sequence_length:
+            raise ValueError("attention_mask sequence length does not match input")
+
+        causal = torch.ones(
+            sequence_length,
+            sequence_length,
+            dtype=torch.bool,
+            device=device,
+        ).tril()
+        valid_keys = attention_mask[:, None, None, :].to(dtype=torch.bool, device=device)
+        return causal[None, None, :, :] & valid_keys
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+        batch_size, sequence_length, _ = hidden_states.shape
+        past_len = 0 if past_key_value is None else past_key_value[0].shape[2]
+        query_and_gate = self.q_gate_proj(hidden_states)
+        query, output_gate = query_and_gate.chunk(2, dim=-1)
+
+        query = query.view(batch_size, sequence_length, self.num_heads, self.head_dim)
+        key = self.k_proj(hidden_states).view(
+            batch_size,
+            sequence_length,
+            self.num_key_value_heads,
+            self.head_dim,
+        )
+        value = self.v_proj(hidden_states).view(
+            batch_size,
+            sequence_length,
+            self.num_key_value_heads,
+            self.head_dim,
+        )
+
+        query = self.q_norm(query).transpose(1, 2)
+        key = self.k_norm(key).transpose(1, 2)
+        value = value.transpose(1, 2)
+
+        cosine, sine = self.rotary(
+            sequence_length,
+            device=hidden_states.device,
+            dtype=query.dtype,
+            position_offset=past_len,
+        )
+        query, key = apply_rotary_embedding(
+            query,
+            key,
+            cosine,
+            sine,
+            rotary_dim=self.rotary_dim,
+        )
+        if past_key_value is not None:
+            key = torch.cat([past_key_value[0], key], dim=2)
+            value = torch.cat([past_key_value[1], value], dim=2)
+        present = (key, value)
+
+        kv_len = key.shape[2]
+        if attention_mask is not None and past_key_value is None:
+            sdpa_mask = self._attention_mask(
+                attention_mask,
+                sequence_length,
+                hidden_states.device,
+            )
+            is_causal = False
+        elif past_key_value is not None and sequence_length == 1:
+            # Decode step: query attends to the cached key/value prefix. Preserve
+            # the prefill padding mask when decoding a left-padded prompt batch.
+            sdpa_mask = (
+                None
+                if attention_mask is None
+                else attention_mask[:, None, None, :].to(
+                    dtype=torch.bool,
+                    device=hidden_states.device,
+                )
+            )
+            is_causal = False
+        elif past_key_value is not None:
+            # Prefill chunk with cache — build causal mask over kv_len.
+            q_idx = torch.arange(
+                past_len, past_len + sequence_length, device=hidden_states.device
+            )[:, None]
+            k_idx = torch.arange(kv_len, device=hidden_states.device)[None, :]
+            sdpa_mask = (k_idx <= q_idx)[None, None, :, :]
+            is_causal = False
+        else:
+            sdpa_mask = None
+            is_causal = True
+        sdpa_options = {
+            "attn_mask": sdpa_mask,
+            "dropout_p": self.dropout if self.training else 0.0,
+            "is_causal": is_causal,
+            "enable_gqa": True,
+        }
+        with sdpa_kernel(attention_sdpa_backends(query.device)):
+            attention_output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                **sdpa_options,
+            )
+        attention_output = attention_output.transpose(1, 2).reshape(
+            batch_size,
+            sequence_length,
+            -1,
+        )
+        attention_output = attention_output * torch.sigmoid(output_gate)
+        return self.o_proj(attention_output), present
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.gate_up_proj = nn.Linear(
+            config.hidden_size,
+            config.intermediate_size * 2,
+            bias=False,
+        )
+        self.down_proj = nn.Linear(
+            config.intermediate_size,
+            config.hidden_size,
+            bias=False,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        gate, up = self.gate_up_proj(hidden_states).chunk(2, dim=-1)
+        return self.down_proj(F.silu(gate) * up)
+
+
+class TinyGDNBlock(nn.Module):
+    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
+        super().__init__()
+        layer_type = config.layer_types[layer_index]
+        self.layer_type = layer_type
+        self.token_mixer_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+
+        if layer_type == "gdn2":
+            if GatedDeltaNet2 is None:
+                raise ImportError(
+                    "Gated DeltaNet-2 requires the pinned flash-linear-attention dependency"
+                ) from FLA_IMPORT_ERROR
+            self.token_mixer = GatedDeltaNet2(
+                hidden_size=config.hidden_size,
+                expand_v=config.linear_expand_v,
+                head_dim=config.linear_head_dim,
+                num_heads=config.linear_num_heads,
+                num_v_heads=config.linear_num_value_heads,
+                mode="chunk",
+                use_short_conv=True,
+                allow_neg_eigval=config.allow_negative_eigenvalues,
+                conv_size=config.linear_conv_kernel_dim,
+                conv_bias=False,
+                layer_idx=layer_index,
+                norm_eps=config.rms_norm_eps,
+            )
+        elif layer_type == "kda":
+            if KimiDeltaAttention is None:
+                raise ImportError(
+                    "Kimi Delta Attention requires the pinned flash-linear-attention dependency"
+                ) from KDA_IMPORT_ERROR
+            self.token_mixer = KimiDeltaAttention(
+                hidden_size=config.hidden_size,
+                expand_v=config.linear_expand_v,
+                head_dim=config.linear_head_dim,
+                num_heads=config.linear_num_heads,
+                num_v_heads=config.linear_num_value_heads,
+                mode="chunk",
+                use_short_conv=True,
+                allow_neg_eigval=config.allow_negative_eigenvalues,
+                safe_gate=config.kda_safe_gate,
+                lower_bound=config.kda_lower_bound,
+                conv_size=config.linear_conv_kernel_dim,
+                conv_bias=False,
+                layer_idx=layer_index,
+                norm_eps=config.rms_norm_eps,
+            )
+        elif layer_type == "diag_kdn":
+            if DiagonalKalmanDeltaNet is None:
+                raise ImportError(
+                    "Diagonal KDN requires Triton and the pinned flash-linear-attention dependency"
+                ) from KDN_IMPORT_ERROR
+            self.token_mixer = DiagonalKalmanDeltaNet(
+                hidden_size=config.hidden_size,
+                head_dim=config.linear_head_dim,
+                num_heads=config.linear_num_heads,
+                conv_size=config.linear_conv_kernel_dim,
+                norm_eps=config.rms_norm_eps,
+                layer_idx=layer_index,
+            )
+        elif layer_type == "full_attention":
+            self.token_mixer = GatedGroupedQueryAttention(config)
+        elif layer_type == "gated_mla":
+            self.token_mixer = GatedMultiheadLatentAttention(config)
+        else:
+            raise ValueError(f"Unsupported layer type: {layer_type}")
+
+        if config.mlp_activation == "swiglu":
+            self.mlp: nn.Module = SwiGLU(config)
+        elif config.mlp_activation == "situ_glu":
+            self.mlp = SiTUGLU(config)
+        else:
+            raise ValueError(f"Unsupported mlp_activation: {config.mlp_activation}")
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        *,
+        past_key_values: Any | None = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+        use_cache: bool = False,
+    ) -> tuple[torch.Tensor, Any]:
+        residual = hidden_states
+        normalized = self.token_mixer_norm(hidden_states)
+        present: Any = None
+        if self.layer_type in LINEAR_LAYER_TYPES:
+            mixed, _, past_key_values = self.token_mixer(
+                normalized,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+            )
+            present = past_key_values
+        elif self.layer_type in {"full_attention", "gated_mla"}:
+            mixed, present = self.token_mixer(
+                normalized,
+                attention_mask=attention_mask,
+                past_key_value=past_key_value,
+            )
+            if not use_cache:
+                present = None
+        else:
+            raise ValueError(f"Unsupported layer type: {self.layer_type}")
+        hidden_states = residual + mixed
+        hidden_states = hidden_states + self.mlp(self.mlp_norm(hidden_states))
+        return hidden_states, present
+
+
+class MultiTokenPredictionAdapter(nn.Module):
+    """A lightweight residual adapter for one additional prediction horizon."""
+
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.down_proj = nn.Linear(
+            config.hidden_size,
+            config.mtp_adapter_rank,
+            bias=False,
+        )
+        self.up_proj = nn.Linear(
+            config.mtp_adapter_rank,
+            config.hidden_size,
+            bias=False,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        adapted = self.up_proj(F.silu(self.down_proj(self.norm(hidden_states))))
+        return hidden_states + adapted
+
+
+class TinyGDNForCausalLM(nn.Module):
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            TinyGDNBlock(config, layer_index)
+            for layer_index in range(config.num_hidden_layers)
+        )
+        self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.mtp_adapters = nn.ModuleList(
+            MultiTokenPredictionAdapter(config)
+            for _ in range(config.mtp_num_heads)
+        )
+        self.block_attn_res = (
+            BlockAttentionResidual(config.hidden_size)
+            if config.use_block_attn_res
+            else None
+        )
+        self.checkpointed_layers: frozenset[int] = frozenset()
+
+        self.apply(self._initialize_module)
+        self._initialize_residual_projections()
+
+    def _initialize_module(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(
+                module.weight,
+                mean=0.0,
+                std=self.config.initializer_range,
+            )
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(
+                module.weight,
+                mean=0.0,
+                std=self.config.initializer_range,
+            )
+        elif DiagonalKalmanDeltaNet is not None and isinstance(
+            module, DiagonalKalmanDeltaNet
+        ):
+            # apply() visits children first, so this runs after the Linear re-init.
+            module.reset_noise_parameters()
+
+    def _initialize_residual_projections(self) -> None:
+        residual_std = self.config.initializer_range / math.sqrt(
+            2 * self.config.num_hidden_layers
+        )
+        for layer in self.layers:
+            nn.init.normal_(
+                layer.token_mixer.o_proj.weight,
+                mean=0.0,
+                std=residual_std,
+            )
+            nn.init.normal_(
+                layer.mlp.down_proj.weight,
+                mean=0.0,
+                std=residual_std,
+            )
+        for adapter in self.mtp_adapters:
+            nn.init.normal_(adapter.up_proj.weight, mean=0.0, std=residual_std)
+
+    def enable_gradient_checkpointing(
+        self,
+        enabled: bool = True,
+        *,
+        layers: int | None = None,
+    ) -> None:
+        """Recompute block activations during backward instead of storing them.
+
+        ``layers`` checkpoints only that many blocks, linear-attention blocks
+        first: replaying softmax attention costs quadratically in sequence
+        length, replaying a linear-attention block costs linearly. ``None``
+        checkpoints every block.
+        """
+
+        if not enabled:
+            if layers is not None:
+                raise ValueError("layers needs gradient checkpointing enabled")
+            self.checkpointed_layers = frozenset()
+            return
+        count = len(self.layers) if layers is None else layers
+        if not 1 <= count <= len(self.layers):
+            raise ValueError(f"layers must be between 1 and {len(self.layers)}")
+        cheapest_replay_first = sorted(
+            range(len(self.layers)),
+            key=lambda index: self.layers[index].layer_type not in LINEAR_LAYER_TYPES,
+        )
+        self.checkpointed_layers = frozenset(cheapest_replay_first[:count])
+
+    def describe_gradient_checkpointing(self) -> str:
+        if not self.checkpointed_layers:
+            return "off"
+        layer_types = sorted(
+            {self.layers[index].layer_type for index in self.checkpointed_layers}
+        )
+        return (
+            f"{len(self.checkpointed_layers)}/{len(self.layers)} blocks "
+            f"({', '.join(layer_types)})"
+        )
+
+    def project_to_vocabulary(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return F.linear(hidden_states, self.embed_tokens.weight)
+
+    def _run_layer(
+        self,
+        layer_index: int,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        *,
+        past_key_values: Any | None = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+        use_cache: bool = False,
+    ) -> tuple[torch.Tensor, Any]:
+        layer = self.layers[layer_index]
+        if self.training and layer_index in self.checkpointed_layers:
+            hidden_states, present = checkpoint(
+                layer,
+                hidden_states,
+                attention_mask,
+                use_reentrant=False,
+            )
+            return hidden_states, present
+        return layer(
+            hidden_states,
+            attention_mask,
+            past_key_values=past_key_values,
+            past_key_value=past_key_value,
+            use_cache=use_cache,
+        )
+
+    def _causal_loss(
+        self,
+        hidden_states: torch.Tensor,
+        labels: torch.Tensor,
+        target_offset: int,
+        adapter: nn.Module | None = None,
+        z_loss_coefficient: float = 0.0,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Return ``(objective, cross_entropy, z_loss)``.
+
+        Only ``objective = cross_entropy + z_loss_coefficient * z_loss`` carries
+        gradients; the other two are detached values for logging.
+        """
+
+        if target_offset < 0:
+            raise ValueError("target_offset cannot be negative")
+        if target_offset and hidden_states.shape[1] <= target_offset:
+            raise ValueError(
+                f"Sequence length must exceed target offset {target_offset}"
+            )
+        if target_offset:
+            prediction_states = hidden_states[:, :-target_offset, :]
+            targets = labels[:, target_offset:].contiguous()
+        else:
+            prediction_states = hidden_states
+            targets = labels.contiguous()
+        if adapter is not None:
+            prediction_states = adapter(prediction_states)
+        return linear_cross_entropy(
+            prediction_states.reshape(-1, prediction_states.shape[-1]),
+            self.embed_tokens.weight,
+            targets.reshape(-1),
+            z_loss_coefficient=z_loss_coefficient,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Any | None = None,
+        *,
+        use_cache: bool = False,
+        return_logits: bool = True,
+        return_hidden_states: bool = False,
+        labels_are_shifted: bool = False,
+        include_mtp_loss: bool = True,
+        mtp_loss_weight: float | None = None,
+        z_loss_coefficient: float = 0.0,
+        logits_to_keep: int | None = None,
+    ) -> TinyGDNOutput:
+        if input_ids.ndim != 2:
+            raise ValueError("input_ids must have shape [batch, sequence]")
+        if input_ids.shape[1] > self.config.max_position_embeddings:
+            raise ValueError("Input exceeds max_position_embeddings")
+        if labels is not None and labels.shape != input_ids.shape:
+            raise ValueError("labels must have the same shape as input_ids")
+        if z_loss_coefficient < 0.0:
+            raise ValueError("z_loss_coefficient cannot be negative")
+        if logits_to_keep is not None and logits_to_keep <= 0:
+            raise ValueError("logits_to_keep must be positive")
+        if use_cache and labels is not None:
+            raise ValueError("use_cache is not supported with labels")
+        effective_mtp_weight = (
+            self.config.mtp_loss_weight
+            if mtp_loss_weight is None
+            else mtp_loss_weight
+        )
+        if not 0.0 <= effective_mtp_weight <= 1.0:
+            raise ValueError("mtp_loss_weight must be between zero and one")
+
+        if use_cache and past_key_values is None:
+            try:
+                from fla.models.utils import Cache as FlaCache
+            except ImportError as import_error:
+                raise ImportError(
+                    "Cached decode requires flash-linear-attention Cache"
+                ) from import_error
+            past_key_values = {
+                "fla": FlaCache(),
+                "gqa": [None] * len(self.layers),
+            }
+        elif past_key_values is not None and not isinstance(past_key_values, dict):
+            raise TypeError("past_key_values must be a TinyGDN cache dict or None")
+
+        fla_cache = None if past_key_values is None else past_key_values["fla"]
+        gqa_cache = None if past_key_values is None else past_key_values["gqa"]
+
+        hidden_states = self.embed_tokens(input_ids)
+        shared_layer_indices = set(self.config.shared_layer_indices)
+        softmax_types = {"full_attention", "gated_mla"}
+        attn_memories: list[torch.Tensor] = [hidden_states]
+        interval = self.config.full_attention_interval
+        for layer_index, layer in enumerate(self.layers):
+            layer_gqa = None if gqa_cache is None else gqa_cache[layer_index]
+            hidden_states, present = self._run_layer(
+                layer_index,
+                hidden_states,
+                attention_mask,
+                past_key_values=fla_cache,
+                past_key_value=layer_gqa,
+                use_cache=use_cache,
+            )
+            if use_cache and layer.layer_type in softmax_types and gqa_cache is not None:
+                gqa_cache[layer_index] = present
+            if layer_index in shared_layer_indices:
+                layer_gqa = None if gqa_cache is None else gqa_cache[layer_index]
+                hidden_states, present = self._run_layer(
+                    layer_index,
+                    hidden_states,
+                    attention_mask,
+                    past_key_values=fla_cache,
+                    past_key_value=layer_gqa,
+                    use_cache=use_cache,
+                )
+                if use_cache and layer.layer_type in softmax_types and gqa_cache is not None:
+                    gqa_cache[layer_index] = present
+            if (
+                self.block_attn_res is not None
+                and (layer_index + 1) % interval == 0
+            ):
+                hidden_states = self.block_attn_res(
+                    hidden_states, [*attn_memories, hidden_states]
+                )
+                attn_memories.append(hidden_states)
+        hidden_states = self.final_norm(hidden_states)
+
+        main_loss = None
+        mtp_loss = None
+        z_loss = None
+        total_loss = None
+        if labels is not None:
+            main_target_offset = 0 if labels_are_shifted else 1
+            main_objective, main_loss, z_loss = self._causal_loss(
+                hidden_states,
+                labels,
+                target_offset=main_target_offset,
+                z_loss_coefficient=z_loss_coefficient,
+            )
+            if z_loss_coefficient == 0.0:
+                z_loss = None
+            if self.mtp_adapters and include_mtp_loss:
+                auxiliary_losses = [
+                    self._causal_loss(
+                        hidden_states,
+                        labels,
+                        target_offset=(
+                            head_index + 1
+                            if labels_are_shifted
+                            else head_index + 2
+                        ),
+                        adapter=adapter,
+                    )[0]
+                    for head_index, adapter in enumerate(self.mtp_adapters)
+                ]
+                mtp_loss = torch.stack(auxiliary_losses).mean()
+                total_loss = main_objective + effective_mtp_weight * mtp_loss
+            else:
+                total_loss = main_objective
+
+        output_states = (
+            hidden_states
+            if logits_to_keep is None
+            else hidden_states[:, -logits_to_keep:, :]
+        )
+        logits = self.project_to_vocabulary(output_states) if return_logits else None
+        return TinyGDNOutput(
+            loss=total_loss,
+            logits=logits,
+            main_loss=main_loss,
+            mtp_loss=mtp_loss,
+            z_loss=z_loss,
+            hidden_states=hidden_states if return_hidden_states else None,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+    def parameter_report(self) -> dict[str, int]:
+        total = sum(parameter.numel() for parameter in self.parameters())
+        mtp = sum(parameter.numel() for parameter in self.mtp_adapters.parameters())
+        embeddings = self.embed_tokens.weight.numel()
+        return {
+            "deployable_core": total - mtp,
+            "training_total": total,
+            "embedding": embeddings,
+            "mtp_auxiliary": mtp,
+            "non_embedding_core": total - mtp - embeddings,
+        }
+
+    def save_checkpoint(self, output_dir: Path) -> None:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        self.config.save_json(output_dir / "config.json")
+        save_model(self, output_dir / "model.safetensors")
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        checkpoint_dir: Path,
+        *,
+        device: str | torch.device = "cpu",
+        dtype: torch.dtype | None = None,
+    ) -> TinyGDNForCausalLM:
+        config = TinyGDNConfig.from_json(checkpoint_dir / "config.json")
+        model = cls(config).to(device=device, dtype=dtype)
+        load_model(model, checkpoint_dir / "model.safetensors", device=str(device))
+        return model
+
+    def extra_repr(self) -> str:
+        report = self.parameter_report()
+        return (
+            f"core_parameters={report['deployable_core']:,}, "
+            f"training_parameters={report['training_total']:,}"
+        )
+
+    def get_architecture_metadata(self) -> dict[str, Any]:
+        return {
+            "architecture": self.config.architecture,
+            "layer_types": list(self.config.layer_types),
+            "effective_num_layers": self.config.effective_num_layers,
+            "shared_layer_indices": list(self.config.shared_layer_indices),
+            "parameter_report": self.parameter_report(),
+            "features": [
+                "hybrid linear + softmax token mixing",
+                "Kimi Delta Attention or Gated DeltaNet-2",
+                "gated MLA or gated GQA",
+                "optional block Attention Residuals",
+                "SwiGLU or SiTU-GLU",
+                "QK normalization",
+                "zero-centered RMSNorm",
+                "tied input-output embeddings",
+                "optional multi-token prediction auxiliaries",
+            ],
+        }

tiny_gdn/ngram.py

@@ -0,0 +1,291 @@
+"""Hashed n-gram addressing shared by Senryu-2's model, its host table store, and serving.
+
+Each table row is addressed with integer ops that give the same result on CPU
+and CUDA: canonical token ids (NFKC-normalized, lowercased, stripped surface
+forms) are multiplied by odd 31-bit constants, XOR-combined, mixed, and reduced
+modulo a distinct prime per table. Token ids stay below 2**16, so every
+product stays below 2**47 and int64 never overflows.
+
+Tables live in one concatenated ``[total_rows, embedding_dim]`` tensor; table
+``t`` owns rows ``[offset_t, offset_t + rows_t)``. The prime row counts, the
+multipliers, and the canonical map are saved beside every table file and
+checked on load, so a table can never be read with a different address map.
+"""
+
+from __future__ import annotations
+
+import unicodedata
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+from tokenizers import Tokenizer
+
+from tiny_gdn.config import TinyGDNConfig
+
+NGRAM_TABLES_FILE = "ngram_tables.safetensors"
+NGRAM_TABLES_EMA_FILE = "ngram_tables_ema.safetensors"
+NGRAM_OPTIMIZER_FILE = "ngram_optimizer.safetensors"
+
+PRIME_STRIDE = 4099
+HASH_SEED = 0x6B696D6F
+MIX_SHIFT = 29
+LOAD_CHUNK_ROWS = 1 << 20
+_MASK64 = (1 << 64) - 1
+_MILLER_RABIN_BASES = (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37)
+
+
+def _splitmix64(state: int) -> tuple[int, int]:
+    state = (state + 0x9E3779B97F4A7C15) & _MASK64
+    mixed = ((state ^ (state >> 30)) * 0xBF58476D1CE4E5B9) & _MASK64
+    mixed = ((mixed ^ (mixed >> 27)) * 0x94D049BB133111EB) & _MASK64
+    return state, mixed ^ (mixed >> 31)
+
+
+def _is_prime(value: int) -> bool:
+    if value < 2:
+        return False
+    for base in _MILLER_RABIN_BASES:
+        if value % base == 0:
+            return value == base
+    odd, twos = value - 1, 0
+    while odd % 2 == 0:
+        odd //= 2
+        twos += 1
+    for base in _MILLER_RABIN_BASES:
+        witness = pow(base, odd, value)
+        if witness in (1, value - 1):
+            continue
+        for _ in range(twos - 1):
+            witness = witness * witness % value
+            if witness == value - 1:
+                break
+        else:
+            return False
+    return True
+
+
+def _next_prime(value: int) -> int:
+    while not _is_prime(value):
+        value += 1
+    return value
+
+
+@dataclass(frozen=True)
+class NGramLayout:
+    """Where every n-gram of every order lands in the concatenated tables."""
+
+    orders: tuple[int, ...]
+    heads_per_order: int
+    table_rows: tuple[int, ...]
+    table_offsets: tuple[int, ...]
+    multipliers: tuple[tuple[int, ...], ...]
+    embedding_dim: int
+    pad_token_id: int
+
+    @classmethod
+    def from_config(cls, config: TinyGDNConfig) -> NGramLayout:
+        if not config.ngram_orders:
+            raise ValueError("config has no n-gram memory")
+        rows: list[int] = []
+        candidate = config.ngram_table_rows
+        for _ in range(config.ngram_table_count):
+            candidate = _next_prime(candidate + PRIME_STRIDE)
+            rows.append(candidate)
+        offsets = tuple(int(value) for value in np.cumsum([0, *rows[:-1]]))
+        state = HASH_SEED
+        multipliers: list[tuple[int, ...]] = []
+        max_order = max(config.ngram_orders)
+        for _ in range(config.ngram_table_count):
+            table: list[int] = []
+            for _ in range(max_order):
+                state, value = _splitmix64(state)
+                table.append((value >> 33) | (1 << 30) | 1)
+            multipliers.append(tuple(table))
+        return cls(
+            orders=tuple(config.ngram_orders),
+            heads_per_order=config.ngram_heads_per_order,
+            table_rows=tuple(rows),
+            table_offsets=offsets,
+            multipliers=tuple(multipliers),
+            embedding_dim=config.ngram_embedding_dim,
+            pad_token_id=config.bos_token_id,
+        )
+
+    @property
+    def table_count(self) -> int:
+        return len(self.table_rows)
+
+    @property
+    def total_rows(self) -> int:
+        return self.table_offsets[-1] + self.table_rows[-1]
+
+    @property
+    def max_order(self) -> int:
+        return max(self.orders)
+
+    @property
+    def memory_dim(self) -> int:
+        return self.table_count * self.embedding_dim
+
+    @property
+    def table_orders(self) -> tuple[int, ...]:
+        return tuple(order for order in self.orders for _ in range(self.heads_per_order))
+
+    def metadata_tensors(self) -> dict[str, torch.Tensor]:
+        return {
+            "table_rows": torch.tensor(self.table_rows, dtype=torch.int64),
+            "hash_multipliers": torch.tensor(self.multipliers, dtype=torch.int64),
+        }
+
+    def check_metadata(self, tensors: dict[str, torch.Tensor], source: Path) -> None:
+        for name, expected in self.metadata_tensors().items():
+            found = tensors.get(name)
+            if found is None or not torch.equal(found.cpu(), expected):
+                raise ValueError(f"{source} was written for a different n-gram layout ({name})")
+
+
+def canonical_token_ids(tokenizer: Tokenizer, vocab_size: int) -> torch.Tensor:
+    """Map every token id to the smallest id sharing its normalized surface form.
+
+    Normalization is NFKC, then strip, then lowercase. Special tokens keep
+    their own ids, whitespace-only tokens are only NFKC-merged, and tokens
+    that decode to partial UTF-8 stay distinct.
+    """
+
+    special = set(tokenizer.get_added_tokens_decoder())
+    canonical = np.arange(vocab_size, dtype=np.int64)
+    first_by_key: dict[tuple[str, str], int] = {}
+    for token_id in range(min(vocab_size, tokenizer.get_vocab_size())):
+        if token_id in special:
+            continue
+        text = tokenizer.decode([token_id], skip_special_tokens=False)
+        if "\ufffd" in text:
+            key = ("bytes", tokenizer.id_to_token(token_id) or str(token_id))
+        else:
+            normalized = unicodedata.normalize("NFKC", text)
+            folded = normalized.strip().lower()
+            key = ("text", folded) if folded else ("space", normalized)
+        canonical[token_id] = first_by_key.setdefault(key, token_id)
+    return torch.from_numpy(canonical)
+
+
+def document_distance(tokens: torch.Tensor, bos_token_id: int, row_starts: torch.Tensor | None = None) -> torch.Tensor:
+    """Offset of every token in a flat ``[N]`` stream from its document's first token.
+
+    Documents start at position 0, at every BOS token, and at ``row_starts``.
+    """
+
+    starts = tokens == bos_token_id
+    starts[0] = True
+    if row_starts is not None:
+        starts[row_starts] = True
+    positions = torch.arange(tokens.numel(), device=tokens.device)
+    first = torch.where(starts, positions, torch.zeros_like(positions)).cummax(dim=0).values
+    return positions - first
+
+
+def ngram_row_indices(
+    canonical: torch.Tensor,
+    layout: NGramLayout,
+    *,
+    distance: torch.Tensor | None = None,
+    history: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Global table rows ``[..., L, tables]`` for the n-grams ending at each token.
+
+    ``canonical`` holds int64 canonical ids ``[..., L]``. Tokens before a
+    document start read as the pad id: with ``distance`` (same shape) a
+    document starts wherever the distance is zero; with ``history`` (the
+    ``max_order - 1`` canonical ids before position 0) the sequence continues
+    a cached one; with neither, position 0 starts the only document.
+    """
+
+    span = layout.max_order
+    length = canonical.shape[-1]
+    pad = layout.pad_token_id
+    if history is not None:
+        if history.shape[-1] != span - 1:
+            raise ValueError(f"history must hold {span - 1} tokens")
+        extended = torch.cat([history, canonical], dim=-1)
+        shifted = [extended[..., span - 1 - back : span - 1 - back + length] for back in range(span)]
+    else:
+        shifted = [canonical]
+        for back in range(1, span):
+            if back >= length:
+                previous = torch.full_like(canonical, pad)
+            else:
+                padding = canonical.new_full((*canonical.shape[:-1], back), pad)
+                previous = torch.cat([padding, canonical[..., :-back]], dim=-1)
+            if distance is not None:
+                previous = torch.where(distance >= back, previous, pad)
+            shifted.append(previous)
+    columns = []
+    for table, order in enumerate(layout.table_orders):
+        multipliers = layout.multipliers[table]
+        mixed = shifted[0] * multipliers[0]
+        for back in range(1, order):
+            mixed = mixed ^ (shifted[back] * multipliers[back])
+        mixed = mixed ^ (mixed >> MIX_SHIFT)
+        columns.append(torch.remainder(mixed, layout.table_rows[table]) + layout.table_offsets[table])
+    return torch.stack(columns, dim=-1)
+
+
+def save_table_file(path: Path, tables: torch.Tensor, layout: NGramLayout, canonical: torch.Tensor) -> None:
+    if tables.shape != (layout.total_rows, layout.embedding_dim):
+        raise ValueError(f"tables have shape {tuple(tables.shape)}, layout needs {(layout.total_rows, layout.embedding_dim)}")
+    save_file(
+        {
+            "tables": tables.contiguous(),
+            "canonical_ids": canonical.to(dtype=torch.int64, device="cpu").contiguous(),
+            **layout.metadata_tensors(),
+        },
+        str(path),
+    )
+
+
+def read_table_metadata(path: Path, layout: NGramLayout) -> torch.Tensor:
+    """Check ``path`` against ``layout`` and return its canonical token map."""
+
+    with safe_open(str(path), framework="pt", device="cpu") as handle:
+        metadata = {name: handle.get_tensor(name) for name in ("table_rows", "hash_multipliers")}
+        layout.check_metadata(metadata, path)
+        # get_tensor returns a view of a mapping of the whole file, which a kept
+        # view would hold open; FUSE mounts such as ntfs-3g then cannot delete it.
+        return handle.get_tensor("canonical_ids").clone()
+
+
+def copy_rows_from_file(path: Path, name: str, destination: torch.Tensor) -> None:
+    """Copy tensor ``name`` of a safetensors file into ``destination`` in row chunks."""
+
+    with safe_open(str(path), framework="pt", device="cpu") as handle:
+        source = handle.get_slice(name)
+        shape = tuple(source.get_shape())
+        if shape != tuple(destination.shape):
+            raise ValueError(f"{path} holds {name} of shape {shape}, expected {tuple(destination.shape)}")
+        for start in range(0, shape[0], LOAD_CHUNK_ROWS):
+            stop = min(start + LOAD_CHUNK_ROWS, shape[0])
+            destination[start:stop].copy_(source[start:stop])
+
+
+def load_table_file(
+    path: Path,
+    layout: NGramLayout,
+    *,
+    device: torch.device | str = "cpu",
+    dtype: torch.dtype = torch.float32,
+    out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """``(tables, canonical_ids)`` from ``path``, copied in row chunks to bound host memory."""
+
+    canonical = read_table_metadata(path, layout)
+    tables = (
+        torch.empty((layout.total_rows, layout.embedding_dim), dtype=dtype, device=device)
+        if out is None
+        else out
+    )
+    copy_rows_from_file(path, "tables", tables)
+    return tables, canonical

tiny_gdn/nn_common.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.attention import SDPBackend
+
+try:
+    from fla.modules.layernorm import rms_norm as fused_rms_norm
+except ImportError:
+    fused_rms_norm = None
+
+
+class RMSNorm(nn.Module):
+    """Zero-centered RMSNorm as used by Qwen3-Next.
+
+    On CUDA this runs FLA's single-pass Triton kernel with the ``1 + weight``
+    scale kept in FP32, which matches the eager FP32 math to rounding.
+    """
+
+    def __init__(self, hidden_size: int, eps: float) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(hidden_size))
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        scale = 1.0 + self.weight.float()
+        if fused_rms_norm is not None and hidden_states.is_cuda:
+            return fused_rms_norm(hidden_states, scale, None, eps=self.eps)
+        input_dtype = hidden_states.dtype
+        normalized = hidden_states.float()
+        normalized = normalized * torch.rsqrt(
+            normalized.square().mean(dim=-1, keepdim=True) + self.eps
+        )
+        return (normalized * scale).to(dtype=input_dtype)
+
+
+def linear_cross_entropy_reference(
+    states: torch.Tensor,
+    weight: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    z_loss_coefficient: float = 0.0,
+    ignore_index: int = -100,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Eager ``(objective, cross_entropy, z_loss)`` over materialized logits."""
+
+    logits = F.linear(states, weight)
+    cross_entropy = F.cross_entropy(logits, targets, ignore_index=ignore_index)
+    reduced_dtype = torch.promote_types(logits.dtype, torch.float32)
+    if z_loss_coefficient > 0.0:
+        # Keep logits in their training dtype for logsumexp. Casting the full
+        # [tokens, vocab] tensor to fp32 materializes it again in the autograd
+        # graph; upcast only the reduced [tokens] partition.
+        log_partition = torch.logsumexp(logits, dim=-1).to(reduced_dtype)
+        z_loss = log_partition.square()[targets.ne(ignore_index)].mean()
+    else:
+        z_loss = cross_entropy.new_zeros((), dtype=reduced_dtype)
+    objective = cross_entropy + z_loss_coefficient * z_loss
+    return objective, cross_entropy.detach(), z_loss.detach()
+
+
+def linear_token_logprobs_reference(
+    states: torch.Tensor,
+    weight: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    ignore_index: int = -100,
+) -> torch.Tensor:
+    """Eager fp32 ``log_softmax(states @ weight.T)[row, target]``; ignored rows are 0."""
+
+    valid = targets.ne(ignore_index)
+    logprobs = F.linear(states, weight).float().log_softmax(dim=-1)
+    picked = logprobs.gather(1, targets.masked_fill(~valid, 0)[:, None]).squeeze(1)
+    return picked.masked_fill(~valid, 0.0)
+
+
+def softcap(values: torch.Tensor, cap: float) -> torch.Tensor:
+    return cap * torch.tanh(values / cap)
+
+
+def situ_glu_reference(gate_up: torch.Tensor, gate_cap: float, up_cap: float) -> torch.Tensor:
+    """``softcap(gate) * sigmoid(gate) * softcap(up)`` over ``[..., gate | up]``."""
+
+    gate, up = gate_up.chunk(2, dim=-1)
+    return softcap(gate, gate_cap) * torch.sigmoid(gate) * softcap(up, up_cap)
+
+
+def attention_sdpa_backends(device: torch.device) -> list[SDPBackend]:
+    if device.type != "cuda":
+        return [SDPBackend.MATH]
+    return [
+        SDPBackend.CUDNN_ATTENTION,
+        SDPBackend.FLASH_ATTENTION,
+        SDPBackend.EFFICIENT_ATTENTION,
+    ]

tiny_gdn/parallax_ops.py

@@ -0,0 +1,394 @@
+"""Parallax attention on fla's Triton kernels, with launch settings that fit the GPU.
+
+fla tiles every Parallax kernel 64 rows deep with two pipeline stages, sized
+for Hopper's 228 KB of shared memory. Its backward key/value kernel then needs
+about 105 KB per block, more than consumer Blackwell (sm_120) and Ampere allow
+(99 KB), so it fails to launch there. The tile is a launch constant, so this
+module launches fla's own kernels with the largest tile that fits and keeps
+fla's math unchanged.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import torch
+import triton
+import triton.language as tl
+from einops import reduce
+from fla.ops.parallax.decode import parallax_decode_kernel, parallax_decode_one_step
+from fla.ops.parallax.parallel import (
+    parallel_parallax_bwd_kernel_dkv,
+    parallel_parallax_bwd_kernel_dqr,
+    parallel_parallax_bwd_kernel_preprocess,
+    parallel_parallax_fwd_kernel,
+)
+from fla.ops.utils import prepare_chunk_indices
+from fla.ops.utils.op import exp2
+from fla.utils import autocast_custom_bwd, autocast_custom_fwd, contiguous
+from triton.runtime.errors import OutOfResources
+
+LaunchConfig = tuple[int, int, int]
+# (tile rows, warps, pipeline stages), fastest first. Measured on an RTX PRO
+# 6000 (sm_120, 99 KB per block): a 32-row tile with 4 warps beats every other
+# setting that launches, forward and backward.
+LARGE_SHARED_MEMORY_CONFIGS: tuple[LaunchConfig, ...] = ((64, 8, 2), (32, 4, 2), (32, 4, 1), (16, 4, 1))
+SMALL_SHARED_MEMORY_CONFIGS: tuple[LaunchConfig, ...] = ((32, 4, 2), (32, 4, 1), (16, 4, 1))
+LARGE_SHARED_MEMORY_BYTES = 160 * 1024
+
+_selected: dict[tuple[str, int, int], LaunchConfig] = {}
+
+
+def _device_index(device: torch.device) -> int:
+    return torch.cuda.current_device() if device.index is None else device.index
+
+
+def _candidates(device: torch.device) -> tuple[LaunchConfig, ...]:
+    properties = torch.cuda.get_device_properties(_device_index(device))
+    if properties.shared_memory_per_block_optin >= LARGE_SHARED_MEMORY_BYTES:
+        return LARGE_SHARED_MEMORY_CONFIGS
+    return SMALL_SHARED_MEMORY_CONFIGS
+
+
+def _launch(name: str, device: torch.device, head_dim: int, launch: Callable[[int, int, int], Any]) -> Any:
+    """Run ``launch`` with the first config that fits; remember it per kernel and GPU."""
+
+    key = (name, _device_index(device), head_dim)
+    chosen = _selected.get(key)
+    candidates = (chosen,) if chosen is not None else _candidates(device)
+    last_error: OutOfResources | None = None
+    for config in candidates:
+        try:
+            result = launch(*config)
+        except OutOfResources as error:
+            last_error = error
+            continue
+        _selected[key] = config
+        return result
+    assert last_error is not None
+    raise last_error
+
+
+def selected_launch_configs() -> dict[tuple[str, int, int], LaunchConfig]:
+    return dict(_selected)
+
+
+def _chunk_indices(cu_seqlens: torch.Tensor | None, tile: int) -> torch.Tensor | None:
+    return None if cu_seqlens is None else prepare_chunk_indices(cu_seqlens, tile)
+
+
+def _grid_rows(length: int, chunk_indices: torch.Tensor | None, tile: int) -> int:
+    return triton.cdiv(length, tile) if chunk_indices is None else len(chunk_indices)
+
+
+def _forward(
+    q: torch.Tensor,
+    r: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    scale: float,
+    cu_seqlens: torch.Tensor | None,
+    window_size_left: int,
+) -> tuple[torch.Tensor, ...]:
+    batch, length, query_heads, head_dim = q.shape
+    kv_heads = k.shape[2]
+    o = torch.empty_like(q)
+    barv = torch.empty_like(q)
+    d1 = torch.empty((batch, length, query_heads), device=q.device, dtype=torch.float32)
+    bart = torch.empty_like(d1)
+    m = torch.empty_like(d1)
+
+    def launch(tile: int, warps: int, stages: int) -> None:
+        chunk_indices = _chunk_indices(cu_seqlens, tile)
+        grid = (_grid_rows(length, chunk_indices, tile), batch * query_heads)
+        parallel_parallax_fwd_kernel[grid](
+            q, r, k, v, o, barv, d1, bart, m,
+            scale, cu_seqlens, chunk_indices, length,
+            HQ=query_heads, H=kv_heads, G=query_heads // kv_heads, K=head_dim,
+            BK=triton.next_power_of_2(head_dim),
+            WINDOW_SIZE_LEFT=window_size_left, BT=tile, BS=tile,
+            num_warps=warps, num_stages=stages,
+        )
+
+    _launch("forward", q.device, head_dim, launch)
+    return o, barv, d1, bart, m
+
+
+def _backward(
+    saved: tuple[torch.Tensor, ...],
+    grad_o: torch.Tensor,
+    scale: float,
+    cu_seqlens: torch.Tensor | None,
+    window_size_left: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, r, k, v, o, barv, d1, bart, m = saved
+    batch, length, query_heads, head_dim = q.shape
+    kv_heads = k.shape[2]
+    groups = query_heads // kv_heads
+    grad_q = torch.empty_like(q)
+    grad_r = torch.empty_like(r)
+    # dK/dV are written per query head, then summed back onto the KV heads.
+    grad_k_heads = torch.empty((batch, length, query_heads, head_dim), device=q.device, dtype=q.dtype)
+    grad_v_heads = torch.empty_like(grad_k_heads)
+    delta_t = torch.empty((batch, length, query_heads), device=q.device, dtype=torch.float32)
+    delta_b = torch.empty_like(delta_t)
+    block_k = triton.next_power_of_2(head_dim)
+
+    def launch(tile: int, warps: int, stages: int) -> None:
+        chunk_indices = _chunk_indices(cu_seqlens, tile)
+        grid = (_grid_rows(length, chunk_indices, tile), batch * query_heads)
+        parallel_parallax_bwd_kernel_preprocess[grid](
+            grad_o, o, barv, delta_t, delta_b,
+            cu_seqlens, chunk_indices, length,
+            HQ=query_heads, K=head_dim, BK=block_k, BT=tile,
+            num_warps=4, num_stages=2,
+        )
+        parallel_parallax_bwd_kernel_dqr[grid](
+            q, r, k, v, d1, bart, m, delta_t, delta_b, grad_o, grad_q, grad_r,
+            scale, cu_seqlens, chunk_indices, length,
+            HQ=query_heads, H=kv_heads, G=groups, K=head_dim, BK=block_k,
+            WINDOW_SIZE_LEFT=window_size_left, BT=tile, BS=tile,
+            num_warps=warps, num_stages=stages,
+        )
+        parallel_parallax_bwd_kernel_dkv[grid](
+            q, r, k, v, d1, bart, m, delta_t, delta_b, grad_o, grad_k_heads, grad_v_heads,
+            scale, cu_seqlens, chunk_indices, length,
+            HQ=query_heads, H=kv_heads, G=groups, K=head_dim, BK=block_k,
+            WINDOW_SIZE_LEFT=window_size_left, BT=tile, BS=tile,
+            num_warps=warps, num_stages=stages,
+        )
+
+    _launch("backward", q.device, head_dim, launch)
+    if groups == 1:
+        return grad_q, grad_r, grad_k_heads, grad_v_heads
+    grad_k = reduce(grad_k_heads, "b t (h g) d -> b t h d", g=groups, reduction="sum")
+    grad_v = reduce(grad_v_heads, "b t (h g) d -> b t h d", g=groups, reduction="sum")
+    return grad_q, grad_r, grad_k, grad_v
+
+
+class _ParallaxFunction(torch.autograd.Function):
+    @staticmethod
+    @contiguous
+    @autocast_custom_fwd
+    def forward(ctx, q, r, k, v, scale, window_size_left, cu_seqlens):
+        o, barv, d1, bart, m = _forward(q, r, k, v, scale, cu_seqlens, window_size_left)
+        ctx.save_for_backward(q, r, k, v, o, barv, d1, bart, m)
+        ctx.scale = scale
+        ctx.window_size_left = window_size_left
+        ctx.cu_seqlens = cu_seqlens
+        return o
+
+    @staticmethod
+    @contiguous
+    @autocast_custom_bwd
+    def backward(ctx, grad_o):
+        saved = ctx.saved_tensors
+        grad_q, grad_r, grad_k, grad_v = _backward(
+            saved, grad_o, ctx.scale, ctx.cu_seqlens, ctx.window_size_left
+        )
+        q, r, k, v = saved[:4]
+        return grad_q.to(q), grad_r.to(r), grad_k.to(k), grad_v.to(v), None, None, None
+
+
+def _validate(q: torch.Tensor, r: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None:
+    if q.dtype not in (torch.bfloat16, torch.float16):
+        raise TypeError(f"Parallax needs bf16 or fp16 inputs, got {q.dtype}")
+    if r.shape != q.shape:
+        raise ValueError("r must have the shape of q")
+    if k.shape != v.shape or k.shape[0] != q.shape[0] or k.shape[-1] != q.shape[-1]:
+        raise ValueError("k and v must share q's batch and head_dim")
+    if q.shape[2] % k.shape[2]:
+        raise ValueError("query heads must be a multiple of key/value heads")
+
+
+def parallax_attention(
+    q: torch.Tensor,
+    r: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    scale: float | None = None,
+    window_size: int | None = None,
+    cu_seqlens: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Causal Parallax over ``[batch, length, heads, head_dim]`` inputs with autograd.
+
+    ``r`` is not scaled; ``scale`` applies to ``q @ kᵀ`` only (default
+    ``head_dim ** -0.5``). With ``cu_seqlens`` the batch must be 1 and each
+    packed sequence attends only within itself.
+    """
+
+    _validate(q, r, k, v)
+    if k.shape[1] != q.shape[1]:
+        raise ValueError("training Parallax needs one key per query")
+    if cu_seqlens is not None and q.shape[0] != 1:
+        raise ValueError("packed sequences (cu_seqlens) need batch size 1")
+    if scale is None:
+        scale = q.shape[-1] ** -0.5
+    window_size_left = -1 if window_size is None else window_size
+    return _ParallaxFunction.apply(q, r, k, v, float(scale), window_size_left, cu_seqlens)
+
+
+@torch.no_grad()
+def parallax_decode_cached(
+    q: torch.Tensor,
+    r: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    scale: float | None = None,
+    window_size: int | None = None,
+    cache_start: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Forward-only Parallax for ``Sq`` new queries over ``Skv >= Sq`` cached keys.
+
+    Query ``i`` sits at absolute position ``Skv - Sq + i``. One query uses
+    fla's vector decode kernel; longer chunks use the tiled prefill kernel.
+    """
+
+    _validate(q, r, k, v)
+    batch, queries, query_heads, head_dim = q.shape
+    keys, kv_heads = k.shape[1], k.shape[2]
+    if keys < queries:
+        raise ValueError("the cache must hold at least as many keys as new queries")
+    if scale is None:
+        scale = head_dim ** -0.5
+    if queries == 1:
+        return parallax_decode_one_step(q, r, k, v, scale=scale, window_size=window_size, cache_start=cache_start)
+    window_size_left = -1 if window_size is None else window_size
+    q, r, k, v = (tensor.contiguous() for tensor in (q, r, k, v))
+    if cache_start is not None:
+        cache_start = cache_start.to(device=q.device, dtype=torch.int32).contiguous()
+    o = torch.empty_like(q)
+
+    def launch(tile: int, warps: int, stages: int) -> None:
+        grid = (triton.cdiv(queries, tile), batch * query_heads)
+        parallax_decode_kernel[grid](
+            q, r, k, v, o, float(scale), cache_start, queries, keys,
+            HQ=query_heads, H=kv_heads, G=query_heads // kv_heads, K=head_dim,
+            BK=triton.next_power_of_2(head_dim),
+            WINDOW_SIZE_LEFT=window_size_left,
+            USE_CACHE_START=cache_start is not None,
+            BT=tile, BS=tile,
+            num_warps=warps, num_stages=stages,
+        )
+
+    _launch("decode", q.device, head_dim, launch)
+    return o
+
+
+@triton.jit(do_not_specialize=["Skv"])
+def _parallax_graph_decode_kernel(
+    q,
+    r,
+    k,
+    v,
+    o,
+    scale,
+    cache_end,
+    Skv,
+    HQ: tl.constexpr,
+    H: tl.constexpr,
+    G: tl.constexpr,
+    K: tl.constexpr,
+    BK: tl.constexpr,
+    BS: tl.constexpr,
+):
+    """One query per row over a fixed ``Skv`` cache, stopping at that row's ``cache_end``.
+
+    ``Skv`` stays the buffer length so a CUDA graph can replay the launch while
+    each row's live prefix, read from ``cache_end``, changes between replays.
+    """
+
+    i_bh = tl.program_id(0)
+    i_b, i_hq = i_bh // HQ, i_bh % HQ
+    i_h = i_hq // G
+    RCP_LN2: tl.constexpr = 1.4426950216
+    kv_hi = tl.load(cache_end + i_b).to(tl.int32)
+    kv_hi = tl.minimum(tl.maximum(kv_hi, 0), Skv)
+
+    p_q = tl.make_block_ptr(q + i_bh * K, (K,), (1,), (0,), (BK,), (0,))
+    p_r = tl.make_block_ptr(r + i_bh * K, (K,), (1,), (0,), (BK,), (0,))
+    p_o = tl.make_block_ptr(o + i_bh * K, (K,), (1,), (0,), (BK,), (0,))
+    b_q = tl.load(p_q, boundary_check=(0,), padding_option="zero").to(tl.float32)
+    b_r = tl.load(p_r, boundary_check=(0,), padding_option="zero").to(tl.float32)
+    scale_log2 = scale * RCP_LN2
+    m = tl.full((1,), -float("inf"), dtype=tl.float32)
+    d1 = tl.zeros((1,), dtype=tl.float32)
+    d2 = tl.zeros((1,), dtype=tl.float32)
+    o1 = tl.zeros((BK,), dtype=tl.float32)
+    o2 = tl.zeros((BK,), dtype=tl.float32)
+
+    p_k = tl.make_block_ptr(k + (i_b * Skv * H + i_h) * K, (Skv, K), (H * K, 1), (0, 0), (BS, BK), (1, 0))
+    p_v = tl.make_block_ptr(v + (i_b * Skv * H + i_h) * K, (Skv, K), (H * K, 1), (0, 0), (BS, BK), (1, 0))
+    for i_s in range(0, tl.cdiv(kv_hi, BS) * BS, BS):
+        col = i_s + tl.arange(0, BS)
+        mask = col < kv_hi
+        b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero")
+        b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero")
+        s1 = tl.sum(b_q[None, :] * b_k, axis=1) * scale_log2
+        s2 = tl.sum(b_r[None, :] * b_k, axis=1)
+        s1 = tl.where(mask, s1, -float("inf"))
+        m_new = tl.maximum(m, tl.max(s1))
+        m_safe = tl.where(m_new == -float("inf"), 0.0, m_new)
+        alpha = exp2(m - m_safe)
+        p1 = exp2(s1 - m_safe)
+        p2 = p1 * s2
+        d1 = d1 * alpha + tl.sum(p1)
+        d2 = d2 * alpha + tl.sum(p2)
+        o1 = o1 * alpha + tl.sum(p1[:, None] * b_v, axis=0)
+        o2 = o2 * alpha + tl.sum(p2[:, None] * b_v, axis=0)
+        m = m_new
+        p_k = tl.advance(p_k, (BS, 0))
+        p_v = tl.advance(p_v, (BS, 0))
+
+    inv_d1 = tl.where(d1 > 0.0, 1.0 / d1, 0.0)
+    out = o1 * inv_d1 * (1.0 + d2 * inv_d1) - o2 * inv_d1
+    tl.store(p_o, out.to(p_o.dtype.element_ty), boundary_check=(0,))
+
+
+def parallax_graph_decode(
+    q: torch.Tensor,
+    r: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cache_end: torch.Tensor,
+    *,
+    scale: float,
+) -> torch.Tensor:
+    """Single-token Parallax over ``k, v`` of ``[rows, capacity, heads, dim]``.
+
+    Row ``i`` attends to keys ``[0, cache_end[i])``. ``capacity`` is fixed, so
+    the launch can be captured in a CUDA graph and replayed as rows grow.
+    """
+
+    _validate(q, r, k, v)
+    batch, queries, query_heads, head_dim = q.shape
+    if queries != 1:
+        raise ValueError(f"graph decode steps one query, got {queries}")
+    if cache_end.shape != (batch,):
+        raise ValueError("cache_end must hold one length per row")
+    kv_heads = k.shape[2]
+    q, r, k, v = (tensor.contiguous() for tensor in (q, r, k, v))
+    cache_end = cache_end.to(device=q.device, dtype=torch.int32).contiguous()
+    out = torch.empty_like(q)
+    _parallax_graph_decode_kernel[(batch * query_heads,)](
+        q,
+        r,
+        k,
+        v,
+        out,
+        float(scale),
+        cache_end,
+        k.shape[1],
+        HQ=query_heads,
+        H=kv_heads,
+        G=query_heads // kv_heads,
+        K=head_dim,
+        BK=triton.next_power_of_2(head_dim),
+        BS=128,
+        num_warps=4,
+        num_stages=2,
+    )
+    return out

tiny_gdn/senryu2.py

@@ -0,0 +1,982 @@
+"""Senryu-2: a Gated DeltaNet-2 + Parallax hybrid on a 4-branch gated residual stream.
+
+Every sublayer reads its input from the residual branches through Qwen's
+Gated Residual (per-branch RMSNorm, elementwise sigmoid gate from a low-rank
+bottleneck over all branches) and writes its output back to every branch with
+one data-dependent scalar per branch. A hashed n-gram memory, whose tables live
+in host memory during training, is injected before block ``ngram_layer_index``
+(docs/architecture.md, "Senryu-2").
+
+Training and scoring run packed: the batch is flattened into one row and split
+into documents at every BOS token, so GDN-2, Parallax, the n-gram hashes and
+the n-gram convolution never look across a document boundary.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from fla.layers.gdn2 import GatedDeltaNet2
+from fla.models.utils import Cache as FlaCache
+from fla.modules.rotary import rotary_embedding
+from safetensors.torch import save_model
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from tiny_gdn.config import TinyGDNConfig
+from tiny_gdn.gated_residual_ops import fused_gated_residual_read, fused_gated_residual_write
+from tiny_gdn.haiku_layers import SiTUGLU
+from tiny_gdn.model import TinyGDNOutput, linear_cross_entropy
+from tiny_gdn.ngram import (
+    NGRAM_TABLES_EMA_FILE,
+    NGRAM_TABLES_FILE,
+    NGramLayout,
+    document_distance,
+    load_table_file,
+    ngram_row_indices,
+)
+from tiny_gdn.nn_common import RMSNorm
+from tiny_gdn.parallax_ops import parallax_attention, parallax_decode_cached
+
+GATED_RESIDUAL_CHUNK_ROWS = 4096
+KV_CACHE_GRANULE = 256
+
+
+def mark_muon_row_blocks(parameter: nn.Parameter, blocks: int) -> None:
+    """Orthogonalize ``parameter`` as ``blocks`` equal row blocks (heads, fused halves)."""
+
+    if parameter.shape[0] % blocks:
+        raise ValueError(f"{blocks} row blocks do not divide {parameter.shape[0]} rows")
+    parameter._muon_row_blocks = blocks
+
+
+def mark_adamw_only(*parameters: nn.Parameter) -> None:
+    """Keep 2D weights that are gates or elongated low-rank maps on AdamW."""
+
+    for parameter in parameters:
+        parameter._adamw_only = True
+
+
+def mark_gain(*parameters: nn.Parameter) -> None:
+    """2D per-branch norm gains and filters: AdamW without weight decay, like 1D gains."""
+
+    mark_adamw_only(*parameters)
+    for parameter in parameters:
+        parameter._no_weight_decay = True
+
+
+# --------------------------------------------------------------------------- layout
+
+
+@dataclass
+class Senryu2Context:
+    """Per-forward state shared by every block (passes through activation checkpointing)."""
+
+    batch: int
+    length: int
+    cu_seqlens: torch.Tensor | None
+    distance: torch.Tensor | None
+    rotary: tuple[torch.Tensor, torch.Tensor]
+    position_offset: int = 0
+    cache: Senryu2Cache | None = None
+
+
+class RotaryTables:
+    """Partial-RoPE cos/sin rows ``[length, rotary_dim / 2]``, built in FP32 and cached."""
+
+    def __init__(self, rotary_dim: int, theta: float) -> None:
+        self.rotary_dim = rotary_dim
+        self.theta = theta
+        self._cached: tuple[torch.Tensor, torch.Tensor] | None = None
+        self._key: tuple[torch.device, torch.dtype, bool] | None = None
+
+    def get(self, length: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
+        key = (device, dtype, torch.is_inference_mode_enabled())
+        cached = self._cached
+        if cached is None or self._key != key or cached[0].shape[0] < length:
+            rows = max(length, 4096 if cached is None else 2 * cached[0].shape[0])
+            inverse_frequency = 1.0 / (
+                self.theta
+                ** (torch.arange(0, self.rotary_dim, 2, device=device, dtype=torch.float32) / self.rotary_dim)
+            )
+            angles = torch.outer(torch.arange(rows, device=device, dtype=torch.float32), inverse_frequency)
+            cached = (angles.cos().to(dtype), angles.sin().to(dtype))
+            self._cached = cached
+            self._key = key
+        return cached
+
+
+def parallax_reference(
+    q: torch.Tensor,
+    r: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    scale: float,
+    cu_seqlens: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """FP32 Parallax for queries at the last ``Tq`` of ``Tk`` key positions.
+
+    ``out = O1 / d1 * (1 + d2 / d1) - O2 / d1`` with ``p1 = exp(scale q·k)``,
+    ``p2 = p1 * (r·k)``; ``cu_seqlens`` (batch 1, ``Tq == Tk``) masks across
+    packed documents.
+    """
+
+    batch, queries, query_heads, head_dim = q.shape
+    keys, kv_heads = k.shape[1], k.shape[2]
+    groups = query_heads // kv_heads
+    qf = q.float().reshape(batch, queries, kv_heads, groups, head_dim)
+    rf = r.float().reshape(batch, queries, kv_heads, groups, head_dim)
+    kf, vf = k.float(), v.float()
+    s1 = torch.einsum("bqhgd,bkhd->bhgqk", qf, kf) * scale
+    s2 = torch.einsum("bqhgd,bkhd->bhgqk", rf, kf)
+    query_positions = torch.arange(keys - queries, keys, device=q.device)[:, None]
+    key_positions = torch.arange(keys, device=q.device)[None, :]
+    masked = key_positions > query_positions
+    if cu_seqlens is not None:
+        segment = torch.zeros(keys, dtype=torch.long, device=q.device)
+        segment[cu_seqlens[1:-1].long()] = 1
+        segment = segment.cumsum(0)
+        masked = masked | (segment[keys - queries :, None] != segment[None, :])
+    s1 = s1.masked_fill(masked, float("-inf"))
+    p1 = (s1 - s1.amax(dim=-1, keepdim=True)).exp()
+    d1 = p1.sum(dim=-1)
+    p2 = p1 * s2
+    d2 = p2.sum(dim=-1)
+    o1 = torch.einsum("bhgqk,bkhd->bqhgd", p1, vf)
+    o2 = torch.einsum("bhgqk,bkhd->bqhgd", p2, vf)
+    inverse_d1 = (1.0 / d1).permute(0, 3, 1, 2)[..., None]
+    correction = (d2 / d1).permute(0, 3, 1, 2)[..., None]
+    out = o1 * inverse_d1 * (1.0 + correction) - o2 * inverse_d1
+    return out.reshape(batch, queries, query_heads, head_dim).to(q.dtype)
+
+
+# ------------------------------------------------------------------ gated residual
+
+
+class _GatedResidualRead(torch.autograd.Function):
+    """``x = (1/n) Σ σ(W_u SiLU(W_d vec(R̂) / n))_i ⊙ R̂_i`` and ``s = 2σ(W_w vec(R̂) / n)``.
+
+    Saves only the branches, the per-branch inverse RMS and the bottleneck
+    pre-activation; the normalized branches and gates are recomputed in
+    chunks during backward, so a 32k-token read keeps ~0.4 GB instead of ~4 GB.
+    CUDA runs the fused kernels in ``gated_residual_ops`` instead.
+    """
+
+    @staticmethod
+    def forward(ctx, residual, gamma, down, up, write, eps):
+        rows, branches, hidden_size = residual.shape
+        dtype = down.dtype
+        gain = 1.0 + gamma.float()
+        out = residual.new_empty(rows, hidden_size)
+        rstd = residual.new_empty(rows, branches, dtype=torch.float32)
+        hidden = residual.new_empty(rows, down.shape[0])
+        scales = residual.new_empty(rows, 0 if write is None else branches, dtype=torch.float32)
+        for start in range(0, rows, GATED_RESIDUAL_CHUNK_ROWS):
+            stop = min(start + GATED_RESIDUAL_CHUNK_ROWS, rows)
+            values = residual[start:stop].float()
+            chunk_rstd = torch.rsqrt(values.square().mean(dim=-1, keepdim=True) + eps)
+            hat = values * chunk_rstd * gain
+            flat = hat.to(dtype).flatten(1)
+            chunk_hidden = (flat @ down.t()) / branches
+            activated = F.silu(chunk_hidden.float()).to(dtype)
+            gates = torch.sigmoid((activated @ up.t()).float()).view(stop - start, branches, hidden_size)
+            out[start:stop] = ((gates * hat).sum(dim=1) / branches).to(out.dtype)
+            rstd[start:stop] = chunk_rstd.squeeze(-1)
+            hidden[start:stop] = chunk_hidden
+            if write is not None:
+                scales[start:stop] = 2.0 * torch.sigmoid((flat @ write.t()).float() / branches)
+        ctx.save_for_backward(residual, gamma, down, up, write, rstd, hidden, scales)
+        scales_out = scales.to(residual.dtype, copy=True)
+        if write is None:
+            ctx.mark_non_differentiable(scales_out)
+        return out, scales_out
+
+    @staticmethod
+    def backward(ctx, grad_out, grad_scales):
+        residual, gamma, down, up, write, rstd, hidden, scales = ctx.saved_tensors
+        rows, branches, hidden_size = residual.shape
+        dtype = down.dtype
+        gain = 1.0 + gamma.float()
+        grad_residual = torch.empty(residual.shape, dtype=residual.dtype, device=residual.device)
+        grad_gamma = torch.zeros(branches, hidden_size, dtype=torch.float32, device=residual.device)
+        grad_down = torch.zeros(down.shape, dtype=torch.float32, device=down.device)
+        grad_up = torch.zeros(up.shape, dtype=torch.float32, device=up.device)
+        grad_write = None if write is None else torch.zeros(write.shape, dtype=torch.float32, device=write.device)
+        use_scales = write is not None and grad_scales is not None
+        for start in range(0, rows, GATED_RESIDUAL_CHUNK_ROWS):
+            stop = min(start + GATED_RESIDUAL_CHUNK_ROWS, rows)
+            count = stop - start
+            normalized = residual[start:stop].float() * rstd[start:stop, :, None]
+            hat = normalized * gain
+            flat = hat.to(dtype).flatten(1)
+            pre = hidden[start:stop].float()
+            sigmoid_pre = torch.sigmoid(pre)
+            activated = (pre * sigmoid_pre).to(dtype)
+            gates = torch.sigmoid((activated @ up.t()).float()).view(count, branches, hidden_size)
+            grad_x = grad_out[start:stop].float().unsqueeze(1) / branches
+            grad_hat = gates * grad_x
+            grad_gate_logits = (grad_x * hat * gates * (1.0 - gates)).flatten(1).to(dtype)
+            grad_up += (grad_gate_logits.t() @ activated).float()
+            grad_activated = (grad_gate_logits @ up).float()
+            grad_pre = grad_activated * sigmoid_pre * (1.0 + pre * (1.0 - sigmoid_pre))
+            grad_pre = (grad_pre / branches).to(dtype)
+            grad_down += (grad_pre.t() @ flat).float()
+            grad_flat = (grad_pre @ down).float()
+            if use_scales:
+                scale_values = scales[start:stop]
+                grad_scale_logits = grad_scales[start:stop].float() * scale_values * (1.0 - scale_values / 2.0)
+                grad_scale_logits = (grad_scale_logits / branches).to(dtype)
+                grad_write += (grad_scale_logits.t() @ flat).float()
+                grad_flat += (grad_scale_logits @ write).float()
+            grad_hat = grad_hat + grad_flat.view(count, branches, hidden_size)
+            grad_gamma += (grad_hat * normalized).sum(dim=0)
+            grad_normalized = grad_hat * gain
+            grad_residual[start:stop] = (
+                rstd[start:stop, :, None]
+                * (grad_normalized - normalized * (grad_normalized * normalized).mean(dim=-1, keepdim=True))
+            ).to(grad_residual.dtype)
+        return (
+            grad_residual,
+            grad_gamma.to(gamma.dtype),
+            grad_down.to(down.dtype),
+            grad_up.to(up.dtype),
+            None if grad_write is None else grad_write.to(write.dtype),
+            None,
+        )
+
+
+def gated_residual_read_reference(
+    residual: torch.Tensor,
+    gamma: torch.Tensor,
+    down: torch.Tensor,
+    up: torch.Tensor,
+    write: torch.Tensor | None,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Eager autograd version of :class:`_GatedResidualRead`, for tests."""
+
+    rows, branches, hidden_size = residual.shape
+    gain = 1.0 + gamma.float()
+    values = residual.float()
+    hat = values * torch.rsqrt(values.square().mean(dim=-1, keepdim=True) + eps) * gain
+    flat = hat.to(down.dtype).flatten(1)
+    hidden = (flat @ down.t()) / branches
+    gates = torch.sigmoid((F.silu(hidden.float()).to(up.dtype) @ up.t()).float()).view(rows, branches, hidden_size)
+    out = ((gates * hat).sum(dim=1) / branches).to(residual.dtype)
+    if write is None:
+        return out, residual.new_empty(rows, 0)
+    scales = 2.0 * torch.sigmoid((flat @ write.t()).float() / branches)
+    return out, scales.to(residual.dtype)
+
+
+class GatedResidual(nn.Module):
+    """Qwen's Gated Residual read (and write scales) over ``branches`` residual branches."""
+
+    def __init__(self, branches: int, hidden_size: int, rank: int, eps: float, *, write: bool) -> None:
+        super().__init__()
+        self.branches = branches
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(branches, hidden_size))
+        self.down_proj = nn.Linear(branches * hidden_size, rank, bias=False)
+        self.up_proj = nn.Linear(rank, branches * hidden_size, bias=False)
+        self.write_proj = nn.Linear(branches * hidden_size, branches, bias=False) if write else None
+        mark_gain(self.weight)
+        mark_adamw_only(self.down_proj.weight, self.up_proj.weight)
+        if self.write_proj is not None:
+            mark_adamw_only(self.write_proj.weight)
+
+    def read(self, residual: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """``(x [N, d], scales [N, n])`` from branches ``[N, n, d]``."""
+
+        write = None if self.write_proj is None else self.write_proj.weight
+        read = fused_gated_residual_read if residual.is_cuda else _GatedResidualRead.apply
+        return read(residual, self.weight, self.down_proj.weight, self.up_proj.weight, write, self.eps)
+
+    @staticmethod
+    def write(residual: torch.Tensor, scales: torch.Tensor, update: torch.Tensor) -> torch.Tensor:
+        if residual.is_cuda:
+            return fused_gated_residual_write(residual, scales, update)
+        return residual + scales.unsqueeze(-1) * update.unsqueeze(-2)
+
+
+# ------------------------------------------------------------------- token mixers
+
+
+class ParallaxAttention(nn.Module):
+    """Parallax attention (softmax plus a learned first-order correction from ``r``).
+
+    Grouped query/key/value heads with per-head zero-centered RMSNorm on q, r
+    and k, partial RoPE on q, r and k, and a full-rank sigmoid output gate.
+    """
+
+    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
+        super().__init__()
+        self.layer_index = layer_index
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.attention_head_dim
+        self.scale = self.head_dim ** -0.5
+        query_width = self.num_heads * self.head_dim
+        kv_width = self.num_kv_heads * self.head_dim
+        hidden = config.hidden_size
+        self.q_proj = nn.Linear(hidden, query_width, bias=False)
+        self.r_proj = nn.Linear(hidden, query_width, bias=False)
+        self.k_proj = nn.Linear(hidden, kv_width, bias=False)
+        self.v_proj = nn.Linear(hidden, kv_width, bias=False)
+        self.g_proj = nn.Linear(hidden, query_width, bias=True)
+        self.o_proj = nn.Linear(query_width, hidden, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        self.r_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        for projection in (self.q_proj, self.r_proj):
+            mark_muon_row_blocks(projection.weight, self.num_heads)
+        for projection in (self.k_proj, self.v_proj):
+            mark_muon_row_blocks(projection.weight, self.num_kv_heads)
+        mark_adamw_only(self.g_proj.weight)
+
+    def forward(self, hidden_states: torch.Tensor, context: Senryu2Context) -> torch.Tensor:
+        batch, length, _ = hidden_states.shape
+        q = self.q_norm(self.q_proj(hidden_states).view(batch, length, self.num_heads, self.head_dim))
+        r = self.r_norm(self.r_proj(hidden_states).view(batch, length, self.num_heads, self.head_dim))
+        k = self.k_norm(self.k_proj(hidden_states).view(batch, length, self.num_kv_heads, self.head_dim))
+        v = self.v_proj(hidden_states).view(batch, length, self.num_kv_heads, self.head_dim)
+        cosine, sine = context.rotary
+        q, r, k = (
+            rotary_embedding(
+                tensor,
+                cosine,
+                sine,
+                seqlen_offsets=context.position_offset,
+                cu_seqlens=context.cu_seqlens,
+            )
+            for tensor in (q, r, k)
+        )
+        use_kernel = q.is_cuda and q.dtype in (torch.bfloat16, torch.float16)
+        cache = context.cache
+        if cache is not None:
+            keys, values = cache.append_attention(self.layer_index, k, v, context.position_offset)
+            if not use_kernel:
+                out = parallax_reference(q, r, keys, values, scale=self.scale)
+            elif keys.shape[1] == length:
+                out = parallax_attention(q, r, keys, values, scale=self.scale)
+            else:
+                out = parallax_decode_cached(q, r, keys, values, scale=self.scale)
+        elif use_kernel:
+            out = parallax_attention(q, r, k, v, scale=self.scale, cu_seqlens=context.cu_seqlens)
+        else:
+            out = parallax_reference(q, r, k, v, scale=self.scale, cu_seqlens=context.cu_seqlens)
+        gate = torch.sigmoid(self.g_proj(hidden_states)).view(batch, length, self.num_heads, self.head_dim)
+        return self.o_proj((out * gate).reshape(batch, length, self.num_heads * self.head_dim))
+
+
+def build_gated_deltanet2(config: TinyGDNConfig, layer_index: int) -> GatedDeltaNet2:
+    mixer = GatedDeltaNet2(
+        hidden_size=config.hidden_size,
+        expand_v=config.linear_expand_v,
+        head_dim=config.linear_head_dim,
+        num_heads=config.linear_num_heads,
+        num_v_heads=config.linear_num_value_heads,
+        mode="chunk",
+        use_short_conv=True,
+        allow_neg_eigval=config.allow_negative_eigenvalues,
+        conv_size=config.linear_conv_kernel_dim,
+        conv_bias=False,
+        layer_idx=layer_index,
+        norm_eps=config.rms_norm_eps,
+    )
+    if config.linear_full_rank_gate:
+        # The layer calls `self.g_proj(hidden_states)`, so a single full-rank
+        # map drops in for fla's low-rank pair.
+        mixer.g_proj = nn.Linear(config.hidden_size, mixer.value_dim, bias=True)
+        mark_adamw_only(mixer.g_proj.weight)
+    else:
+        mark_adamw_only(*(layer.weight for layer in mixer.g_proj))
+    mark_adamw_only(*(layer.weight for layer in mixer.f_proj))
+    for projection in (mixer.q_proj, mixer.k_proj, mixer.b_proj):
+        mark_muon_row_blocks(projection.weight, mixer.num_heads)
+    for projection in (mixer.v_proj, mixer.w_proj):
+        mark_muon_row_blocks(projection.weight, mixer.num_v_heads)
+    return mixer
+
+
+# ------------------------------------------------------------------- n-gram memory
+
+
+class BranchRMSNorm(nn.Module):
+    """Zero-centered RMSNorm over the last dim with one gain vector per branch."""
+
+    def __init__(self, branches: int, hidden_size: int, eps: float) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(branches, hidden_size))
+        self.eps = eps
+        mark_gain(self.weight)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        values = hidden_states.float()
+        values = values * torch.rsqrt(values.square().mean(dim=-1, keepdim=True) + self.eps)
+        return values * (1.0 + self.weight.float())
+
+
+class NGramMemory(nn.Module):
+    """Engram-style memory: hashed n-gram rows, gated per residual branch, then a short conv.
+
+    ``e`` concatenates one row per table. Each branch gates the shared value
+    ``W_V e`` with ``σ(RMSNorm(R_m) · RMSNorm(W_K,m e) / √d)``; the gated values
+    pass a depthwise causal convolution (dilation = largest order, zero init)
+    and are added to their branches.
+    """
+
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.layout = NGramLayout.from_config(config)
+        self.branches = config.residual_branches
+        self.hidden_size = config.hidden_size
+        self.kernel = config.ngram_conv_kernel
+        self.dilation = self.layout.max_order
+        memory_dim = self.layout.memory_dim
+        self.value_proj = nn.Linear(memory_dim, config.hidden_size, bias=False)
+        self.key_proj = nn.Linear(memory_dim, self.branches * config.hidden_size, bias=False)
+        self.hidden_norm = BranchRMSNorm(self.branches, config.hidden_size, config.rms_norm_eps)
+        self.key_norm = BranchRMSNorm(self.branches, config.hidden_size, config.rms_norm_eps)
+        self.conv_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.conv_weight = nn.Parameter(torch.zeros(config.hidden_size, self.kernel))
+        mark_muon_row_blocks(self.key_proj.weight, self.branches)
+        mark_gain(self.conv_weight)
+        self.register_buffer("canonical_ids", torch.arange(config.vocab_size, dtype=torch.int64), persistent=False)
+        self.register_buffer("tables", None, persistent=False)
+
+    @property
+    def conv_history(self) -> int:
+        return (self.kernel - 1) * self.dilation
+
+    def set_canonical_ids(self, canonical: torch.Tensor) -> None:
+        if canonical.shape != self.canonical_ids.shape:
+            raise ValueError("canonical map does not match the vocabulary")
+        self.canonical_ids = canonical.to(device=self.canonical_ids.device, dtype=torch.int64)
+
+    def attach_tables(self, tables: torch.Tensor, canonical: torch.Tensor) -> None:
+        """Look rows up in ``tables [total_rows, embedding_dim]``.
+
+        The tables may stay in host memory while the model runs on the GPU:
+        decoding reads only ``table_count`` rows per token.
+        """
+
+        if tables.shape != (self.layout.total_rows, self.layout.embedding_dim):
+            raise ValueError(f"tables have shape {tuple(tables.shape)}")
+        self.set_canonical_ids(canonical)
+        self.tables = tables
+
+    def load_tables(self, path: Path, *, device: torch.device | str, dtype: torch.dtype) -> None:
+        tables, canonical = load_table_file(path, self.layout, device=device, dtype=dtype)
+        self.attach_tables(tables, canonical)
+
+    def lookup(
+        self,
+        input_ids: torch.Tensor,
+        *,
+        distance: torch.Tensor | None = None,
+        history: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Concatenated rows ``[..., memory_dim]`` for every position of ``input_ids``."""
+
+        if self.tables is None:
+            raise RuntimeError(
+                "Senryu-2 needs n-gram embeddings: pass ngram_embeddings or attach tables "
+                "with model.ngram.load_tables(...)"
+            )
+        canonical = self.canonical_ids[input_ids]
+        indices = ngram_row_indices(canonical, self.layout, distance=distance, history=history)
+        tables = self.tables
+        if tables.device == indices.device:
+            return F.embedding(indices, tables).flatten(-2)
+        rows = F.embedding(indices.to(tables.device), tables).flatten(-2)
+        return rows.to(indices.device, non_blocking=True)
+
+    def _convolve(
+        self,
+        values: torch.Tensor,
+        distance: torch.Tensor | None,
+        history: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Causal dilated depthwise conv over ``[B, T, n, d]``; returns ``(out, new_history)``."""
+
+        batch, length = values.shape[:2]
+        span = self.conv_history
+        if history is None:
+            history = values.new_zeros(batch, span, self.branches, self.hidden_size)
+        extended = torch.cat([history, values], dim=1)
+        weight = self.conv_weight.to(values.dtype)
+        out = values * weight[:, -1]
+        for tap in range(1, self.kernel):
+            shift = tap * self.dilation
+            shifted = extended[:, span - shift : span - shift + length]
+            if distance is not None:
+                shifted = shifted * (distance >= shift).to(values.dtype)[..., None, None]
+            out = out + shifted * weight[:, -1 - tap]
+        return out, extended[:, -span:]
+
+    def forward(
+        self,
+        residual: torch.Tensor,
+        embeddings: torch.Tensor,
+        *,
+        batch: int,
+        distance: torch.Tensor | None,
+        conv_state: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Branch updates ``[N, n, d]`` for branches ``[N, n, d]`` and rows ``[N, memory_dim]``."""
+
+        rows = residual.shape[0]
+        value = self.value_proj(embeddings)
+        keys = self.key_proj(embeddings).view(rows, self.branches, self.hidden_size)
+        similarity = (self.hidden_norm(residual) * self.key_norm(keys)).sum(dim=-1)
+        alpha = torch.sigmoid(similarity / math.sqrt(self.hidden_size)).to(value.dtype)
+        gated = alpha.unsqueeze(-1) * value.unsqueeze(1)
+        normalized = self.conv_norm(gated).view(batch, rows // batch, self.branches, self.hidden_size)
+        convolved, new_state = self._convolve(
+            normalized,
+            None if distance is None else distance.view(batch, rows // batch),
+            conv_state,
+        )
+        return F.silu(convolved).view(rows, self.branches, self.hidden_size) + gated, new_state
+
+
+# ------------------------------------------------------------------------- blocks
+
+
+class Senryu2Block(nn.Module):
+    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
+        super().__init__()
+        self.layer_index = layer_index
+        self.layer_type = config.layer_types[layer_index]
+        residual = (config.residual_branches, config.hidden_size, config.residual_gate_rank, config.rms_norm_eps)
+        self.mixer_residual = GatedResidual(*residual, write=True)
+        self.mlp_residual = GatedResidual(*residual, write=True)
+        if self.layer_type == "gdn2":
+            self.token_mixer: nn.Module = build_gated_deltanet2(config, layer_index)
+        elif self.layer_type == "parallax":
+            self.token_mixer = ParallaxAttention(config, layer_index)
+        else:
+            raise ValueError(f"Senryu-2 has no {self.layer_type} layers")
+        self.mlp = SiTUGLU(config)
+        mark_muon_row_blocks(self.mlp.gate_up_proj.weight, 2)
+
+    def _mix(self, hidden_states: torch.Tensor, context: Senryu2Context) -> torch.Tensor:
+        batched = hidden_states.view(context.batch, context.length, -1)
+        if self.layer_type == "parallax":
+            mixed = self.token_mixer(batched, context)
+        else:
+            cache = context.cache
+            mixed, _, _ = self.token_mixer(
+                batched,
+                past_key_values=None if cache is None else cache.linear,
+                use_cache=cache is not None,
+                cu_seqlens=context.cu_seqlens,
+            )
+        return mixed.reshape(hidden_states.shape)
+
+    def forward(self, residual: torch.Tensor, context: Senryu2Context) -> torch.Tensor:
+        hidden_states, scales = self.mixer_residual.read(residual)
+        residual = GatedResidual.write(residual, scales, self._mix(hidden_states, context))
+        hidden_states, scales = self.mlp_residual.read(residual)
+        return GatedResidual.write(residual, scales, self.mlp(hidden_states))
+
+
+# -------------------------------------------------------------------------- cache
+
+
+@dataclass
+class Senryu2Cache:
+    """Decode state: GDN-2 states, Parallax keys/values, and the n-gram tails."""
+
+    linear: Any = field(default_factory=FlaCache)
+    keys: dict[int, torch.Tensor] = field(default_factory=dict)
+    values: dict[int, torch.Tensor] = field(default_factory=dict)
+    length: int = 0
+    ngram_tokens: torch.Tensor | None = None
+    ngram_conv: torch.Tensor | None = None
+
+    def append_attention(
+        self,
+        layer_index: int,
+        keys: torch.Tensor,
+        values: torch.Tensor,
+        start: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Write ``[B, T, H, D]`` keys/values at ``start``; return views of the whole prefix."""
+
+        stop = start + keys.shape[1]
+        stored_keys = self.keys.get(layer_index)
+        if stored_keys is None or stored_keys.shape[1] < stop:
+            capacity = max(stop, 0 if stored_keys is None else 2 * stored_keys.shape[1])
+            capacity = -(-capacity // KV_CACHE_GRANULE) * KV_CACHE_GRANULE
+            grown_keys = keys.new_empty(keys.shape[0], capacity, *keys.shape[2:])
+            grown_values = values.new_empty(values.shape[0], capacity, *values.shape[2:])
+            if stored_keys is not None:
+                grown_keys[:, :start] = stored_keys[:, :start]
+                grown_values[:, :start] = self.values[layer_index][:, :start]
+            self.keys[layer_index] = stored_keys = grown_keys
+            self.values[layer_index] = grown_values
+        stored_values = self.values[layer_index]
+        stored_keys[:, start:stop] = keys
+        stored_values[:, start:stop] = values
+        return stored_keys[:, :stop], stored_values[:, :stop]
+
+
+# -------------------------------------------------------------------------- model
+
+
+@dataclass
+class _Packing:
+    """How a ``[B, T]`` batch maps onto one packed row of ``valid`` tokens."""
+
+    batch: int
+    length: int
+    index: torch.Tensor | None
+    row_starts: torch.Tensor
+
+    def pack(self, tensor: torch.Tensor) -> torch.Tensor:
+        flat = tensor.reshape(self.batch * self.length, *tensor.shape[2:])
+        return flat if self.index is None else flat.index_select(0, self.index)
+
+    def unpack(self, tensor: torch.Tensor) -> torch.Tensor:
+        if self.index is None:
+            return tensor.view(self.batch, self.length, *tensor.shape[1:])
+        out = tensor.new_zeros(self.batch * self.length, *tensor.shape[1:])
+        out.index_copy_(0, self.index, tensor)
+        return out.view(self.batch, self.length, *tensor.shape[1:])
+
+
+class Senryu2ForCausalLM(nn.Module):
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        if config.residual_mode != "gated_residual":
+            raise ValueError("Senryu2ForCausalLM needs residual_mode 'gated_residual'")
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(Senryu2Block(config, index) for index in range(config.num_hidden_layers))
+        self.final_residual = GatedResidual(
+            config.residual_branches,
+            config.hidden_size,
+            config.residual_gate_rank,
+            config.rms_norm_eps,
+            write=False,
+        )
+        self.ngram = NGramMemory(config) if config.ngram_orders else None
+        self.rotary = RotaryTables(config.rotary_dim, config.rope_theta)
+        self.checkpointed_layers: frozenset[int] = frozenset()
+        self.apply(self._initialize_module)
+        self._initialize_residual_projections()
+
+    # ---------------------------------------------------------------- setup
+    def _initialize_module(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+
+    def _initialize_residual_projections(self) -> None:
+        residual_std = self.config.initializer_range / math.sqrt(2 * self.config.num_hidden_layers)
+        for layer in self.layers:
+            nn.init.normal_(layer.token_mixer.o_proj.weight, mean=0.0, std=residual_std)
+            nn.init.normal_(layer.mlp.down_proj.weight, mean=0.0, std=residual_std)
+
+    def enable_gradient_checkpointing(self, enabled: bool = True, *, layers: int | None = None) -> None:
+        """Recompute block activations in backward; ``layers`` picks GDN-2 blocks first."""
+
+        if not enabled:
+            if layers is not None:
+                raise ValueError("layers needs gradient checkpointing enabled")
+            self.checkpointed_layers = frozenset()
+            return
+        count = len(self.layers) if layers is None else layers
+        if not 1 <= count <= len(self.layers):
+            raise ValueError(f"layers must be between 1 and {len(self.layers)}")
+        cheapest_replay_first = sorted(
+            range(len(self.layers)),
+            key=lambda index: self.layers[index].layer_type != "gdn2",
+        )
+        self.checkpointed_layers = frozenset(cheapest_replay_first[:count])
+
+    def describe_gradient_checkpointing(self) -> str:
+        if not self.checkpointed_layers:
+            return "off"
+        layer_types = sorted({self.layers[index].layer_type for index in self.checkpointed_layers})
+        return f"{len(self.checkpointed_layers)}/{len(self.layers)} blocks ({', '.join(layer_types)})"
+
+    def project_to_vocabulary(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return F.linear(hidden_states, self.embed_tokens.weight)
+
+    # -------------------------------------------------------------- packing
+    def _packing(self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None) -> _Packing:
+        batch, length = input_ids.shape
+        device = input_ids.device
+        if attention_mask is None or bool(attention_mask.all()):
+            row_starts = torch.arange(0, batch * length, length, device=device)
+            return _Packing(batch, length, None, row_starts)
+        if attention_mask.shape != input_ids.shape:
+            raise ValueError("attention_mask must match input_ids")
+        valid = attention_mask.reshape(-1).bool()
+        index = valid.nonzero().flatten()
+        lengths = attention_mask.bool().sum(dim=1)
+        if bool((lengths == 0).any()):
+            raise ValueError("every row needs at least one unmasked token")
+        row_starts = torch.cumsum(lengths, dim=0) - lengths
+        return _Packing(batch, length, index, row_starts)
+
+    def packed_layout(
+        self,
+        tokens: torch.Tensor,
+        row_starts: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """``(cu_seqlens int32, distance int64)`` for a flat token stream split at BOS."""
+
+        distance = document_distance(tokens, self.config.bos_token_id, row_starts)
+        starts = (distance == 0).nonzero().flatten()
+        cu_seqlens = torch.cat([starts, starts.new_tensor([tokens.numel()])]).to(torch.int32)
+        return cu_seqlens, distance
+
+    # -------------------------------------------------------------- forward
+    def _run_block(self, index: int, residual: torch.Tensor, context: Senryu2Context) -> torch.Tensor:
+        layer = self.layers[index]
+        if self.training and index in self.checkpointed_layers:
+            return checkpoint(layer, residual, context, use_reentrant=False)
+        return layer(residual, context)
+
+    def _inject_ngrams(
+        self,
+        residual: torch.Tensor,
+        embeddings: torch.Tensor,
+        context: Senryu2Context,
+    ) -> torch.Tensor:
+        assert self.ngram is not None
+        cache = context.cache
+        conv_state = None if cache is None else cache.ngram_conv
+        distance = context.distance if cache is None else None
+        if self.training and self.checkpointed_layers:
+            update, new_state = checkpoint(
+                self.ngram,
+                residual,
+                embeddings,
+                batch=context.batch,
+                distance=distance,
+                conv_state=conv_state,
+                use_reentrant=False,
+            )
+        else:
+            update, new_state = self.ngram(
+                residual,
+                embeddings,
+                batch=context.batch,
+                distance=distance,
+                conv_state=conv_state,
+            )
+        if cache is not None:
+            cache.ngram_conv = new_state
+        return residual + update
+
+    def _ngram_rows(
+        self,
+        input_ids: torch.Tensor,
+        packing: _Packing | None,
+        distance: torch.Tensor | None,
+        ngram_embeddings: torch.Tensor | None,
+        cache: Senryu2Cache | None,
+    ) -> torch.Tensor:
+        assert self.ngram is not None
+        layout = self.ngram.layout
+        if cache is not None:
+            history = cache.ngram_tokens
+            if history is None:
+                history = input_ids.new_full((input_ids.shape[0], layout.max_order - 1), layout.pad_token_id)
+            canonical = self.ngram.canonical_ids[input_ids]
+            cache.ngram_tokens = torch.cat([history, canonical], dim=1)[:, -(layout.max_order - 1) :]
+            if ngram_embeddings is not None:
+                return ngram_embeddings.reshape(-1, layout.memory_dim)
+            return self.ngram.lookup(input_ids, history=history).reshape(-1, layout.memory_dim)
+        assert packing is not None
+        if ngram_embeddings is not None:
+            if ngram_embeddings.shape != (*input_ids.shape, layout.memory_dim):
+                raise ValueError(
+                    f"ngram_embeddings must have shape {(*input_ids.shape, layout.memory_dim)}"
+                )
+            return packing.pack(ngram_embeddings)
+        return self.ngram.lookup(packing.pack(input_ids), distance=distance)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Senryu2Cache | None = None,
+        *,
+        use_cache: bool = False,
+        return_logits: bool = True,
+        return_hidden_states: bool = False,
+        labels_are_shifted: bool = False,
+        include_mtp_loss: bool = True,
+        mtp_loss_weight: float | None = None,
+        z_loss_coefficient: float = 0.0,
+        logits_to_keep: int | None = None,
+        ngram_embeddings: torch.Tensor | None = None,
+    ) -> TinyGDNOutput:
+        del include_mtp_loss, mtp_loss_weight
+        if input_ids.ndim != 2:
+            raise ValueError("input_ids must have shape [batch, sequence]")
+        if labels is not None and labels.shape != input_ids.shape:
+            raise ValueError("labels must have the same shape as input_ids")
+        if z_loss_coefficient < 0.0:
+            raise ValueError("z_loss_coefficient cannot be negative")
+        if logits_to_keep is not None and logits_to_keep <= 0:
+            raise ValueError("logits_to_keep must be positive")
+        if use_cache and labels is not None:
+            raise ValueError("use_cache is not supported with labels")
+        if past_key_values is not None and not isinstance(past_key_values, Senryu2Cache):
+            raise TypeError("past_key_values must be a Senryu2Cache or None")
+        cached = use_cache or past_key_values is not None
+        batch, length = input_ids.shape
+        offset = 0 if past_key_values is None else past_key_values.length
+        if offset + length > self.config.max_position_embeddings:
+            raise ValueError("Input exceeds max_position_embeddings")
+
+        embeddings = self.embed_tokens(input_ids)
+        device, dtype = embeddings.device, embeddings.dtype
+        packing: _Packing | None = None
+        distance: torch.Tensor | None = None
+        if cached:
+            if attention_mask is not None and not bool(attention_mask.all()):
+                raise ValueError("Senryu-2 decodes unpadded rows; drop the padding instead of masking it")
+            cache = past_key_values if past_key_values is not None else Senryu2Cache()
+            context = Senryu2Context(
+                batch=batch,
+                length=length,
+                cu_seqlens=None,
+                distance=None,
+                rotary=self.rotary.get(offset + length, device, dtype),
+                position_offset=offset,
+                cache=cache,
+            )
+            hidden = embeddings.reshape(batch * length, -1)
+        else:
+            cache = None
+            packing = self._packing(input_ids, attention_mask)
+            tokens = packing.pack(input_ids)
+            cu_seqlens, distance = self.packed_layout(tokens, packing.row_starts)
+            context = Senryu2Context(
+                batch=1,
+                length=tokens.numel(),
+                cu_seqlens=cu_seqlens,
+                distance=distance,
+                rotary=self.rotary.get(tokens.numel(), device, dtype),
+            )
+            hidden = packing.pack(embeddings)
+        residual = hidden.unsqueeze(1).expand(-1, self.config.residual_branches, -1)
+
+        for index in range(len(self.layers)):
+            if self.ngram is not None and index == self.config.ngram_layer_index:
+                rows = self._ngram_rows(input_ids, packing, distance, ngram_embeddings, cache)
+                residual = self._inject_ngrams(residual, rows.to(dtype), context)
+            residual = self._run_block(index, residual, context)
+        hidden_states, _ = self.final_residual.read(residual)
+        if cache is not None:
+            cache.length = offset + length
+
+        main_loss = None
+        z_loss = None
+        total_loss = None
+        if labels is not None:
+            assert packing is not None
+            if labels_are_shifted:
+                targets = labels
+            else:
+                targets = torch.cat([labels[:, 1:], labels.new_full((batch, 1), -100)], dim=1)
+            total_loss, main_loss, z_loss = linear_cross_entropy(
+                hidden_states,
+                self.embed_tokens.weight,
+                packing.pack(targets).reshape(-1),
+                z_loss_coefficient=z_loss_coefficient,
+            )
+            if z_loss_coefficient == 0.0:
+                z_loss = None
+
+        batched_states = (
+            hidden_states.view(batch, length, -1) if packing is None else packing.unpack(hidden_states)
+        )
+        output_states = batched_states if logits_to_keep is None else batched_states[:, -logits_to_keep:]
+        logits = self.project_to_vocabulary(output_states) if return_logits else None
+        return TinyGDNOutput(
+            loss=total_loss,
+            logits=logits,
+            main_loss=main_loss,
+            mtp_loss=None,
+            z_loss=z_loss,
+            hidden_states=batched_states if return_hidden_states else None,
+            past_key_values=cache if use_cache else None,
+        )
+
+    # ------------------------------------------------------------ reporting
+    def parameter_report(self) -> dict[str, int]:
+        total = sum(parameter.numel() for parameter in self.parameters())
+        embeddings = self.embed_tokens.weight.numel()
+        tables = 0 if self.ngram is None else self.ngram.layout.total_rows * self.ngram.layout.embedding_dim
+        return {
+            "deployable_core": total,
+            "training_total": total,
+            "embedding": embeddings,
+            "mtp_auxiliary": 0,
+            "non_embedding_core": total - embeddings,
+            "ngram_table": tables,
+        }
+
+    def save_checkpoint(self, output_dir: Path) -> None:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        self.config.save_json(output_dir / "config.json")
+        save_model(self, output_dir / "model.safetensors")
+
+    def load_ngram_tables(
+        self,
+        checkpoint_dir: Path,
+        *,
+        weights: str,
+        device: str | torch.device,
+        dtype: torch.dtype,
+    ) -> None:
+        """Attach the checkpoint's tables that match ``weights`` ("normal" or "ema")."""
+
+        if self.ngram is None:
+            return
+        table_file = checkpoint_dir / (NGRAM_TABLES_EMA_FILE if weights == "ema" else NGRAM_TABLES_FILE)
+        if not table_file.is_file():
+            raise FileNotFoundError(f"Senryu-2 checkpoint has no n-gram tables: {table_file}")
+        self.ngram.load_tables(table_file, device=device, dtype=dtype)
+
+    def extra_repr(self) -> str:
+        report = self.parameter_report()
+        return f"core_parameters={report['deployable_core']:,}, ngram_table_parameters={report['ngram_table']:,}"
+
+    def get_architecture_metadata(self) -> dict[str, Any]:
+        return {
+            "architecture": self.config.architecture,
+            "layer_types": list(self.config.layer_types),
+            "effective_num_layers": self.config.effective_num_layers,
+            "shared_layer_indices": [],
+            "parameter_report": self.parameter_report(),
+            "features": [
+                "3:1 Gated DeltaNet-2 + Parallax attention hybrid",
+                "4-branch Gated Residual stream (gated read, per-branch write scales)",
+                "hashed n-gram memory (orders 2-3) with per-branch context gating",
+                "full-rank sigmoid output gates",
+                "partial RoPE, per-head QK/R normalization",
+                "SiTU-GLU",
+                "document-masked packed training",
+                "tied input-output embeddings",
+            ],
+        }

tokenizer_config.json

@@ -0,0 +1,69 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "bos_token": "<|begin_of_text|>",
+  "eos_token": "<|end_of_text|>",
+  "pad_token": "<|padding|>",
+  "unk_token": "<|unknown|>",
+  "model_max_length": 2048,
+  "clean_up_tokenization_spaces": false,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "chat_template": "{%- set ns = namespace(xml_tools=none, tools_emitted=false) -%}\n{%- if xml_tools is defined and xml_tools -%}\n{%- set ns.xml_tools = xml_tools -%}\n{%- elif tools is defined and tools -%}\n{%- set ns.xml_tools = tools -%}\n{%- endif -%}\n{%- set tools_preamble = 'You may call one or more functions to assist with the user query.\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\\n' -%}\n{%- set tools_epilogue = '</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\\n</tool_call>' -%}\n{%- for message in messages -%}\n{%- if loop.first -%}{{- bos_token -}}{%- endif -%}\n{%- if loop.first and ns.xml_tools and message['role'] != 'system' -%}\n{{- '<|im_start|>system\\n' + tools_preamble -}}\n{%- for tool in ns.xml_tools -%}\n{%- if tool is string -%}{{- (tool | replace('<tools>', '') | replace('</tools>', '') | trim) + '\\n' -}}\n{%- else -%}{{- tool | tojson + '\\n' -}}\n{%- endif -%}\n{%- endfor -%}\n{{- tools_epilogue + '<|im_end|>\\n' -}}\n{%- set ns.tools_emitted = true -%}\n{%- endif -%}\n{%- set raw_role = message['role'] -%}\n{%- set role = 'user' if raw_role == 'tool' or raw_role == 'function' else raw_role -%}\n{%- set content_text = namespace(value='') -%}\n{%- if message['content'] is string -%}\n{%- set content_text.value = message['content'] -%}\n{%- elif message['content'] is iterable -%}\n{%- for item in message['content'] -%}\n{%- if item['type'] == 'text' -%}{%- set content_text.value = content_text.value + item['text'] -%}{%- endif -%}\n{%- endfor -%}\n{%- endif -%}\n{%- set is_tool = raw_role == 'tool' or raw_role == 'function' -%}\n{%- set prev_is_tool = loop.previtem is defined and (loop.previtem['role'] == 'tool' or loop.previtem['role'] == 'function') -%}\n{%- set next_is_tool = loop.nextitem is defined and (loop.nextitem['role'] == 'tool' or loop.nextitem['role'] == 'function') -%}\n{%- if raw_role == 'system' and not (content_text.value | trim) and not (ns.xml_tools and not ns.tools_emitted) -%}\n{%- else -%}\n{%- if is_tool and prev_is_tool -%}\n{{- '\\n' -}}\n{%- else -%}\n{{- '<|im_start|>' + role + '\\n' -}}\n{%- endif -%}\n{%- if raw_role == 'assistant' and enable_thinking is defined -%}\n{{- ('<|think|>\\n' if enable_thinking else '<|no_think|>\\n') -}}\n{%- endif -%}\n{%- if is_tool -%}\n{{- '<|tool_response|>\\n' + (content_text.value | replace('<tool_response>', '') | replace('</tool_response>', '') | trim) -}}\n{%- elif raw_role == 'system' -%}\n{{- content_text.value | replace('/system_override', '') | replace('/no_think', '') | replace('/think', '') | trim -}}\n{%- else -%}\n{{- content_text.value -}}\n{%- endif -%}\n{%- if raw_role == 'system' and ns.xml_tools and not ns.tools_emitted and '<tools>' not in content_text.value -%}\n{%- if content_text.value | trim -%}{{- '\\n\\n' -}}{%- endif -%}\n{{- tools_preamble -}}\n{%- for tool in ns.xml_tools -%}\n{%- if tool is string -%}{{- (tool | replace('<tools>', '') | replace('</tools>', '') | trim) + '\\n' -}}\n{%- else -%}{{- tool | tojson + '\\n' -}}\n{%- endif -%}\n{%- endfor -%}\n{{- tools_epilogue -}}\n{%- set ns.tools_emitted = true -%}\n{%- endif -%}\n{%- if raw_role == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] -%}\n{%- if '<tool_call>' not in content_text.value -%}\n{%- for tool_call in message['tool_calls'] -%}\n{%- set fn = tool_call['function'] if tool_call['function'] is defined else tool_call -%}\n{%- if loop.first and not (content_text.value | trim) -%}\n{{- '<tool_call>\\n{\"name\": \"' + fn['name'] + '\", \"arguments\": ' -}}\n{%- else -%}\n{{- '\\n<tool_call>\\n{\"name\": \"' + fn['name'] + '\", \"arguments\": ' -}}\n{%- endif -%}\n{%- if fn['arguments'] is string -%}{{- fn['arguments'] -}}\n{%- else -%}{{- fn['arguments'] | tojson -}}\n{%- endif -%}\n{{- '}\\n</tool_call>' -}}\n{%- endfor -%}\n{%- endif -%}\n{%- endif -%}\n{%- if not (is_tool and next_is_tool) -%}\n{{- '<|im_end|>\\n' -}}\n{%- endif -%}\n{%- endif -%}\n{%- endfor -%}\n{%- if add_generation_prompt -%}\n{{- '<|im_start|>assistant\\n' -}}\n{%- if enable_thinking is defined -%}{{- ('<|think|>\\n' if enable_thinking else '<|no_think|>\\n') -}}{%- endif -%}\n{%- endif -%}",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|tool_call|>",
+    "<|tool_response|>",
+    "<think>",
+    "</think>",
+    "<|fim_prefix|>",
+    "<|fim_middle|>",
+    "<|fim_suffix|>",
+    "<|fim_pad|>",
+    "<|no_think|>",
+    "<|think|>",
+    "<|mask|>",
+    "<|reserved_009|>",
+    "<|reserved_010|>",
+    "<|reserved_011|>",
+    "<|reserved_012|>",
+    "<|reserved_013|>",
+    "<|reserved_014|>",
+    "<|reserved_015|>",
+    "<|reserved_016|>",
+    "<|reserved_017|>",
+    "<|reserved_018|>",
+    "<|reserved_019|>",
+    "<|reserved_020|>",
+    "<|reserved_021|>",
+    "<|reserved_022|>",
+    "<|reserved_023|>",
+    "<|reserved_024|>",
+    "<|reserved_025|>",
+    "<|reserved_026|>",
+    "<|reserved_027|>",
+    "<|reserved_028|>",
+    "<|reserved_029|>",
+    "<|reserved_030|>",
+    "<|reserved_031|>",
+    "<|reserved_032|>",
+    "<|reserved_033|>",
+    "<|reserved_034|>",
+    "<|reserved_035|>",
+    "<|reserved_036|>",
+    "<|reserved_037|>",
+    "<|reserved_038|>",
+    "<|reserved_039|>",
+    "<|reserved_040|>",
+    "<|reserved_041|>",
+    "<|reserved_042|>",
+    "<|reserved_043|>",
+    "<|reserved_044|>",
+    "<|reserved_045|>",
+    "<|reserved_046|>",
+    "<|reserved_047|>",
+    "<|reserved_048|>",
+    "<|reserved_049|>"
+  ],
+  "mask_token": "<|mask|>"
+}