Upload Haiku-base (pretrain EMA @ step 8400)

README.md

@@ -0,0 +1,275 @@
+---
+license: apache-2.0
+language:
+  - en
+tags:
+  - text-generation
+  - causal-lm
+  - pytorch
+  - pretrain
+  - hybrid
+  - kimi-delta-attention
+  - gated-mla
+  - haiku
+pipeline_tag: text-generation
+library_name: tiny_gdn
+datasets:
+  - HuggingFaceFW/fineweb-edu
+model-index:
+  - name: Haiku-base
+    results: []
+---
+
+<div align="center">
+
+# Haiku-base
+
+### Pretrained base model for the Haiku family (~655M)
+
+[![Model](https://img.shields.io/badge/Model-~655M_params-blue)](.)
+[![Stage](https://img.shields.io/badge/Stage-Pretrain_(base)-orange.svg)](.)
+[![License](https://img.shields.io/badge/License-Apache_2.0-green.svg)](LICENSE)
+[![Architecture](https://img.shields.io/badge/Arch-KDA_+_Gated_MLA-purple.svg)](.)
+[![Demo](https://img.shields.io/badge/Space-haiku--demo-indigo.svg)](https://huggingface.co/spaces/kerzgrr/haiku-demo)
+
+*A larger TinyGDN hybrid: Kimi Delta Attention memory plus gated multi-head latent attention*
+
+</div>
+
+---
+
+## What this is
+
+**Haiku-base** is the **pretrained (base) checkpoint** for **Haiku**, the ~655M successor to the Tercet family.
+
+- Scales [`kerzgrr/Tercet-base`](https://huggingface.co/kerzgrr/Tercet-base) from ~502M to ~655M parameters
+- Hybrid **Kimi Delta Attention (KDA)** recurrent layers + **gated MLA** (NoPE) full-attention layers
+- Own **65,536** BPE tokenizer (not the Tercet 49k vocab)
+- This repo is **pretrain-only** raw text continuation
+- Chat / instruction SFT is **not released**
+
+This base model is for continuation and research. It will not follow instructions reliably.
+
+---
+
+## Model Architecture
+
+**Pipeline:** `Text Prompt` → `BPE-65K Tokenizer` → `Haiku Hybrid Decoder (36L)` → `Next-token Prediction`
+
+### Hybrid block schedule (×36)
+
+Every 4th layer is gated MLA; the rest are Kimi Delta Attention:
+
+`KDA, KDA, KDA, MLA, …` (3:1 recurrent-to-attention)
+
+| Component | Details |
+|-----------|---------|
+| **Kimi Delta Attention** | Linear-time recurrent memory (`flash-linear-attention`) |
+| **Gated MLA** | DeepSeek-style latent KV, content-only QK (NoPE), full-rank output gate |
+| **MLP** | SiTU-GLU |
+| **Residuals** | Block attention residual |
+| **Norm** | Zero-centered RMSNorm |
+| **Embeddings** | Tied input / output |
+
+### Technical specifications
+
+| | |
+|--|--|
+| **Architecture** | Haiku hybrid (KDA + gated MLA) |
+| **Parameters** | 655,270,488 deployable |
+| **Hidden size** | 1,024 |
+| **Intermediate (MLP)** | 3,840 |
+| **Layers** | 36 (27 KDA + 9 gated MLA) |
+| **Attention** | 8 heads, Q LoRA rank 512, KV LoRA rank 256 |
+| **Linear (KDA)** | 8 heads × 128 dim |
+| **Context (trained)** | 2,048 |
+| **Max position embeddings** | 32,768 |
+| **Vocabulary** | 65,536 (BPE) |
+| **RoPE θ** | 1,000,000 (partial factor 0.5; used by KDA) |
+| **Precision (Hub weights)** | bfloat16 EMA |
+| **Weight file** | `model.safetensors` (~1.22 GiB) |
+
+---
+
+## Training (pretrain)
+
+| | |
+|--|--|
+| **Dataset** | [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) (10.13B packed train tokens) |
+| **Tokens seen** | 4,404,019,200 |
+| **Sequence length** | 2,048 |
+| **Objective** | Next-token prediction (+ MTP during training; not used at decode) |
+| **Optimizer** | Hybrid Muon + AdamW — β₁=0.9, β₂=0.95 |
+| **Peak LR** | 2 × 10⁻⁴ |
+| **Warmup** | 1% of steps |
+| **Grad clip** | 1.0 |
+| **EMA** | Karras power EMA (γ=1.0, p=0.75, max decay 0.9999) — **this Hub file is the EMA weights** |
+| **Checkpoint** | optimizer step 8,400 |
+| **Val loss (EMA)** | 3.6904 (ppl 40.06) |
+
+---
+
+## Install
+
+### 1) System requirements
+
+- Python **3.10+**
+- **CUDA GPU strongly recommended**
+- PyTorch with CUDA matching your driver
+
+### 2) Create an environment
+
+```bash
+python -m venv .venv
+# Windows
+.venv\Scripts\activate
+# Linux / macOS
+source .venv/bin/activate
+```
+
+### 3) Install PyTorch
+
+Pick the build for your platform from https://pytorch.org. Example:
+
+```bash
+pip install torch --index-url https://download.pytorch.org/whl/cu124
+```
+
+CPU-only:
+
+```bash
+pip install torch
+```
+
+### 4) Install Python deps
+
+```bash
+pip install safetensors tokenizers huggingface_hub
+```
+
+**Flash Linear Attention is installed automatically by `inference.py`** on first run (pinned commit + Windows import patches when needed). Git must be on `PATH`.
+
+### 5) Download the inference script
+
+```bash
+curl -L -o inference.py https://huggingface.co/kerzgrr/Haiku-base/resolve/main/inference.py
+
+# or Hugging Face CLI
+hf download kerzgrr/Haiku-base inference.py --local-dir .
+```
+
+The script auto-downloads `model.safetensors`, `config.json`, `tokenizer.json`, and the `tiny_gdn/` package from this repo.
+
+---
+
+## Quick start
+
+**Single prompt (streams tokens):**
+
+```bash
+python inference.py --prompt "The history of computing begins"
+```
+
+**Interactive REPL:**
+
+```bash
+python inference.py
+```
+
+**Common options:**
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--prompt` | *(none)* | One-shot continuation; omit for REPL |
+| `--temperature` | `0.8` | Sampling temperature |
+| `--top-p` | `0.95` | Nucleus sampling |
+| `--top-k` | `50` | Top-k (0 disables) |
+| `--max-new-tokens` | `256` | Generation length |
+| `--repetition-penalty` | `1.08` | Repetition penalty |
+| `--context-length` | `2048` | Tokens kept in the window |
+| `--seed` | `42` | RNG seed |
+| `--device` | `cuda` if available | `cuda` or `cpu` |
+| `--no-stream` | off | Print the full completion at once |
+| `--no-bos` | off | Do not prepend `<\|begin_of_text\|>` |
+| `--local-dir` | *(none)* | Use a local snapshot directory |
+
+---
+
+## Files
+
+```
+kerzgrr/Haiku-base/
+  README.md
+  inference.py
+  requirements.txt
+  model.safetensors
+  config.json
+  tokenizer.json
+  tokenizer_config.json
+  special_tokens_map.json
+  special_token_ids.json
+  merges.txt
+  vocab.json
+  chat_template.jinja
+  tiny_gdn/
+    __init__.py
+    config.py
+    model.py
+    haiku_layers.py
+    nn_common.py
+```
+
+---
+
+## Limitations
+
+- **Base model**: not instruction-tuned; may ramble or fail at Q&A format
+- **Scale**: ~655M parameters — research / edge prototype, not a frontier model
+- **Dependency**: requires `flash-linear-attention` (KDA); not GGUF / llama.cpp compatible today
+- **Context**: trained at 2,048; longer windows are experimental
+- **Partial epoch**: ~4.4B of 10.13B packed FineWeb-Edu tokens
+
+---
+
+## Model family
+
+| Model | Parameters | Architecture | Stage | Hub |
+|-------|------------|--------------|-------|-----|
+| **Monostich** | ~100M | LLaMA-style | SFT | [`kerzgrr/Monostich`](https://huggingface.co/kerzgrr/Monostich) |
+| **Monostich-2-base** | ~150M | TinyGDN hybrid | Pretrain | [`kerzgrr/Monostich-2-base`](https://huggingface.co/kerzgrr/Monostich-2-base) |
+| **Monostich-2** | ~150M | TinyGDN hybrid | SFT | [`kerzgrr/Monostich-2`](https://huggingface.co/kerzgrr/Monostich-2) |
+| **Couplet-base** | ~268M | TinyGDN hybrid | Pretrain | [`kerzgrr/Couplet-base`](https://huggingface.co/kerzgrr/Couplet-base) |
+| **Couplet** | ~268M | TinyGDN hybrid | SFT | [`kerzgrr/Couplet`](https://huggingface.co/kerzgrr/Couplet) |
+| **Tercet-base** | ~502M | TinyGDN hybrid | Pretrain | [`kerzgrr/Tercet-base`](https://huggingface.co/kerzgrr/Tercet-base) |
+| **Tercet** | ~502M | TinyGDN hybrid | SFT | [`kerzgrr/Tercet`](https://huggingface.co/kerzgrr/Tercet) |
+| **Haiku-base** | ~655M | KDA + gated MLA | Pretrain | *this repo* |
+
+---
+
+## Citation
+
+```bibtex
+@misc{haikubase2026,
+  title={Haiku-base: A 655M Hybrid KDA + Gated-MLA Language Model},
+  author={kerzgrr},
+  year={2026},
+  url={https://huggingface.co/kerzgrr/Haiku-base}
+}
+```
+
+---
+
+## Acknowledgments
+
+- [flash-linear-attention](https://github.com/fla-org/flash-linear-attention) (Kimi Delta Attention)
+- [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu)
+- Tercet family: [`kerzgrr/Tercet-base`](https://huggingface.co/kerzgrr/Tercet-base)
+- PyTorch SDPA / Hugging Face Hub + tokenizers
+
+---
+
+<div align="center">
+
+*A haiku is three lines — larger than a tercet, still compact.*
+
+</div>

chat_template.jinja

@@ -0,0 +1,86 @@
+{%- set ns = namespace(xml_tools=none, tools_emitted=false) -%}
+{%- if xml_tools is defined and xml_tools -%}
+{%- set ns.xml_tools = xml_tools -%}
+{%- elif tools is defined and tools -%}
+{%- set ns.xml_tools = tools -%}
+{%- endif -%}
+{%- 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' -%}
+{%- 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>' -%}
+{%- for message in messages -%}
+{%- if loop.first -%}{{- bos_token -}}{%- endif -%}
+{%- if loop.first and ns.xml_tools and message['role'] != 'system' -%}
+{{- '<|im_start|>system\n' + tools_preamble -}}
+{%- for tool in ns.xml_tools -%}
+{%- if tool is string -%}{{- (tool | replace('<tools>', '') | replace('</tools>', '') | trim) + '\n' -}}
+{%- else -%}{{- tool | tojson + '\n' -}}
+{%- endif -%}
+{%- endfor -%}
+{{- tools_epilogue + '<|im_end|>\n' -}}
+{%- set ns.tools_emitted = true -%}
+{%- endif -%}
+{%- set raw_role = message['role'] -%}
+{%- set role = 'user' if raw_role == 'tool' or raw_role == 'function' else raw_role -%}
+{%- set content_text = namespace(value='') -%}
+{%- if message['content'] is string -%}
+{%- set content_text.value = message['content'] -%}
+{%- elif message['content'] is iterable -%}
+{%- for item in message['content'] -%}
+{%- if item['type'] == 'text' -%}{%- set content_text.value = content_text.value + item['text'] -%}{%- endif -%}
+{%- endfor -%}
+{%- endif -%}
+{%- set is_tool = raw_role == 'tool' or raw_role == 'function' -%}
+{%- set prev_is_tool = loop.previtem is defined and (loop.previtem['role'] == 'tool' or loop.previtem['role'] == 'function') -%}
+{%- set next_is_tool = loop.nextitem is defined and (loop.nextitem['role'] == 'tool' or loop.nextitem['role'] == 'function') -%}
+{%- if raw_role == 'system' and not (content_text.value | trim) and not (ns.xml_tools and not ns.tools_emitted) -%}
+{%- else -%}
+{%- if is_tool and prev_is_tool -%}
+{{- '\n' -}}
+{%- else -%}
+{{- '<|im_start|>' + role + '\n' -}}
+{%- endif -%}
+{%- if raw_role == 'assistant' and enable_thinking is defined -%}
+{{- ('<|think|>\n' if enable_thinking else '<|no_think|>\n') -}}
+{%- endif -%}
+{%- if is_tool -%}
+{{- '<|tool_response|>\n' + (content_text.value | replace('<tool_response>', '') | replace('</tool_response>', '') | trim) -}}
+{%- elif raw_role == 'system' -%}
+{{- content_text.value | replace('/system_override', '') | replace('/no_think', '') | replace('/think', '') | trim -}}
+{%- else -%}
+{{- content_text.value -}}
+{%- endif -%}
+{%- if raw_role == 'system' and ns.xml_tools and not ns.tools_emitted and '<tools>' not in content_text.value -%}
+{%- if content_text.value | trim -%}{{- '\n\n' -}}{%- endif -%}
+{{- tools_preamble -}}
+{%- for tool in ns.xml_tools -%}
+{%- if tool is string -%}{{- (tool | replace('<tools>', '') | replace('</tools>', '') | trim) + '\n' -}}
+{%- else -%}{{- tool | tojson + '\n' -}}
+{%- endif -%}
+{%- endfor -%}
+{{- tools_epilogue -}}
+{%- set ns.tools_emitted = true -%}
+{%- endif -%}
+{%- if raw_role == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] -%}
+{%- if '<tool_call>' not in content_text.value -%}
+{%- for tool_call in message['tool_calls'] -%}
+{%- set fn = tool_call['function'] if tool_call['function'] is defined else tool_call -%}
+{%- if loop.first and not (content_text.value | trim) -%}
+{{- '<tool_call>\n{"name": "' + fn['name'] + '", "arguments": ' -}}
+{%- else -%}
+{{- '\n<tool_call>\n{"name": "' + fn['name'] + '", "arguments": ' -}}
+{%- endif -%}
+{%- if fn['arguments'] is string -%}{{- fn['arguments'] -}}
+{%- else -%}{{- fn['arguments'] | tojson -}}
+{%- endif -%}
+{{- '}\n</tool_call>' -}}
+{%- endfor -%}
+{%- endif -%}
+{%- endif -%}
+{%- if not (is_tool and next_is_tool) -%}
+{{- '<|im_end|>\n' -}}
+{%- endif -%}
+{%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+{{- '<|im_start|>assistant\n' -}}
+{%- if enable_thinking is defined -%}{{- ('<|think|>\n' if enable_thinking else '<|no_think|>\n') -}}{%- endif -%}
+{%- endif -%}

config.json

@@ -0,0 +1,94 @@
+{
+  "allow_negative_eigenvalues": false,
+  "architecture": "HaikuForCausalLM",
+  "architectures": [
+    "HaikuForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "attention_head_dim": 128,
+  "attention_layer_type": "gated_mla",
+  "bos_token_id": 0,
+  "checkpoint_step": 8400,
+  "eos_token_id": 1,
+  "full_attention_interval": 4,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 3840,
+  "kda_lower_bound": -5.0,
+  "kda_safe_gate": true,
+  "layer_types": [
+    "kda",
+    "kda",
+    "kda",
+    "gated_mla",
+    "kda",
+    "kda",
+    "kda",
+    "gated_mla",
+    "kda",
+    "kda",
+    "kda",
+    "gated_mla",
+    "kda",
+    "kda",
+    "kda",
+    "gated_mla",
+    "kda",
+    "kda",
+    "kda",
+    "gated_mla",
+    "kda",
+    "kda",
+    "kda",
+    "gated_mla",
+    "kda",
+    "kda",
+    "kda",
+    "gated_mla",
+    "kda",
+    "kda",
+    "kda",
+    "gated_mla",
+    "kda",
+    "kda",
+    "kda",
+    "gated_mla"
+  ],
+  "linear_conv_kernel_dim": 4,
+  "linear_expand_v": 1.0,
+  "linear_head_dim": 128,
+  "linear_layer_type": "kda",
+  "linear_num_heads": 8,
+  "linear_num_value_heads": 8,
+  "max_position_embeddings": 32768,
+  "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_family": "Haiku",
+  "model_name": "Haiku-base",
+  "model_type": "haiku",
+  "mtp_adapter_rank": 128,
+  "mtp_loss_weight": 0.2,
+  "mtp_num_heads": 2,
+  "num_attention_heads": 8,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 8,
+  "pad_token_id": 2,
+  "partial_rotary_factor": 0.5,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000.0,
+  "shared_layer_indices": [],
+  "situ_gate_cap": 4.0,
+  "situ_up_cap": 25.0,
+  "stage": "pretrain",
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "training_sequence_length": 2048,
+  "transformers_version": "4.45.0",
+  "unk_token_id": 3,
+  "use_block_attn_res": true,
+  "vocab_size": 65536,
+  "weights": "ema"
+}

inference.py

@@ -0,0 +1,501 @@
+#!/usr/bin/env python3
+"""Standalone pretrain inference for kerzgrr/Haiku-base.
+
+Downloads model assets from the Hub (cached after first run), loads TinyGDN,
+and streams raw text continuations. This is the pretrained base model — not
+instruction-tuned. Chat SFT is not released yet — this repo is the pretrained base only.
+
+Examples:
+  python inference.py --prompt "The history of computing begins"
+  python inference.py
+  python inference.py --prompt "Once upon a time" --temperature 0.9 --top-p 0.95
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import platform
+import shutil
+import subprocess
+import sys
+import time
+import warnings
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+from tokenizers import Tokenizer
+
+
+def _silence_runtime_warnings() -> None:
+    """Hide known-noisy Triton / SDPA warnings that don't affect results."""
+    patterns = (
+        r"tl\.make_block_ptr is deprecated",
+        r"Memory efficient kernel not used because",
+        r"Memory Efficient attention has been runtime disabled",
+        r"Flash attention kernel not used because",
+        r"Torch was not compiled with flash attention",
+        r"cuDNN attention kernel not used because",
+        r"cuDNN attention has been runtime disabled",
+    )
+    for pattern in patterns:
+        warnings.filterwarnings("ignore", message=pattern)
+
+REPO_ID = "kerzgrr/Haiku-base"
+FLA_COMMIT = "cbb0a72efb55c18ca0ef4f298298317573ad2cb3"
+FLA_REPO = "https://github.com/fla-org/flash-linear-attention.git"
+PATCH_FILES = (
+    "fla/__init__.py",
+    "fla/ops/__init__.py",
+    "fla/layers/__init__.py",
+    "fla/ops/simple_gla/__init__.py",
+)
+
+
+def _download(filename: str, local_dir: Path | None) -> Path:
+    from huggingface_hub import hf_hub_download
+
+    return Path(
+        hf_hub_download(
+            repo_id=REPO_ID,
+            filename=filename,
+            local_dir=str(local_dir) if local_dir else None,
+        )
+    )
+
+
+def _run(cmd: list[str], *, cwd: Path | None = None, env: dict | None = None) -> None:
+    print("+", " ".join(cmd), flush=True)
+    merged = os.environ.copy()
+    if env:
+        merged.update(env)
+    # Avoid Windows cp1252 crashes while reading FLA's README during setup.
+    merged.setdefault("PYTHONUTF8", "1")
+    merged.setdefault("PYTHONIOENCODING", "utf-8")
+    subprocess.check_call(cmd, cwd=str(cwd) if cwd else None, env=merged)
+
+
+def _pip_install(*args: str) -> None:
+    _run([sys.executable, "-m", "pip", "install", *args])
+
+
+def _fla_importable() -> tuple[bool, str]:
+    try:
+        from fla.layers.kda import KimiDeltaAttention  # noqa: F401
+    except Exception as error:  # noqa: BLE001
+        return False, str(error)
+    return True, ""
+
+
+def _cache_root() -> Path:
+    override = os.environ.get("MONOSTICH_CACHE")
+    if override:
+        path = Path(override).expanduser().resolve()
+    else:
+        path = Path.home() / ".cache" / "haiku-base"
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def _ensure_git() -> None:
+    if shutil.which("git") is None:
+        raise RuntimeError(
+            "git is required to auto-install flash-linear-attention. "
+            "Install Git and ensure it is on PATH."
+        )
+
+
+def _apply_windows_fla_patches(fla_root: Path, local_dir: Path | None) -> None:
+    print("Applying Windows FLA import patches from the Hub …", flush=True)
+    for relative in PATCH_FILES:
+        source = _download(f"windows_fla_patches/{relative}", local_dir)
+        target = fla_root / relative
+        target.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(source, target)
+        print(f"  patched {relative}", flush=True)
+
+
+def _install_fla(local_dir: Path | None) -> None:
+    print("flash-linear-attention missing/broken — installing automatically …", flush=True)
+    _pip_install("einops", "numpy")
+    is_windows = platform.system() == "Windows"
+
+    if not is_windows:
+        _pip_install(
+            "--no-deps",
+            f"git+{FLA_REPO}@{FLA_COMMIT}",
+        )
+        return
+
+    # Windows: editable clone + Hub patches (stock FLA + Triton 3.7 breaks on import).
+    _ensure_git()
+    fla_root = _cache_root() / "flash-linear-attention"
+    if (fla_root / ".git").is_dir():
+        _run(["git", "fetch", "--depth", "1", "origin", FLA_COMMIT], cwd=fla_root)
+        _run(["git", "checkout", "--force", FLA_COMMIT], cwd=fla_root)
+    else:
+        if fla_root.exists():
+            shutil.rmtree(fla_root)
+        _run(
+            [
+                "git",
+                "clone",
+                "--filter=blob:none",
+                FLA_REPO,
+                str(fla_root),
+            ]
+        )
+        _run(["git", "fetch", "--depth", "1", "origin", FLA_COMMIT], cwd=fla_root)
+        _run(["git", "checkout", "--force", FLA_COMMIT], cwd=fla_root)
+
+    _apply_windows_fla_patches(fla_root, local_dir)
+    _pip_install("--no-build-isolation", "--no-deps", "-e", str(fla_root))
+
+
+def _ensure_fla(local_dir: Path | None) -> None:
+    ok, error = _fla_importable()
+    if ok:
+        return
+    print(f"FLA not ready ({error})", flush=True)
+    try:
+        _install_fla(local_dir)
+    except Exception as install_error:  # noqa: BLE001
+        raise RuntimeError(
+            "Automatic flash-linear-attention install failed.\n"
+            f"Original import error: {error}\n"
+            f"Install error: {install_error}"
+        ) from install_error
+
+    # Drop cached failed imports so the newly installed package is picked up.
+    for name in list(sys.modules):
+        if name == "fla" or name.startswith("fla."):
+            del sys.modules[name]
+
+    ok, error = _fla_importable()
+    if not ok:
+        raise RuntimeError(
+            "flash-linear-attention installed but still failed to import "
+            f"KimiDeltaAttention: {error}"
+        )
+    print("flash-linear-attention ready.", flush=True)
+
+
+def _ensure_tiny_gdn(local_dir: Path | None) -> Path:
+    """Return a directory that contains the tiny_gdn package on sys.path."""
+    # Prefer files next to this script (local clone / Hub snapshot).
+    here = Path(__file__).resolve().parent
+    if (here / "tiny_gdn" / "__init__.py").is_file():
+        return here
+    if local_dir and (local_dir / "tiny_gdn" / "__init__.py").is_file():
+        return local_dir
+
+    # Pull package modules from the Hub into the cache.
+    for name in (
+        "tiny_gdn/__init__.py",
+        "tiny_gdn/config.py",
+        "tiny_gdn/model.py",
+        "tiny_gdn/haiku_layers.py",
+        "tiny_gdn/nn_common.py",
+    ):
+        _download(name, local_dir)
+    # hf_hub_download with local_dir=None caches under hub/; resolve via download of init.
+    init_path = _download("tiny_gdn/__init__.py", local_dir)
+    return init_path.parent.parent
+
+
+def _sample(
+    logits: torch.Tensor,
+    *,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    generator: torch.Generator,
+) -> int:
+    logits = logits.float()
+    if temperature <= 1e-5:
+        return int(torch.argmax(logits).item())
+    logits = logits / temperature
+    if 0 < top_k < logits.shape[-1]:
+        threshold = torch.topk(logits, top_k).values[-1]
+        logits = logits.masked_fill(logits < threshold, -torch.inf)
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        probs = torch.softmax(sorted_logits, dim=-1)
+        remove = torch.cumsum(probs, dim=-1) > top_p
+        remove[1:] = remove[:-1].clone()
+        remove[0] = False
+        sorted_logits = sorted_logits.masked_fill(remove, -torch.inf)
+        logits = torch.full_like(logits, -torch.inf)
+        logits.scatter_(0, sorted_indices, sorted_logits)
+    probs = torch.softmax(logits, dim=-1)
+    return int(torch.multinomial(probs, 1, generator=generator).item())
+
+
+def _apply_repetition_penalty(
+    logits: torch.Tensor,
+    token_ids: list[int],
+    penalty: float,
+    window: int,
+) -> torch.Tensor:
+    if penalty == 1.0 or not token_ids:
+        return logits
+    recent = token_ids[-window:] if window > 0 else token_ids
+    unique = torch.tensor(list(set(recent)), dtype=torch.long, device=logits.device)
+    score = logits[unique]
+    logits[unique] = torch.where(score > 0, score / penalty, score * penalty)
+    return logits
+
+
+@torch.inference_mode()
+def generate(
+    model,
+    tokenizer: Tokenizer,
+    prompt_ids: list[int],
+    *,
+    max_new_tokens: int,
+    context_length: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    repetition_penalty: float,
+    repetition_window: int,
+    seed: int,
+    stream: bool,
+    device: torch.device,
+) -> tuple[str, int, str]:
+    eos_id = int(model.config.eos_token_id)
+    token_ids = list(prompt_ids[-context_length:])
+    generated: list[int] = []
+    decoded = ""
+    stop_reason = "max_new_tokens"
+    generator = torch.Generator(device=device)
+    generator.manual_seed(seed)
+    started = time.perf_counter()
+
+    for _ in range(max_new_tokens):
+        context = token_ids[-context_length:]
+        input_ids = torch.tensor([context], dtype=torch.long, device=device)
+        output = model(input_ids, return_logits=True, logits_to_keep=1)
+        if output.logits is None:
+            raise RuntimeError("Model returned no logits")
+        next_logits = _apply_repetition_penalty(
+            output.logits[0, -1],
+            token_ids,
+            repetition_penalty,
+            repetition_window,
+        )
+        next_id = _sample(
+            next_logits,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            generator=generator,
+        )
+        if next_id == eos_id:
+            stop_reason = "eos"
+            break
+        token_ids.append(next_id)
+        generated.append(next_id)
+        current = tokenizer.decode(generated, skip_special_tokens=True)
+        delta = (
+            current[len(decoded) :]
+            if current.startswith(decoded)
+            else current
+        )
+        decoded = current
+        if stream and delta:
+            print(delta, end="", flush=True)
+
+    if stream:
+        print(flush=True)
+    elapsed = time.perf_counter() - started
+    tps = len(generated) / max(elapsed, 1e-9)
+    if stream:
+        print(
+            f"[done] tokens={len(generated)} stop={stop_reason} "
+            f"{tps:.1f} tok/s",
+            file=sys.stderr,
+            flush=True,
+        )
+    return decoded, len(generated), stop_reason
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Haiku-base pretrained text continuation"
+    )
+    parser.add_argument(
+        "--prompt",
+        default=None,
+        help="Single prompt to continue (omit for interactive REPL)",
+    )
+    parser.add_argument("--max-new-tokens", type=int, default=256)
+    parser.add_argument("--temperature", type=float, default=0.8)
+    parser.add_argument("--top-p", type=float, default=0.95)
+    parser.add_argument("--top-k", type=int, default=50)
+    parser.add_argument("--repetition-penalty", type=float, default=1.08)
+    parser.add_argument("--repetition-window", type=int, default=256)
+    parser.add_argument("--context-length", type=int, default=2048)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--device",
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        choices=["cuda", "cpu"],
+    )
+    parser.add_argument(
+        "--no-stream",
+        action="store_true",
+        help="Disable token streaming (print full completion at once)",
+    )
+    parser.add_argument(
+        "--no-bos",
+        action="store_true",
+        help="Do not prepend <|begin_of_text|> to the prompt",
+    )
+    parser.add_argument(
+        "--local-dir",
+        default=None,
+        help="Optional local snapshot directory (skips re-download when populated)",
+    )
+    parser.add_argument(
+        "--repo-id",
+        default=REPO_ID,
+        help=f"Hub repo id (default: {REPO_ID})",
+    )
+    return parser.parse_args()
+
+
+def main() -> int:
+    _silence_runtime_warnings()
+    args = parse_args()
+    global REPO_ID
+    REPO_ID = args.repo_id
+    local_dir = Path(args.local_dir).resolve() if args.local_dir else None
+
+    print(f"Loading Haiku-base from huggingface.co/{REPO_ID} …", flush=True)
+    try:
+        package_root = _ensure_tiny_gdn(local_dir)
+    except Exception as error:  # noqa: BLE001
+        print(f"Failed to resolve tiny_gdn package: {error}", file=sys.stderr)
+        print(
+            "Make sure huggingface_hub is installed and you can reach the Hub.",
+            file=sys.stderr,
+        )
+        return 1
+
+    if str(package_root) not in sys.path:
+        sys.path.insert(0, str(package_root))
+
+    try:
+        _ensure_fla(local_dir)
+    except Exception as error:  # noqa: BLE001
+        print(str(error), file=sys.stderr)
+        return 1
+
+    try:
+        from tiny_gdn import TinyGDNConfig, TinyGDNForCausalLM
+    except ImportError as error:
+        print(f"Could not import tiny_gdn: {error}", file=sys.stderr)
+        return 1
+
+    weights_path = _download("model.safetensors", local_dir)
+    tok_path = _download("tokenizer.json", local_dir)
+    cfg_path = _download("config.json", local_dir)
+
+    raw = json.loads(cfg_path.read_text(encoding="utf-8"))
+    # Config.from_json rejects unknown HF-only keys — filter to dataclass fields.
+    from dataclasses import fields
+
+    allowed = {item.name for item in fields(TinyGDNConfig)}
+    payload = {key: value for key, value in raw.items() if key in allowed}
+    if "shared_layer_indices" in payload:
+        payload["shared_layer_indices"] = tuple(payload["shared_layer_indices"])
+    config = TinyGDNConfig(**payload)
+
+    device = torch.device(args.device)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        print("CUDA requested but unavailable; falling back to CPU.", flush=True)
+        device = torch.device("cpu")
+    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
+
+    print(
+        f"Building Haiku ({config.num_hidden_layers}L / {config.hidden_size}d) "
+        f"on {device} …",
+        flush=True,
+    )
+    model = TinyGDNForCausalLM(config)
+    state = load_file(str(weights_path), device="cpu")
+    model.load_state_dict(state, strict=True)
+    del state
+    model = model.to(device=device, dtype=dtype)
+    model.eval()
+    model.requires_grad_(False)
+
+    tokenizer = Tokenizer.from_file(str(tok_path))
+    bos_id = config.bos_token_id
+    context_length = min(args.context_length, config.max_position_embeddings)
+    stream = not args.no_stream
+
+    def encode_prompt(text: str) -> list[int]:
+        ids = tokenizer.encode(text, add_special_tokens=False).ids
+        if not args.no_bos and (not ids or ids[0] != bos_id):
+            ids = [bos_id] + ids
+        return ids
+
+    def run_once(prompt: str) -> None:
+        prompt_ids = encode_prompt(prompt)
+        if stream:
+            print(prompt, end="", flush=True)
+        text, _, _ = generate(
+            model,
+            tokenizer,
+            prompt_ids,
+            max_new_tokens=args.max_new_tokens,
+            context_length=context_length,
+            temperature=args.temperature,
+            top_p=args.top_p,
+            top_k=args.top_k,
+            repetition_penalty=args.repetition_penalty,
+            repetition_window=args.repetition_window,
+            seed=args.seed,
+            stream=stream,
+            device=device,
+        )
+        if not stream:
+            print(prompt + text)
+
+    if args.prompt is not None:
+        run_once(args.prompt)
+        return 0
+
+    print(
+        "Interactive pretrain continuation. Type a prompt and press Enter.\n"
+        "Commands: /exit  /quit  /reset (clears nothing persistent; just a noop marker)\n"
+        "Note: this is the BASE model — raw continuation, not chat.",
+        flush=True,
+    )
+    while True:
+        try:
+            user_input = input("prompt> ")
+        except (EOFError, KeyboardInterrupt):
+            print()
+            break
+        text = user_input.strip()
+        if not text:
+            continue
+        if text.lower() in {"/exit", "/quit"}:
+            break
+        if text.lower() == "/reset":
+            print("(history not kept in pretrain mode)", flush=True)
+            continue
+        if stream:
+            print("completion> ", end="", flush=True)
+        run_once(text)
+        print(flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

model.safetensors

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

requirements.txt

@@ -0,0 +1,9 @@
+torch>=2.4.0
+safetensors>=0.4.0
+tokenizers>=0.20.0
+huggingface_hub>=0.26.0
+einops>=0.8.0
+numpy>=1.26.0
+
+# Install separately (pinned commit used for Haiku):
+# pip install --no-deps "git+https://github.com/fla-org/flash-linear-attention.git@cbb0a72efb55c18ca0ef4f298298317573ad2cb3"

special_token_ids.json

@@ -0,0 +1,10 @@
+{
+  "bos_token": "<|begin_of_text|>",
+  "eos_token": "<|end_of_text|>",
+  "pad_token": "<|padding|>",
+  "unk_token": "<|unknown|>",
+  "bos_token_id": 0,
+  "eos_token_id": 1,
+  "pad_token_id": 2,
+  "unk_token_id": 3
+}

special_tokens_map.json

@@ -0,0 +1,68 @@
+{
+  "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|>",
+    "<|reserved_002|>",
+    "<|reserved_003|>",
+    "<|reserved_004|>",
+    "<|reserved_005|>",
+    "<|reserved_006|>",
+    "<|reserved_007|>",
+    "<|reserved_008|>",
+    "<|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|>"
+  ]
+}

tiny_gdn/__init__.py

@@ -0,0 +1,9 @@
+from tiny_gdn.config import TinyGDNConfig, haiku_config
+from tiny_gdn.model import TinyGDNForCausalLM, TinyGDNOutput
+
+__all__ = [
+    "TinyGDNConfig",
+    "TinyGDNForCausalLM",
+    "TinyGDNOutput",
+    "haiku_config",
+]

tiny_gdn/config.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, fields
+from pathlib import Path
+from typing import Any
+
+
+@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
+
+    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
+
+    def __post_init__(self) -> None:
+        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 {"gdn2", "kda"}:
+            raise ValueError("linear_layer_type must be 'gdn2' or 'kda'")
+        if self.attention_layer_type not in {"full_attention", "gated_mla"}:
+            raise ValueError("attention_layer_type must be 'full_attention' or 'gated_mla'")
+        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)
+
+    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)
+        if "shared_layer_indices" in payload:
+            payload["shared_layer_indices"] = tuple(payload["shared_layer_indices"])
+        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)

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
+
+
+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).chunk(2, dim=-1)
+        gated = _softcap(gate, self.gate_cap) * torch.sigmoid(gate)
+        return self.down_proj(gated * _softcap(up, 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
+
+
+def _softcap(values: torch.Tensor, cap: float) -> torch.Tensor:
+    return cap * torch.tanh(values / cap)

tiny_gdn/model.py

@@ -0,0 +1,713 @@
+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 TinyGDNConfig
+from tiny_gdn.haiku_layers import (
+    BlockAttentionResidual,
+    GatedMultiheadLatentAttention,
+    SiTUGLU,
+)
+from tiny_gdn.nn_common import RMSNorm, attention_sdpa_backends
+
+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
+
+
+@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 == "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 {"gdn2", "kda"}:
+            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.gradient_checkpointing = False
+
+        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,
+            )
+
+    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) -> None:
+        self.gradient_checkpointing = enabled
+
+    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: TinyGDNBlock,
+        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]:
+        if self.gradient_checkpointing and self.training:
+            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,
+        compute_z_loss: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        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)
+        logits = self.project_to_vocabulary(prediction_states)
+        cross_entropy = F.cross_entropy(
+            logits.reshape(-1, self.config.vocab_size),
+            targets.reshape(-1),
+            ignore_index=-100,
+        )
+        z_loss = None
+        if compute_z_loss:
+            valid_targets = targets.ne(-100)
+            # Keep logits in their training dtype for logsumexp. Casting the
+            # full [B, S, V] tensor to fp32 materializes ~3 GiB in the autograd
+            # graph at 16k; upcast only the reduced [B, S] partition.
+            log_partition = torch.logsumexp(logits, dim=-1).float()
+            z_loss = log_partition.square()[valid_targets].mean()
+        return cross_entropy, z_loss
+
+    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,
+                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,
+                    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_loss, z_loss = self._causal_loss(
+                hidden_states,
+                labels,
+                target_offset=main_target_offset,
+                compute_z_loss=z_loss_coefficient > 0.0,
+            )
+            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_loss + effective_mtp_weight * mtp_loss
+            else:
+                total_loss = main_loss
+            if z_loss is not None:
+                total_loss = total_loss + z_loss_coefficient * z_loss
+
+        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/nn_common.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn.attention import SDPBackend
+
+
+class RMSNorm(nn.Module):
+    """Zero-centered RMSNorm as used by Qwen3-Next."""
+
+    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:
+        input_dtype = hidden_states.dtype
+        normalized = hidden_states.float()
+        normalized = normalized * torch.rsqrt(
+            normalized.square().mean(dim=-1, keepdim=True) + self.eps
+        )
+        normalized = normalized * (1.0 + self.weight.float())
+        return normalized.to(dtype=input_dtype)
+
+
+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,
+    ]

tokenizer_config.json

@@ -0,0 +1,74 @@
+{
+  "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|>",
+    "<|reserved_002|>",
+    "<|reserved_003|>",
+    "<|reserved_004|>",
+    "<|reserved_005|>",
+    "<|reserved_006|>",
+    "<|reserved_007|>",
+    "<|reserved_008|>",
+    "<|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|>"
+  ]
+}

validation.json

@@ -0,0 +1,21 @@
+{
+  "checkpoint": "checkpoint-00008400",
+  "ema": {
+    "batches": 16,
+    "elapsed_seconds": 7.330645765992813,
+    "loss": 3.6904296875,
+    "perplexity": 40.06205742476025,
+    "tokens": 524288,
+    "tokens_per_second": 71520.02930385683
+  },
+  "normal": {
+    "batches": 16,
+    "elapsed_seconds": 7.311252611019881,
+    "loss": 3.736328125,
+    "perplexity": 41.943695056893915,
+    "tokens": 524288,
+    "tokens_per_second": 71709.73674329993
+  },
+  "optimizer_step": 8400,
+  "type": "validation"
+}

windows_fla_patches/fla/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# Keep package import light. Eagerly importing fla.layers pulls every Triton
+# kernel and breaks on Windows + Triton 3.7.
+
+from pkgutil import extend_path
+
+__path__ = extend_path(__path__, __name__)
+__version__ = "0.5.2"
+__all__: list[str] = []

windows_fla_patches/fla/layers/__init__.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# Lazy layer exports — avoid compiling every Triton kernel at import time.
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+_EXPORTS: dict[str, tuple[str, str]] = {
+    "ABCAttention": (".abc", "ABCAttention"),
+    "Attention": (".attn", "Attention"),
+    "BasedLinearAttention": (".based", "BasedLinearAttention"),
+    "BitAttention": (".bitattn", "BitAttention"),
+    "Comba": (".comba", "Comba"),
+    "DeltaNet": (".delta_net", "DeltaNet"),
+    "DeltaFormerAttention": (".deltaformer", "DeltaFormerAttention"),
+    "ForgettingAttention": (".forgetting_attn", "ForgettingAttention"),
+    "GatedDeltaNet": (".gated_deltanet", "GatedDeltaNet"),
+    "GatedDeltaProduct": (".gated_deltaproduct", "GatedDeltaProduct"),
+    "GatedDeltaNet2": (".gdn2", "GatedDeltaNet2"),
+    "GatedLinearAttention": (".gla", "GatedLinearAttention"),
+    "GatedSlotAttention": (".gsa", "GatedSlotAttention"),
+    "HGRNAttention": (".hgrn", "HGRNAttention"),
+    "HGRN2Attention": (".hgrn2", "HGRN2Attention"),
+    "KimiDeltaAttention": (".kda", "KimiDeltaAttention"),
+    "LightNetAttention": (".lightnet", "LightNetAttention"),
+    "LinearAttention": (".linear_attn", "LinearAttention"),
+    "LogLinearMamba2": (".log_linear_mamba2", "LogLinearMamba2"),
+    "Mamba": (".mamba", "Mamba"),
+    "Mamba2": (".mamba2", "Mamba2"),
+    "Mamba3": (".mamba3", "Mamba3"),
+    "MesaNet": (".mesa_net", "MesaNet"),
+    "MultiheadLatentAttention": (".mla", "MultiheadLatentAttention"),
+    "MoBA": (".moba", "MoBA"),
+    "MomAttention": (".mom", "MomAttention"),
+    "MultiScaleRetention": (".multiscale_retention", "MultiScaleRetention"),
+    "NativeSparseAttention": (".nsa", "NativeSparseAttention"),
+    "Parallax": (".parallax", "Parallax"),
+    "PaTHAttention": (".path_attn", "PaTHAttention"),
+    "Raven": (".raven", "Raven"),
+    "ReBasedLinearAttention": (".rebased", "ReBasedLinearAttention"),
+    "RodimusAttention": (".rodimus", "RodimusAttention"),
+    "SlidingWindowSharedKeyAttention": (".rodimus", "SlidingWindowSharedKeyAttention"),
+    "RWKV6Attention": (".rwkv6", "RWKV6Attention"),
+    "RWKV7Attention": (".rwkv7", "RWKV7Attention"),
+    "WallAttention": (".wall_attn", "WallAttention"),
+    "YOCOCrossAttention": (".yoco", "YOCOCrossAttention"),
+    "YOCOGatedRetention": (".yoco", "YOCOGatedRetention"),
+    "YOCOSharedKVBuilder": (".yoco", "YOCOSharedKVBuilder"),
+}
+
+__all__ = list(_EXPORTS)
+
+
+def __getattr__(name: str) -> Any:
+    spec = _EXPORTS.get(name)
+    if spec is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module_name, attr = spec
+    value = getattr(importlib.import_module(module_name, __name__), attr)
+    globals()[name] = value
+    return value

windows_fla_patches/fla/ops/__init__.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# Lazy public exports so importing fla.ops.utils / fla.ops.gdn2 does not
+# eagerly compile every Triton kernel (needed on Windows + Triton 3.7).
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+_EXPORTS: dict[str, str] = {
+    "chunk_abc": "fla.ops.abc",
+    "parallel_attn": "fla.ops.attn",
+    "fused_attnres": "fla.ops.attnres",
+    "fused_chunk_based": "fla.ops.based",
+    "parallel_based": "fla.ops.based",
+    "chunk_comba": "fla.ops.comba",
+    "fused_recurrent_comba": "fla.ops.comba",
+    "chunk_delta_rule": "fla.ops.delta_rule",
+    "fused_chunk_delta_rule": "fla.ops.delta_rule",
+    "fused_recurrent_delta_rule": "fla.ops.delta_rule",
+    "parallel_forgetting_attn": "fla.ops.forgetting_attn",
+    "chunk_gated_delta_rule": "fla.ops.gated_delta_rule",
+    "chunk_gdn": "fla.ops.gated_delta_rule",
+    "fused_recurrent_gated_delta_rule": "fla.ops.gated_delta_rule",
+    "fused_recurrent_gdn": "fla.ops.gated_delta_rule",
+    "chunk_dplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "chunk_iplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "fused_recurrent_dplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "fused_recurrent_iplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "chunk_gla": "fla.ops.gla",
+    "fused_chunk_gla": "fla.ops.gla",
+    "fused_recurrent_gla": "fla.ops.gla",
+    "chunk_gsa": "fla.ops.gsa",
+    "fused_recurrent_gsa": "fla.ops.gsa",
+    "fused_recurrent_hgrn": "fla.ops.hgrn",
+    "chunk_kda": "fla.ops.kda",
+    "fused_recurrent_kda": "fla.ops.kda",
+    "chunk_lightning_attn": "fla.ops.lightning_attn",
+    "fused_recurrent_lightning_attn": "fla.ops.lightning_attn",
+    "chunk_linear_attn": "fla.ops.linear_attn",
+    "fused_chunk_linear_attn": "fla.ops.linear_attn",
+    "fused_recurrent_linear_attn": "fla.ops.linear_attn",
+    "chunk_log_linear_attn": "fla.ops.log_linear_attn",
+    "chunk_mesa_net": "fla.ops.mesa_net",
+    "parallel_nsa": "fla.ops.nsa",
+    "parallel_parallax": "fla.ops.parallax",
+    "parallel_path_attn": "fla.ops.path_attn",
+    "chunk_retention": "fla.ops.retention",
+    "fused_chunk_retention": "fla.ops.retention",
+    "fused_recurrent_retention": "fla.ops.retention",
+    "parallel_retention": "fla.ops.retention",
+    "chunk_rwkv6": "fla.ops.rwkv6",
+    "fused_recurrent_rwkv6": "fla.ops.rwkv6",
+    "chunk_rwkv7": "fla.ops.rwkv7",
+    "fused_recurrent_rwkv7": "fla.ops.rwkv7",
+    "chunk_simple_gla": "fla.ops.simple_gla",
+    "fused_chunk_simple_gla": "fla.ops.simple_gla",
+    "fused_recurrent_simple_gla": "fla.ops.simple_gla",
+    "parallel_simple_gla": "fla.ops.simple_gla",
+    "parallel_wall_attn": "fla.ops.wall_attn",
+    "parallel_wall_attn_decode": "fla.ops.wall_attn",
+    "chunk_gdn2": "fla.ops.gdn2",
+    "fused_recurrent_gdn2": "fla.ops.gdn2",
+}
+
+__all__ = list(_EXPORTS)
+
+
+def __getattr__(name: str) -> Any:
+    module_name = _EXPORTS.get(name)
+    if module_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    value = getattr(importlib.import_module(module_name), name)
+    globals()[name] = value
+    return value

windows_fla_patches/fla/ops/simple_gla/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+from .chunk import chunk_simple_gla
+from .fused_chunk import fused_chunk_simple_gla
+from .fused_recurrent import fused_recurrent_simple_gla
+
+# Triton 3.7 on Windows can fail while decorating parallel kernels at import time.
+try:
+    from .parallel import parallel_simple_gla
+except Exception:  # noqa: BLE001
+    parallel_simple_gla = None
+
+__all__ = [
+    'chunk_simple_gla',
+    'fused_chunk_simple_gla',
+    'fused_recurrent_simple_gla',
+    'parallel_simple_gla',
+]