Upload Senryu (apo_zero apo-v2 EMA @ step 704, from SFT step 12800)

README.md

@@ -0,0 +1,222 @@
+---
+license: apache-2.0
+language:
+  - en
+tags:
+  - text-generation
+  - causal-lm
+  - pytorch
+  - sft
+  - preference-optimization
+  - apo
+  - on-policy
+  - instruction-tuned
+  - chat
+  - reasoning
+  - tool-use
+  - hybrid
+  - gated-deltanet-2
+  - gated-mla
+  - senryu
+pipeline_tag: text-generation
+library_name: tiny_gdn
+datasets:
+  - HuggingFaceFW/fineweb-edu
+  - HuggingFaceTB/smoltalk2
+  - HuggingFaceTB/smoltalk
+  - nvidia/Nemotron-Post-Training-Dataset-v1
+  - nvidia/Nemotron-Agentic-v1
+  - nvidia/Nemotron-SFT-Agentic-v2
+  - nvidia/Nemotron-SFT-Instruction-Following-Chat-v2
+  - nvidia/Nemotron-Instruction-Following-Chat-v1
+  - nvidia/OpenMathReasoning
+  - nvidia/OpenCodeReasoning
+  - nvidia/When2Call
+  - Agent-Ark/Toucan-1.5M
+  - interstellarninja/hermes_reasoning_tool_use
+  - interstellarninja/tool-use-multiturn-reasoning
+  - NousResearch/Hermes-3-Dataset
+  - glaiveai/glaive-function-calling-v2
+  - allenai/WildChat-4.8M
+  - QuixiAI/SystemChat-2.0
+  - Open-Orca/SlimOrca
+  - HuggingFaceH4/ultrachat_200k
+  - ise-uiuc/Magicoder-Evol-Instruct-110K
+  - bigcode/self-oss-instruct-sc2-exec-filter-50k
+  - OpenAssistant/oasst2
+  - allenai/tulu-3-sft-personas-instruction-following
+  - HuggingFaceH4/no_robots
+---
+
+<div align="center">
+
+# Senryu
+
+### Reasoning + tool-call chat model (803M), preference-tuned on its own samples
+
+[![Model](https://img.shields.io/badge/Model-803M_params-blue)](.)
+[![Stage](https://img.shields.io/badge/Stage-SFT_+_on--policy_APO-green.svg)](.)
+[![License](https://img.shields.io/badge/License-Apache_2.0-green.svg)](LICENSE)
+[![Architecture](https://img.shields.io/badge/Arch-Gated_DeltaNet--2_+_Gated_MLA-purple.svg)](.)
+[![Demo](https://img.shields.io/badge/Space-senryu--demo-indigo.svg)](https://huggingface.co/spaces/kerzgrr/senryu-demo)
+
+*An 803M hybrid Gated DeltaNet-2 + gated MLA model: 32k-context SFT on reasoning, tool and chat data, then APO-zero on pairs sampled from the model itself*
+
+</div>
+
+---
+
+## What this is
+
+**Senryu** is the 803M successor to [Haiku](https://huggingface.co/kerzgrr/Haiku): eight more layers, Gated DeltaNet-2 in place of Kimi Delta Attention, and a full-rank MLA query.
+
+- Pretraining on **1.15B** FineWeb-Edu tokens
+- Supervised fine-tuning at a **32,768-token** sequence length on 6.7B tokens (5.39B assistant tokens) of reasoning, tool-call, instruction-following and chat data
+- Then **APO-zero** on **22,517 on-policy preference pairs**: four samples per prompt from the SFT model itself, judged by [Skywork-Reward-V2-Llama-3.1-8B](https://huggingface.co/Skywork/Skywork-Reward-V2-Llama-3.1-8B) and rule checks for tool calls, answers and loops
+- Live demo: [`kerzgrr/senryu-demo`](https://huggingface.co/spaces/kerzgrr/senryu-demo)
+- Hub weights are the **EMA** snapshot in **bfloat16**
+
+This upload is the final checkpoint of the preference run (optimizer step **704**, one full epoch of 22,517 pairs). It is the only released checkpoint of the Senryu generation; the pretrain and SFT checkpoints behind it are not published.
+
+---
+
+## Chat contract
+
+### Thinking
+
+Each assistant turn is prefixed with a **zero-loss control token**:
+
+| Mode | Prefix | Typical body |
+|------|--------|----------------|
+| think | `<\|think\|>\n` | `<think>…</think>` then the answer |
+| no-think | `<\|no_think\|>\n` | answer only |
+
+`inference.py` streams the `<think>` region live (dim yellow) and hides the control tokens.
+
+### Tool calls (Hermes JSON)
+
+Tool schemas go in a Hermes-style `<tools>` block in the system turn. The model calls them with:
+
+```
+<tool_call>
+{"name": "web-search", "arguments": {"query": "…"}}
+</tool_call>
+```
+
+Senryu learned schema-following calls from Glaive function calling v2, When2Call, Hermes reasoning tool-use, SmolTalk and the stage-3 tool mix (Toucan 1.5M, Nemotron PTD v1, Nemotron Agentic v1/v2). `web-search` is the only built-in it saw by name, in the Nemotron Agentic conversations; to Senryu, `stateful_python_code_exec` and `calculator` are ordinary schemas it reads from the `<tools>` block. On prompts that offer tools, the preference pairs reward calling and answering directly equally often (6,794 pairs each). `inference.py` auto-runs:
+
+| Built-in | Tool name | Observation |
+|----------|-----------|-------------|
+| `--tools web-search` | `web-search` | Tavily-shaped JSON (Tavily if `TAVILY_API_KEY` is set, else DuckDuckGo + Wikipedia) |
+| `--tools python` | `stateful_python_code_exec` | Jupyter-style stdout / last value from a restricted interpreter |
+| `--tools calculator` | `calculator` | Numeric result of a math expression |
+
+### Tool results
+
+Each observation is a `tool` (or `user`) turn prefixed with:
+
+```
+<|tool_response|>
+{observation}
+```
+
+---
+
+## Install & run
+
+```bash
+pip install torch safetensors tokenizers huggingface_hub
+hf download kerzgrr/Senryu inference.py --local-dir .
+python inference.py
+```
+
+```bash
+python inference.py --prompt "What is the capital of France?"
+python inference.py --tools web-search,python,calculator
+python inference.py --no-think --prompt "Reply in one sentence."
+```
+
+`inference.py` auto-downloads weights / tokenizer / `tiny_gdn/` and auto-installs pinned `flash-linear-attention` (plus `transformers`, which its decode cache imports). Git is required on `PATH`. A CUDA GPU is required: the Gated DeltaNet-2 layers run on Triton kernels.
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--prompt` | — | One-shot user message |
+| `--system` | — | System prompt, used verbatim |
+| `--think` / `--no-think` | think | Assistant control prefix |
+| `--tools` | — | Built-ins: `web-search`, `python`, `calculator` (comma-separated) |
+| `--temperature` | `0.6` | Sampling temperature (0.3 suits tool-calling turns) |
+| `--top-p` | `0.95` | Nucleus sampling |
+| `--repetition-penalty` | `1.08` | Over the last 256 tokens |
+| `--max-new-tokens` | `4096` | Max generation length |
+| `--context-length` | `32768` | Conversation tokens kept (oldest dropped first) |
+| `--device` | `cuda` | CUDA device, e.g. `cuda:1` |
+
+Interactive commands: `/think` `/no_think` `/system …` `/reset` `/exit`. Ctrl+C stops the current reply and keeps the session.
+
+---
+
+## Model architecture
+
+803,000,968 parameters:
+
+| | |
+|--|--|
+| **Layers** | 44 (Gated DeltaNet-2 ×3 + gated MLA every 4th) |
+| **Hidden** | 1,024 |
+| **MLP** | SiTU-GLU 3,392 |
+| **Attention** | Gated MLA (NoPE), 8 heads, full-rank query, KV LoRA rank 512 |
+| **Linear** | Gated DeltaNet-2, 8 heads × 128, channel-wise erase and write gates |
+| **Residuals** | Block attention residuals |
+| **Vocab** | 65,536 BPE (the Haiku tokenizer) |
+| **Context** | 32,768 (SFT sequence length) |
+
+---
+
+## Training
+
+| Stage | Details |
+|-------|---------|
+| **Pretrain** | [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu), seq 2,048, 2,200 steps = 1.15B tokens, hybrid Muon + AdamW at peak 5×10⁻⁴, 10.4 hours. Val loss (EMA) 3.06 |
+| **SFT** | Natural-proportion mixture with a full shuffle: [SmolTalk2](https://huggingface.co/datasets/HuggingFaceTB/smoltalk2) Mid + SFT, the stage-3 tool/agentic mix (Nemotron PTD v1, Agentic v1/v2, IF-Chat v1/v2, Toucan 1.5M), OpenMathReasoning, OpenCodeReasoning, WildChat-4.8M, SmolTalk, Hermes-3, Hermes reasoning tool-use, SystemChat 2.0, SlimOrca, UltraChat 200k, Magicoder Evol-Instruct, Self-OSS-Instruct, Glaive function calling v2, When2Call, OASST2, Tulu 3 IF personas, No Robots. Seq 32,768, hybrid Muon + AdamW 4×10⁻⁴, 12,800 steps = 6.7B tokens (5.39B assistant tokens), 11.8 hours. Val loss (EMA) 0.827 (ppl 2.29) |
+| **Preference (APO)** | APO-zero, β = 0.05, on on-policy pairs. 20,000 prompts drawn from the SFT mixture, four samples each from the SFT EMA (temperature 0.8, top-p 0.95). The chosen reply is the best-judged sample, or the dataset's reference reply when no sample passed; pairs need a reward margin of at least 3. Prompts that offer tools have 6,794 call pairs and 6,794 answer-directly pairs. 3,697 loop-stop pairs prefer ending the turn exactly where a repeat begins over one more repeat. 22,517 pairs, one epoch, 32 pairs per step, 704 steps, AdamW 1×10⁻⁶ cosine to 1×10⁻⁷, 10% warmup, clip 0.2, max length 16,512. Reference model: the SFT EMA. 1.8 hours |
+| **Weights** | EMA (this repo's `model.safetensors`) |
+| **Preference val (EMA)** | 903 held-out pairs: reward accuracy 79.0%, reward margin 1.33, chosen NLL 0.56 per token |
+
+`preference_config.json`, `validation.json`, `sft_config.json` and `sft_validation.json` hold the exact settings and held-out results. `preference_data.json` describes the pair pack: prompt selection, sampling, judging, and the composition by source and label.
+
+---
+
+## Limitations
+
+- **Scale**: 803M is a research / edge model, not a frontier system
+---
+
+## Model family
+
+| Model | Stage | Hub |
+|-------|-------|-----|
+| **Senryu** | **SFT + on-policy APO (chat, reasoning, tools)** | **this repo** |
+| Demo | ZeroGPU Space (chat) | [`kerzgrr/senryu-demo`](https://huggingface.co/spaces/kerzgrr/senryu-demo) |
+| Haiku | Previous generation (~655M) | [`kerzgrr/Haiku`](https://huggingface.co/kerzgrr/Haiku) |
+| Haiku-base | Previous generation, pretrain | [`kerzgrr/Haiku-base`](https://huggingface.co/kerzgrr/Haiku-base) |
+
+---
+
+## Citation
+
+```bibtex
+@misc{senryu2026,
+  title={Senryu: An 803M Hybrid Gated DeltaNet-2 + Gated-MLA Reasoning Model},
+  author={kerzgrr},
+  year={2026},
+  url={https://huggingface.co/kerzgrr/Senryu}
+}
+```
+
+---
+
+<div align="center">
+
+*Five, seven, five — about people this time.*
+
+</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,126 @@
+{
+  "allow_negative_eigenvalues": false,
+  "architecture": "SenryuForCausalLM",
+  "architectures": [
+    "SenryuForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "attention_head_dim": 128,
+  "attention_layer_type": "gated_mla",
+  "bos_token_id": 0,
+  "chat_template": "smollm3",
+  "checkpoint_step": 704,
+  "eos_token_id": 1,
+  "full_attention_interval": 4,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 3392,
+  "kda_lower_bound": -5.0,
+  "kda_safe_gate": true,
+  "layer_types": [
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "gated_mla",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "gated_mla",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "gated_mla",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "gated_mla",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "gated_mla",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "gated_mla",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "gated_mla",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "gated_mla",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "gated_mla",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "gated_mla",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "gated_mla"
+  ],
+  "linear_conv_kernel_dim": 4,
+  "linear_expand_v": 1.0,
+  "linear_head_dim": 128,
+  "linear_layer_type": "gdn2",
+  "linear_num_heads": 8,
+  "linear_num_value_heads": 8,
+  "max_position_embeddings": 32768,
+  "mla_kv_lora_rank": 512,
+  "mla_q_lora_rank": null,
+  "mla_qk_nope_head_dim": 128,
+  "mla_v_head_dim": 128,
+  "mlp_activation": "situ_glu",
+  "model_family": "Senryu",
+  "model_name": "Senryu",
+  "model_type": "senryu",
+  "mtp_adapter_rank": 128,
+  "mtp_loss_weight": 0.0,
+  "mtp_num_heads": 0,
+  "num_attention_heads": 8,
+  "num_hidden_layers": 44,
+  "num_key_value_heads": 8,
+  "pad_token_id": 2,
+  "partial_rotary_factor": 0.5,
+  "preference_beta": 0.05,
+  "preference_hours": 1.7632968792019852,
+  "preference_loss": "apo_zero",
+  "preference_reward_model": "Skywork/Skywork-Reward-V2-Llama-3.1-8B",
+  "preference_run": "apo-v2",
+  "preference_train_pairs": 22517,
+  "preference_val_loss_ema": 0.8922630808546166,
+  "preference_val_pairs": 903,
+  "preference_val_reward_accuracy_ema": 0.7895902547065338,
+  "preference_val_reward_margin_ema": 1.3302929308987403,
+  "pretrain_checkpoint_step": 2200,
+  "pretrain_hours": 10.380536459482197,
+  "pretrain_run": "pretrain-v1",
+  "pretrain_tokens": 1153433600,
+  "pretrain_val_loss_ema": 3.06275936961174,
+  "reasoning": true,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000.0,
+  "sft_assistant_tokens": 5388440922,
+  "sft_checkpoint_step": 12800,
+  "sft_hours": 11.803919963316359,
+  "sft_run": "sft-v1",
+  "sft_val_loss_ema": 0.8266990829440273,
+  "sft_val_ppl_ema": 2.2857611656454533,
+  "shared_layer_indices": [],
+  "situ_gate_cap": 4.0,
+  "situ_up_cap": 25.0,
+  "stage": "preference",
+  "tie_word_embeddings": true,
+  "tool_calls": "hermes_json",
+  "torch_dtype": "bfloat16",
+  "training_sequence_length": 32768,
+  "transformers_version": "4.45.0",
+  "unk_token_id": 3,
+  "use_block_attn_res": true,
+  "vocab_size": 65536,
+  "weights": "ema"
+}

inference.py

@@ -0,0 +1,593 @@
+#!/usr/bin/env python3
+"""Standalone chat inference for kerzgrr/Senryu.
+
+Streams a live reasoning trace (`<think>` / `<|think|>`) and runs Hermes JSON
+tool-call rounds. Built-in `web-search`, `python`, and `calculator` tools
+execute automatically. Other observations are sent as `<|tool_response|>`.
+
+Examples:
+  python inference.py --prompt "What is the capital of France?"
+  python inference.py --tools web-search,python,calculator
+  python inference.py --no-think --prompt "Reply in one sentence."
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+import os
+import platform
+import shutil
+import subprocess
+import sys
+import time
+import warnings
+from dataclasses import fields
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+import torch
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from tokenizers import Tokenizer
+
+
+def _silence_runtime_warnings() -> None:
+    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/Senryu"
+FLA_COMMIT = "cbb0a72efb55c18ca0ef4f298298317573ad2cb3"
+FLA_REPO = "https://github.com/fla-org/flash-linear-attention.git"
+FLA_RUNTIME_REQUIREMENT = "transformers>=4.45.0"
+PATCH_FILES = (
+    "fla/__init__.py",
+    "fla/ops/__init__.py",
+    "fla/layers/__init__.py",
+    "fla/ops/simple_gla/__init__.py",
+)
+TINY_GDN_FILES = (
+    "tiny_gdn/__init__.py",
+    "tiny_gdn/config.py",
+    "tiny_gdn/model.py",
+    "tiny_gdn/haiku_layers.py",
+    "tiny_gdn/nn_common.py",
+    "tiny_gdn/fused_ops.py",
+    "tiny_gdn/tools.py",
+    "tiny_gdn/smoltalk_chat.py",
+    "tiny_gdn/chatml.py",
+    "tiny_gdn/cli_chat.py",
+    "tiny_gdn/detokenize.py",
+    "tiny_gdn/code_exec.py",
+    "tiny_gdn/web_search.py",
+)
+
+
+def _download(filename: str, local_dir: Path | None) -> Path:
+    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)
+    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]:
+    # fla.models.utils holds the decode cache and is the module that needs
+    # transformers at runtime; the GDN-2 layer alone imports without it.
+    try:
+        importlib.import_module("fla.layers.gdn2")
+        importlib.import_module("fla.models.utils")
+    except Exception as error:  # noqa: BLE001
+        return False, str(error)
+    return True, ""
+
+
+def _cache_root() -> Path:
+    override = os.environ.get("SENRYU_CACHE") or os.environ.get("MONOSTICH_CACHE")
+    if override:
+        path = Path(override).expanduser().resolve()
+    else:
+        path = Path.home() / ".cache" / "senryu"
+    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", FLA_RUNTIME_REQUIREMENT)
+    if platform.system() != "Windows":
+        _pip_install("--no-deps", f"git+{FLA_REPO}@{FLA_COMMIT}")
+        return
+
+    _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
+
+    for name in list(sys.modules):
+        if name == "fla" or name.startswith("fla."):
+            del sys.modules[name]
+    importlib.invalidate_caches()
+
+    ok, error = _fla_importable()
+    if not ok:
+        raise RuntimeError(
+            "flash-linear-attention installed but still failed to import "
+            f"GatedDeltaNet2 and its decode cache: {error}"
+        )
+    print("flash-linear-attention ready.", flush=True)
+
+
+def _has_tiny_gdn(root: Path) -> bool:
+    return all((root / name).is_file() for name in TINY_GDN_FILES)
+
+
+def _ensure_tiny_gdn(local_dir: Path | None) -> Path:
+    here = Path(__file__).resolve().parent
+    if _has_tiny_gdn(here):
+        return here
+    if local_dir and _has_tiny_gdn(local_dir):
+        return local_dir
+    for name in TINY_GDN_FILES:
+        _download(name, local_dir)
+    return _download("tiny_gdn/__init__.py", local_dir).parent.parent
+
+
+def _import_chat_stack() -> SimpleNamespace:
+    # tiny_gdn is fetched from the Hub and FLA may be installed at runtime, so
+    # these modules only become importable after _ensure_tiny_gdn/_ensure_fla.
+    tiny_gdn = importlib.import_module("tiny_gdn")
+    chatml = importlib.import_module("tiny_gdn.chatml")
+    cli_chat = importlib.import_module("tiny_gdn.cli_chat")
+    detokenize = importlib.import_module("tiny_gdn.detokenize")
+    tools = importlib.import_module("tiny_gdn.tools")
+    return SimpleNamespace(
+        TinyGDNConfig=tiny_gdn.TinyGDNConfig,
+        TinyGDNForCausalLM=tiny_gdn.TinyGDNForCausalLM,
+        encode_chatml_messages=chatml.encode_chatml_messages,
+        ChatLoopState=cli_chat.ChatLoopState,
+        LiveReasoningStreamer=cli_chat.LiveReasoningStreamer,
+        append_tool_round=cli_chat.append_tool_round,
+        apply_slash_command=cli_chat.apply_slash_command,
+        prompt_tool_results=cli_chat.prompt_tool_results,
+        resolve_cli_tools=cli_chat.resolve_cli_tools,
+        IncrementalUtf8Decoder=detokenize.IncrementalUtf8Decoder,
+        parse_tool_calls=tools.parse_tool_calls,
+    )
+
+
+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).cpu()
+    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_turn(
+    model,
+    tokenizer: Tokenizer,
+    stack: SimpleNamespace,
+    streamer,
+    messages: list[dict[str, Any]],
+    *,
+    tools: Any | None,
+    enable_thinking: bool,
+    max_new_tokens: int,
+    context_length: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    repetition_penalty: float,
+    repetition_window: int,
+    seed: int,
+    device: torch.device,
+) -> tuple[str, tuple]:
+    prompt_ids = stack.encode_chatml_messages(
+        tokenizer,
+        messages,
+        tools=tools,
+        enable_thinking=enable_thinking,
+    )
+    if len(prompt_ids) > context_length:
+        print(
+            f"[warn] conversation is {len(prompt_ids)} tokens; "
+            f"keeping the last {context_length} (--context-length)",
+            file=sys.stderr,
+            flush=True,
+        )
+    eos_id = int(model.config.eos_token_id)
+    im_end = tokenizer.token_to_id("<|im_end|>")
+    stop = {eos_id}
+    if im_end is not None:
+        stop.add(im_end)
+
+    token_ids = list(prompt_ids[-context_length:])
+    generated: list[int] = []
+    decoder = stack.IncrementalUtf8Decoder(tokenizer)
+    stop_reason = "max_new_tokens"
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(seed)
+    started = time.perf_counter()
+
+    try:
+        input_ids = torch.tensor([token_ids], dtype=torch.long, device=device)
+        output = model(
+            input_ids,
+            return_logits=True,
+            logits_to_keep=1,
+            use_cache=True,
+        )
+
+        for step in range(max_new_tokens):
+            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 in stop:
+                stop_reason = "stop"
+                break
+            token_ids.append(next_id)
+            generated.append(next_id)
+            decoder.push(next_id)
+            if streamer is not None:
+                streamer.update(decoder.text)
+            if step + 1 < max_new_tokens:
+                if output.past_key_values is None:
+                    raise RuntimeError("Model returned no decode cache")
+                input_ids = torch.tensor([[next_id]], dtype=torch.long, device=device)
+                output = model(
+                    input_ids,
+                    return_logits=True,
+                    logits_to_keep=1,
+                    past_key_values=output.past_key_values,
+                    use_cache=True,
+                )
+    except KeyboardInterrupt:
+        stop_reason = "interrupted"
+
+    decoded = decoder.finalize()
+    if streamer is not None:
+        streamer.finish()
+    elapsed = time.perf_counter() - started
+    tps = len(generated) / max(elapsed, 1e-9)
+    print(
+        f"[done] tokens={len(generated)} stop={stop_reason} {tps:.1f} tok/s",
+        file=sys.stderr,
+        flush=True,
+    )
+    if stop_reason == "interrupted":
+        return decoded, ()
+    return decoded, tuple(stack.parse_tool_calls(decoded))
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Senryu chat inference")
+    parser.add_argument("--prompt", default=None, help="Single user prompt")
+    parser.add_argument("--system", default="", help="Optional system prompt")
+    parser.add_argument(
+        "--think",
+        dest="enable_thinking",
+        action="store_true",
+        default=True,
+    )
+    parser.add_argument(
+        "--no-think",
+        dest="enable_thinking",
+        action="store_false",
+    )
+    parser.add_argument(
+        "--tools",
+        default="",
+        help="Built-in tools: web-search, python, calculator (comma-separated)",
+    )
+    parser.add_argument("--tools-json", default=None)
+    parser.add_argument("--max-tool-rounds", type=int, default=8)
+    parser.add_argument("--max-new-tokens", type=int, default=4096)
+    parser.add_argument("--temperature", type=float, default=0.6)
+    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=32768)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--device",
+        default="cuda",
+        help="CUDA device (cuda, cuda:1, …); Senryu needs a GPU",
+    )
+    parser.add_argument("--no-stream", action="store_true")
+    parser.add_argument("--no-color", action="store_true")
+    parser.add_argument("--local-dir", default=None)
+    parser.add_argument("--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
+
+    try:
+        device = torch.device(args.device)
+    except RuntimeError as error:
+        print(f"Invalid --device {args.device!r}: {error}", file=sys.stderr)
+        return 1
+    if device.type != "cuda" or not torch.cuda.is_available():
+        print(
+            "Senryu needs a CUDA GPU: its Gated DeltaNet-2 layers run on "
+            "flash-linear-attention Triton kernels. "
+            f"torch {torch.__version__} reports CUDA available={torch.cuda.is_available()}.",
+            file=sys.stderr,
+        )
+        return 1
+
+    print(f"Loading Senryu 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)
+        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:
+        stack = _import_chat_stack()
+    except ImportError as error:
+        print(f"Could not import tiny_gdn chat stack: {error}", file=sys.stderr)
+        return 1
+
+    try:
+        raw_json = (
+            Path(args.tools_json).read_text(encoding="utf-8")
+            if args.tools_json
+            else None
+        )
+        tools = stack.resolve_cli_tools(args.tools, raw_json)
+    except (OSError, ValueError, json.JSONDecodeError) as error:
+        print(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"))
+    allowed = {item.name for item in fields(stack.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 = stack.TinyGDNConfig(**payload)
+
+    print(
+        f"Building Senryu ({config.num_hidden_layers}L / {config.hidden_size}d) "
+        f"on {device} …",
+        flush=True,
+    )
+    model = stack.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=torch.bfloat16)
+    model.eval()
+    model.requires_grad_(False)
+
+    tokenizer = Tokenizer.from_file(str(tok_path))
+    context_length = min(args.context_length, config.max_position_embeddings)
+    color = not args.no_color and sys.stdout.isatty()
+
+    def read_line(prompt: str) -> str:
+        return input(prompt)
+
+    def run_turn(state) -> None:
+        for round_index in range(max(1, args.max_tool_rounds + 1)):
+            streamer = (
+                None
+                if args.no_stream
+                else stack.LiveReasoningStreamer(sys.stdout, color=color)
+            )
+            if streamer is not None:
+                print("assistant> ", end="", flush=True)
+            text, calls = generate_turn(
+                model,
+                tokenizer,
+                stack,
+                streamer,
+                state.messages,
+                tools=state.tools,
+                enable_thinking=state.enable_thinking,
+                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,
+                device=device,
+            )
+            if args.no_stream:
+                print(f"assistant> {text}")
+            if not calls or round_index >= args.max_tool_rounds:
+                state.messages.append({"role": "assistant", "content": text})
+                return
+            results = stack.prompt_tool_results(calls, read_line=read_line)
+            stack.append_tool_round(state.messages, text, results)
+
+    state = stack.ChatLoopState(
+        system=args.system,
+        enable_thinking=args.enable_thinking,
+        tools=tools,
+    )
+    state.reset()
+
+    if args.prompt is not None:
+        state.add_user(args.prompt)
+        try:
+            run_turn(state)
+        except KeyboardInterrupt:
+            print("\n(interrupted)", file=sys.stderr)
+            return 130
+        return 0
+
+    print(
+        "Interactive Senryu chat. "
+        "Commands: /think  /no_think  /system …  /reset  /exit\n"
+        "web-search, python, and calculator run automatically when enabled. "
+        "Other <tool_call> observations can be pasted at the prompt. "
+        "Ctrl+C stops a reply.",
+        flush=True,
+    )
+    while True:
+        try:
+            user_input = input("user> ")
+        except (EOFError, KeyboardInterrupt):
+            print()
+            return 0
+        text = user_input.strip()
+        if not text:
+            continue
+        status = stack.apply_slash_command(state, text)
+        if status == "exit":
+            return 0
+        if status is not None:
+            print(f"({status})", flush=True)
+            continue
+        state.add_user(text)
+        try:
+            run_turn(state)
+        except KeyboardInterrupt:
+            print("\n(interrupted)", flush=True)
+        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:57f7960df1a6e4b9cee141de3a909de43790afd51d2ea8258f34e67847936668
+size 1606092704

preference_config.json

@@ -0,0 +1,43 @@
+{
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.999,
+  "adam_epsilon": 1e-08,
+  "beta": 0.05,
+  "dataset_manifest": "/media/nvidia/4TB2/llm-data/senryu/rl-v1/pack-v2/manifest.json",
+  "ema_inv_gamma": 1.0,
+  "ema_max_decay": 0.9999,
+  "ema_power": 0.75,
+  "enable_gradient_checkpointing": true,
+  "epochs": 1,
+  "forward_overhead_tokens": 1024,
+  "fp32_master_weights": true,
+  "gradient_checkpointing_layers": 22,
+  "gradient_clip_norm": 0.2,
+  "initial_checkpoint": "checkpoint-00012800",
+  "initial_run_dir": "/media/nvidia/4TB2/llm-runs/senryu/sft-v1",
+  "initial_weights": "ema",
+  "label_smoothing": 0.0,
+  "learning_rate": 1e-06,
+  "log_every_optimizer_steps": 1,
+  "loss_type": "apo_zero",
+  "max_length": 16512,
+  "maximum_checkpoints": 5,
+  "maximum_optimizer_steps": null,
+  "micro_batch_tokens": 32768,
+  "minimum_learning_rate_ratio": 0.1,
+  "muon_learning_rate": null,
+  "muon_lr_adjustment": "match_rms_adamw",
+  "muon_momentum": 0.95,
+  "optimizer": "adamw",
+  "output_dir": "/media/nvidia/4TB2/llm-runs/senryu/apo-v2",
+  "pairs_per_step": 32,
+  "reference_batch_tokens": 32768,
+  "save_every_optimizer_steps": 100,
+  "seed": 42,
+  "sft_weight": 0.0,
+  "sources": null,
+  "validation_pairs": 1000,
+  "warmup_ratio": 0.1,
+  "warmup_steps": null,
+  "weight_decay": 0.0
+}

preference_data.json

@@ -0,0 +1,287 @@
+{
+  "kind": "on-policy preference pairs from SFT-mixture prompts",
+  "prompt_selection": {
+    "excluded_components": [
+      "nemotron-cascade-2"
+    ],
+    "settings": {
+      "prompts": 20000,
+      "offered_share": 0.3333333333333333,
+      "floor": 40,
+      "min_weight": 1000,
+      "max_prompt_tokens": 8192,
+      "max_reference_tokens": 8192,
+      "validation_permyriad": 400,
+      "max_scan": 600000,
+      "seed": 20261018
+    }
+  },
+  "policy": {
+    "run": "sft-v1",
+    "checkpoint": "checkpoint-00012800",
+    "weights": "ema"
+  },
+  "samples_per_prompt": 4,
+  "sampling": {
+    "temperature": 0.8,
+    "top_p": 0.95,
+    "top_k": 0,
+    "think_tokens": 8192,
+    "direct_tokens": 3072,
+    "loop_ngram": 16,
+    "loop_window": 1536,
+    "loop_threshold": 0.5,
+    "seed": 20261018
+  },
+  "reward_model": "Skywork/Skywork-Reward-V2-Llama-3.1-8B",
+  "pairing": {
+    "margin": 3.0,
+    "max_length_ratio": 2.0,
+    "max_group_ratio": 1.25,
+    "seed": 20261018,
+    "tool_abstention": {
+      "train": {
+        "tool_call_pairs": 6794,
+        "tool_abstain_pairs": 6794,
+        "abstentions_added": 4328
+      },
+      "validation": {
+        "tool_call_pairs": 277,
+        "tool_abstain_pairs": 277,
+        "abstentions_added": 182
+      }
+    },
+    "loop_stops": true
+  },
+  "sources": [
+    "glaive-fc/direct-answer",
+    "glaive-fc/direct-answer+tools",
+    "glaive-fc/direct-call",
+    "hermes-3/direct-answer",
+    "hermes-tool-use/reason-answer+tools",
+    "hermes-tool-use/reason-call",
+    "magicoder/direct-answer",
+    "no-robots/direct-answer",
+    "oasst2/direct-answer",
+    "opencode/reason-answer",
+    "openmath/reason-answer",
+    "self-oss-instruct/direct-answer",
+    "slimorca/direct-answer",
+    "smoltalk/direct-answer",
+    "smoltalk/direct-answer+tools",
+    "smoltalk/direct-call",
+    "smoltalk2-mid/reason-answer",
+    "smoltalk2-sft/direct-answer",
+    "smoltalk2-sft/direct-answer+tools",
+    "smoltalk2-sft/direct-call",
+    "smoltalk2-sft/reason-answer",
+    "smoltalk2-sft/reason-call",
+    "stage3-v2/direct-answer",
+    "stage3-v2/direct-answer+tools",
+    "stage3-v2/direct-call",
+    "stage3-v2/reason-answer",
+    "stage3-v2/reason-answer+tools",
+    "stage3-v2/reason-call",
+    "systemchat-2/direct-answer",
+    "tulu3-if/direct-answer",
+    "ultrachat-200k/direct-answer",
+    "when2call/direct-answer",
+    "when2call/direct-answer+tools",
+    "wildchat/direct-answer",
+    "wildchat/reason-answer"
+  ],
+  "chat_template": {
+    "prompt": "BOS through the assistant header and its zero-loss <|think|> or <|no_think|> line, as SFT packed it",
+    "response": "reply, <|im_end|>, newline, EOS",
+    "truncated_rejected": "replies cut at their token budget, in a loop, or ending at EOS keep no turn end"
+  },
+  "composition": {
+    "train": {
+      "sources": {
+        "glaive-fc/direct-answer": 79,
+        "glaive-fc/direct-answer+tools": 306,
+        "glaive-fc/direct-call": 154,
+        "hermes-3/direct-answer": 358,
+        "hermes-tool-use/reason-answer+tools": 512,
+        "hermes-tool-use/reason-call": 264,
+        "magicoder/direct-answer": 168,
+        "no-robots/direct-answer": 70,
+        "oasst2/direct-answer": 147,
+        "opencode/reason-answer": 821,
+        "openmath/reason-answer": 1025,
+        "self-oss-instruct/direct-answer": 85,
+        "slimorca/direct-answer": 185,
+        "smoltalk/direct-answer": 303,
+        "smoltalk/direct-answer+tools": 680,
+        "smoltalk/direct-call": 447,
+        "smoltalk2-mid/reason-answer": 1207,
+        "smoltalk2-sft/direct-answer": 343,
+        "smoltalk2-sft/direct-answer+tools": 263,
+        "smoltalk2-sft/direct-call": 426,
+        "smoltalk2-sft/reason-answer": 340,
+        "smoltalk2-sft/reason-call": 403,
+        "stage3-v2/direct-answer": 711,
+        "stage3-v2/direct-answer+tools": 1899,
+        "stage3-v2/direct-call": 2226,
+        "stage3-v2/reason-answer": 777,
+        "stage3-v2/reason-answer+tools": 3098,
+        "stage3-v2/reason-call": 3579,
+        "systemchat-2/direct-answer": 127,
+        "tulu3-if/direct-answer": 116,
+        "ultrachat-200k/direct-answer": 165,
+        "when2call/direct-answer": 57,
+        "when2call/direct-answer+tools": 310,
+        "wildchat/direct-answer": 661,
+        "wildchat/reason-answer": 205
+      },
+      "basis": {
+        "reward": 4971,
+        "rules": 6282,
+        "reference": 7490,
+        "partial": 77,
+        "loop_stop": 3697
+      },
+      "chosen": {
+        "ok": 8403,
+        "reference:correct_call": 4145,
+        "loop_stop": 3697,
+        "correct_call": 2572,
+        "reference:ok": 2180,
+        "reference:correct_answer": 1165,
+        "correct_answer": 278,
+        "wrong_function": 48,
+        "missed_call": 15,
+        "wrong_arguments": 14
+      },
+      "rejected": {
+        "loop": 5930,
+        "ok": 4971,
+        "unneeded_call": 2492,
+        "wrong_arguments": 2099,
+        "wrong_function": 2079,
+        "unfinished": 1640,
+        "missed_call": 1209,
+        "unknown_function": 964,
+        "malformed_call": 504,
+        "missing_arguments": 255,
+        "wrong_answer": 154,
+        "repetitive": 92,
+        "no_final_answer": 48,
+        "unclosed_think": 44,
+        "call_without_tools": 19,
+        "think_in_direct": 10,
+        "malformed_think": 7
+      },
+      "balance": {
+        "available": {
+          "reason-call": 3715,
+          "reason-answer": 4522,
+          "direct-call": 3079,
+          "direct-answer": 4066
+        },
+        "cap": 3849
+      },
+      "loop_stops": 3697
+    },
+    "validation": {
+      "sources": {
+        "glaive-fc/direct-answer": 3,
+        "glaive-fc/direct-answer+tools": 8,
+        "glaive-fc/direct-call": 5,
+        "hermes-3/direct-answer": 11,
+        "hermes-tool-use/reason-answer+tools": 28,
+        "hermes-tool-use/reason-call": 10,
+        "magicoder/direct-answer": 2,
+        "no-robots/direct-answer": 3,
+        "oasst2/direct-answer": 1,
+        "opencode/reason-answer": 39,
+        "openmath/reason-answer": 59,
+        "self-oss-instruct/direct-answer": 7,
+        "slimorca/direct-answer": 6,
+        "smoltalk/direct-answer": 12,
+        "smoltalk/direct-answer+tools": 23,
+        "smoltalk/direct-call": 25,
+        "smoltalk2-mid/reason-answer": 30,
+        "smoltalk2-sft/direct-answer": 14,
+        "smoltalk2-sft/direct-answer+tools": 18,
+        "smoltalk2-sft/direct-call": 19,
+        "smoltalk2-sft/reason-answer": 19,
+        "smoltalk2-sft/reason-call": 26,
+        "stage3-v2/direct-answer": 32,
+        "stage3-v2/direct-answer+tools": 78,
+        "stage3-v2/direct-call": 79,
+        "stage3-v2/reason-answer": 31,
+        "stage3-v2/reason-answer+tools": 121,
+        "stage3-v2/reason-call": 141,
+        "systemchat-2/direct-answer": 5,
+        "tulu3-if/direct-answer": 1,
+        "ultrachat-200k/direct-answer": 8,
+        "when2call/direct-answer": 1,
+        "when2call/direct-answer+tools": 9,
+        "wildchat/direct-answer": 23,
+        "wildchat/reason-answer": 6
+      },
+      "basis": {
+        "rules": 281,
+        "reward": 205,
+        "reference": 278,
+        "partial": 3,
+        "loop_stop": 136
+      },
+      "chosen": {
+        "ok": 363,
+        "reference:correct_call": 163,
+        "loop_stop": 136,
+        "correct_call": 111,
+        "reference:ok": 69,
+        "reference:correct_answer": 46,
+        "correct_answer": 12,
+        "wrong_function": 2,
+        "missed_call": 1
+      },
+      "rejected": {
+        "loop": 219,
+        "ok": 205,
+        "wrong_arguments": 100,
+        "unneeded_call": 84,
+        "wrong_function": 65,
+        "unfinished": 63,
+        "missed_call": 52,
+        "malformed_call": 42,
+        "unknown_function": 38,
+        "missing_arguments": 15,
+        "wrong_answer": 9,
+        "unclosed_think": 5,
+        "no_final_answer": 4,
+        "call_without_tools": 1,
+        "malformed_think": 1
+      },
+      "balance": {
+        "available": {
+          "reason-call": 189,
+          "reason-answer": 203,
+          "direct-call": 123,
+          "direct-answer": 179
+        },
+        "cap": 154
+      },
+      "loop_stops": 136
+    }
+  },
+  "splits": {
+    "train": {
+      "pairs": 22517,
+      "tokens": 86450384,
+      "prompt_tokens": 29781563,
+      "chosen_tokens": 23379623,
+      "rejected_tokens": 33289198
+    },
+    "validation": {
+      "pairs": 903,
+      "tokens": 3556933,
+      "prompt_tokens": 1193370,
+      "chosen_tokens": 972868,
+      "rejected_tokens": 1390695
+    }
+  }
+}

requirements.txt

@@ -0,0 +1,10 @@
+torch>=2.4.0
+safetensors>=0.4.0
+tokenizers>=0.20.0
+huggingface_hub>=0.26.0
+transformers>=4.45.0
+einops>=0.8.0
+numpy>=1.26.0
+sympy>=1.13.0
+
+# flash-linear-attention is auto-installed by inference.py on first run.

sft_config.json

@@ -0,0 +1,51 @@
+{
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.95,
+  "adam_epsilon": 1e-08,
+  "cpu_prefetch_factor": 2,
+  "cpu_workers": 2,
+  "dataset_manifest": "/home/nvidia/llm-data/haiku/tokenized/sft-all-v3/manifest.json",
+  "ema_inv_gamma": 1.0,
+  "ema_max_decay": 0.9999,
+  "ema_power": 0.75,
+  "enable_gradient_checkpointing": true,
+  "exclude_components": [
+    "nemotron-cascade-2"
+  ],
+  "fp32_master_weights": true,
+  "full_shuffle": true,
+  "gradient_accumulation_steps": 16,
+  "gradient_checkpointing_layers": 22,
+  "gradient_clip_norm": 1.0,
+  "initial_run_dir": "/media/nvidia/4TB2/llm-runs/senryu/pretrain-v1",
+  "initial_weights": "ema",
+  "learning_rate": 0.0004,
+  "log_every_optimizer_steps": 1,
+  "logit_z_loss_coefficient": 0.0001,
+  "maximum_checkpoints": 5,
+  "maximum_optimizer_steps": null,
+  "micro_batch_size": 1,
+  "minimum_learning_rate_ratio": 0.1,
+  "muon_learning_rate": null,
+  "muon_lr_adjustment": "match_rms_adamw",
+  "muon_momentum": 0.95,
+  "optimizer": "hybrid_muon",
+  "output_dir": "/media/nvidia/4TB2/llm-runs/senryu/sft-v1",
+  "require_complete_pretraining": false,
+  "require_complete_sft_source": false,
+  "rewrite_empty_think": true,
+  "save_every_optimizer_steps": 100,
+  "seed": 20261008,
+  "sequence_length": 32768,
+  "shuffle_block_sequences": 2048,
+  "stream_download_threads": 8,
+  "stream_english_share": null,
+  "stream_open_languages": null,
+  "stream_pairs_per_row": 1,
+  "tokenizer": null,
+  "unclosed_think": "mask_end",
+  "validation_batches": 16,
+  "warmup_ratio": 0.0,
+  "warmup_steps": 200,
+  "weight_decay": 0.1
+}

sft_validation.json

@@ -0,0 +1,24 @@
+{
+  "checkpoint": "checkpoint-00012800",
+  "ema": {
+    "assistant_targets": 393188,
+    "batches": 16,
+    "elapsed_seconds": 6.345391824841499,
+    "loss": 0.8266990829440273,
+    "perplexity": 2.2857611656454533,
+    "targets_per_second": 61964.33740477821
+  },
+  "excluded_components": [
+    "nemotron-cascade-2"
+  ],
+  "normal": {
+    "assistant_targets": 393188,
+    "batches": 16,
+    "elapsed_seconds": 6.353875299915671,
+    "loss": 0.8744940579562701,
+    "perplexity": 2.3976619090753144,
+    "targets_per_second": 61881.60475941641
+  },
+  "optimizer_step": 12800,
+  "type": "sft_validation"
+}

special_token_ids.json

@@ -0,0 +1,24 @@
+{
+  "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,
+  "im_start": "<|im_start|>",
+  "im_start_id": 4,
+  "im_end": "<|im_end|>",
+  "im_end_id": 5,
+  "think_control": "<|think|>",
+  "think_control_id": 15,
+  "no_think_control": "<|no_think|>",
+  "no_think_control_id": 14,
+  "tool_response": "<|tool_response|>",
+  "tool_response_id": 7,
+  "think_open": "<think>",
+  "think_open_id": 8,
+  "think_close": "</think>",
+  "think_close_id": 9
+}

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,11 @@
+from tiny_gdn.config import TinyGDNConfig, haiku_config, haiku_kdn_config, senryu_config
+from tiny_gdn.model import TinyGDNForCausalLM, TinyGDNOutput
+
+__all__ = [
+    "TinyGDNConfig",
+    "TinyGDNForCausalLM",
+    "TinyGDNOutput",
+    "haiku_config",
+    "haiku_kdn_config",
+    "senryu_config",
+]

tiny_gdn/chatml.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+from typing import Any
+
+from tiny_gdn.smoltalk_chat import (
+    NO_THINK_PREFIX,
+    THINK_CONTROL_PREFIX,
+    ToolFormatChoice,
+    kwargs_from_inference_request,
+    materialize_smoltalk_messages,
+    merge_kwargs_with_system_message,
+    resolve_reasoning_mode,
+)
+from tiny_gdn.tools import prepare_inference_messages
+
+THINK_START = "<think>"
+THINK_END = "</think>"
+
+
+def decode_chat_completion(tokenizer: Any, token_ids: list[int]) -> str:
+    return tokenizer.decode(token_ids, skip_special_tokens=False)
+
+
+def append_assistant_generation_prompt(
+    tokenizer: Any,
+    token_ids: list[int],
+    *,
+    enable_thinking: bool | None = None,
+) -> None:
+    im_start = tokenizer.token_to_id("<|im_start|>")
+    if im_start is None:
+        raise RuntimeError("Tokenizer is missing <|im_start|>")
+    token_ids.append(im_start)
+    token_ids.extend(
+        tokenizer.encode("assistant\n", add_special_tokens=False).ids
+    )
+    if enable_thinking is True:
+        token_ids.extend(
+            tokenizer.encode(THINK_CONTROL_PREFIX, add_special_tokens=False).ids
+        )
+    elif enable_thinking is False:
+        token_ids.extend(
+            tokenizer.encode(NO_THINK_PREFIX, add_special_tokens=False).ids
+        )
+
+
+def encode_chatml_messages(
+    tokenizer: Any,
+    messages: list[dict[str, Any]],
+    *,
+    tools: Any | None = None,
+    enable_thinking: bool | None = None,
+    xml_tools: Any | None = None,
+    python_tools: Any | None = None,
+    tool_format: ToolFormatChoice = "json",
+) -> list[int]:
+    """``tool_format`` picks the SmolTalk tool layout; models without the thinking toggle only know XML."""
+
+    if enable_thinking is None:
+        prepared = prepare_inference_messages(messages, tools=tools)
+        assistant_thinking_control = None
+    else:
+        kwargs = kwargs_from_inference_request(
+            enable_thinking=enable_thinking,
+            xml_tools=xml_tools,
+            python_tools=python_tools,
+            tools=tools,
+            tool_format=tool_format,
+        )
+        prepared = materialize_smoltalk_messages(messages, kwargs)
+        _, resolved = merge_kwargs_with_system_message(messages, kwargs)
+        assistant_thinking_control = (
+            resolve_reasoning_mode(
+                resolved.enable_thinking,
+                resolved.custom_instructions,
+            )
+            == "/think"
+        )
+    required_tokens = {
+        token: tokenizer.token_to_id(token)
+        for token in (
+            "<|begin_of_text|>",
+            "<|im_start|>",
+            "<|im_end|>",
+        )
+    }
+    if any(token_id is None for token_id in required_tokens.values()):
+        raise RuntimeError("Tokenizer is missing ChatML special tokens")
+    bos_id = required_tokens["<|begin_of_text|>"]
+    im_start_id = required_tokens["<|im_start|>"]
+    im_end_id = required_tokens["<|im_end|>"]
+    if bos_id is None or im_start_id is None or im_end_id is None:
+        raise RuntimeError("Tokenizer ChatML IDs could not be resolved")
+
+    token_ids = [bos_id]
+    newline_ids = tokenizer.encode("\n", add_special_tokens=False).ids
+    allowed_roles = {"system", "user", "assistant"}
+    for message_index, message in enumerate(prepared):
+        role = message["role"]
+        content = message["content"]
+        masked_prefix = message.get("masked_prefix", "")
+        if role not in allowed_roles:
+            raise ValueError(
+                f"Unsupported chat role at index {message_index}: {role!r}"
+            )
+        if not isinstance(masked_prefix, str):
+            raise ValueError(
+                f"Masked prefix at index {message_index} must be a string"
+            )
+        if not content.strip():
+            raise ValueError(
+                f"Chat content at index {message_index} must be non-empty"
+            )
+        token_ids.append(im_start_id)
+        token_ids.extend(
+            tokenizer.encode(f"{role}\n", add_special_tokens=False).ids
+        )
+        if masked_prefix:
+            token_ids.extend(
+                tokenizer.encode(masked_prefix, add_special_tokens=False).ids
+            )
+        token_ids.extend(
+            tokenizer.encode(content, add_special_tokens=False).ids
+        )
+        token_ids.append(im_end_id)
+        token_ids.extend(newline_ids)
+    if prepared[-1]["role"] != "user":
+        raise ValueError("The final chat message must have role 'user'")
+    append_assistant_generation_prompt(
+        tokenizer,
+        token_ids,
+        enable_thinking=assistant_thinking_control,
+    )
+    return token_ids

tiny_gdn/cli_chat.py

@@ -0,0 +1,484 @@
+"""Terminal chat: live reasoning stream plus SmolTalk tool-call rounds.
+
+Tercet-R assistant turns may contain:
+
+- a zero-loss `<|think|>` / `<|no_think|>` control prefix
+- a `<think>…</think>` reasoning block
+- one or more SmolTalk JSON `<tool_call>` blocks (NVIDIA XML is also parsed)
+
+This module colours those regions as tokens arrive and, after a completed
+turn, collects tool observations (prefixed with `<|tool_response|>`) so the
+model can continue the same conversation.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+from typing import Any, Literal, TextIO
+
+from tiny_gdn.smoltalk_chat import (
+    TOOL_RESPONSE_TOKEN,
+    format_smoltalk_tool_call,
+    wrap_smoltalk_tool_result,
+)
+from tiny_gdn.code_exec import (
+    CALCULATOR_TOOL,
+    NEMOTRON_PYTHON_EXEC_TOOL,
+    execute_math_tool,
+    is_auto_math_tool_name,
+    reset_default_python_session,
+)
+from tiny_gdn.tools import ParsedToolCall, parse_tool_calls
+from tiny_gdn.web_search import (
+    NEMOTRON_WEB_SEARCH_TOOL,
+    is_web_search_tool_name,
+    query_from_arguments,
+    search_web,
+)
+
+
+SegmentKind = Literal["answer", "think", "tool_call"]
+MarkerKind = Literal[
+    "think_open",
+    "think_close",
+    "think_control",
+    "no_think",
+    "tool_open",
+    "tool_close",
+]
+
+THINK_OPEN = "<think>"
+THINK_CLOSE = "</think>"
+TOOL_CALL_OPEN = "<tool_call>"
+TOOL_CALL_CLOSE = "</tool_call>"
+THINK_CONTROL = "<|think|>"
+NO_THINK_CONTROL = "<|no_think|>"
+
+MARKERS: tuple[tuple[str, MarkerKind], ...] = (
+    (THINK_CLOSE, "think_close"),
+    (TOOL_CALL_CLOSE, "tool_close"),
+    (THINK_OPEN, "think_open"),
+    (TOOL_CALL_OPEN, "tool_open"),
+    (NO_THINK_CONTROL, "no_think"),
+    (THINK_CONTROL, "think_control"),
+)
+
+ANSI = {
+    "think": "\033[2;33m",
+    "tool_call": "\033[36m",
+    "answer": "\033[0m",
+    "reset": "\033[0m",
+}
+
+DEFAULT_MAX_TOOL_ROUNDS = 8
+
+
+@dataclass(frozen=True)
+class ContentSegment:
+    kind: SegmentKind
+    text: str
+    open: bool = False
+
+
+@dataclass(frozen=True)
+class GeneratedTurn:
+    text: str
+    token_count: int
+    stop_reason: str
+    tool_calls: tuple[ParsedToolCall, ...]
+
+
+def first_marker(text: str) -> tuple[int, str, MarkerKind] | None:
+    best: tuple[int, str, MarkerKind] | None = None
+    for marker, kind in MARKERS:
+        at = text.find(marker)
+        if at < 0:
+            continue
+        if (
+            best is None
+            or at < best[0]
+            or (at == best[0] and len(marker) > len(best[1]))
+        ):
+            best = (at, marker, kind)
+    return best
+
+
+def holdback_prefix_length(text: str) -> int:
+    if not text:
+        return 0
+    keep = 0
+    for marker, _kind in MARKERS:
+        limit = min(len(marker) - 1, len(text))
+        for size in range(1, limit + 1):
+            if marker.startswith(text[-size:]):
+                keep = max(keep, size)
+    return keep
+
+
+def mode_after_marker(kind: MarkerKind, current: SegmentKind) -> SegmentKind:
+    if kind in {"think_open", "think_control"}:
+        return "think"
+    if kind in {"think_close", "no_think", "tool_close"}:
+        return "answer"
+    if kind == "tool_open":
+        return "tool_call"
+    return current
+
+
+def _segment_text(segment: ContentSegment) -> str:
+    return segment.text.strip()
+
+
+def format_assistant_markdown(source: str) -> str:
+    """Render a streamed assistant turn as a single markdown block."""
+
+    think_parts: list[str] = []
+    answer_parts: list[str] = []
+    tool_parts: list[str] = []
+    for segment in split_assistant_segments(source):
+        text = _segment_text(segment)
+        if not text:
+            continue
+        if segment.kind == "think":
+            think_parts.append(text)
+        elif segment.kind == "tool_call":
+            tool_parts.append(text)
+        else:
+            answer_parts.append(text)
+    blocks: list[str] = []
+    if think_parts:
+        blocks.append("**Reasoning**\n\n" + "\n\n".join(think_parts))
+    if answer_parts:
+        blocks.append("\n\n".join(answer_parts))
+    for tool in tool_parts:
+        blocks.append(f"```tool_call\n{tool}\n```")
+    return "\n\n".join(blocks).strip()
+
+
+def format_assistant_chat_messages(source: str) -> list[dict[str, Any]]:
+    """Split a turn into Gradio thoughts plus a normal assistant chat message.
+
+    Messages with ``metadata.title`` render as collapsible thoughts. The reply
+    has no metadata so Gradio shows it as the chat bubble.
+    """
+
+    think_parts: list[str] = []
+    answer_parts: list[str] = []
+    tool_parts: list[str] = []
+    think_pending = False
+    for segment in split_assistant_segments(source):
+        text = _segment_text(segment)
+        if segment.kind == "think":
+            think_pending = segment.open
+            if text:
+                think_parts.append(text)
+            continue
+        think_pending = False
+        if not text:
+            continue
+        if segment.kind == "tool_call":
+            tool_parts.append(text)
+        else:
+            answer_parts.append(text)
+
+    messages: list[dict[str, Any]] = []
+    if think_parts:
+        messages.append(
+            {
+                "role": "assistant",
+                "content": "\n\n".join(think_parts),
+                "metadata": {
+                    "title": "Reasoning",
+                    "status": "pending" if think_pending else "done",
+                },
+            }
+        )
+    for tool in tool_parts:
+        messages.append(
+            {
+                "role": "assistant",
+                "content": f"```json\n{tool}\n```",
+                "metadata": {"title": "Tool call", "status": "done"},
+            }
+        )
+    answer = "\n\n".join(answer_parts)
+    if answer or not messages:
+        messages.append({"role": "assistant", "content": answer})
+    return messages
+
+
+def split_assistant_segments(source: str) -> list[ContentSegment]:
+    """Split a completed (or in-progress) assistant turn for tests / replay."""
+
+    segments: list[ContentSegment] = []
+    mode: SegmentKind = "answer"
+    cursor = 0
+    while cursor < len(source):
+        found = first_marker(source[cursor:])
+        if found is None:
+            tail = source[cursor:]
+            if tail:
+                segments.append(ContentSegment(kind=mode, text=tail, open=True))
+            break
+        at, marker, kind = found
+        at += cursor
+        if at > cursor:
+            segments.append(
+                ContentSegment(kind=mode, text=source[cursor:at], open=False)
+            )
+        mode = mode_after_marker(kind, mode)
+        cursor = at + len(marker)
+    return [segment for segment in segments if segment.text]
+
+
+class LiveReasoningStreamer:
+    """Colour reasoning and tool-call regions as decoded text grows."""
+
+    def __init__(
+        self,
+        writer: TextIO | None = None,
+        *,
+        color: bool | None = None,
+    ) -> None:
+        self.writer = writer if writer is not None else sys.stdout
+        if color is None:
+            color = bool(getattr(self.writer, "isatty", lambda: False)())
+        self.color = color
+        self._seen = ""
+        self._hold = ""
+        self._mode: SegmentKind = "answer"
+        self._style: SegmentKind | None = None
+        self._emitted_think_label = False
+        self._emitted_tool_label = False
+
+    def update(self, decoded: str) -> None:
+        if decoded.startswith(self._seen):
+            delta = decoded[len(self._seen) :]
+        else:
+            delta = decoded
+        self._seen = decoded
+        if delta:
+            self._consume(delta, final=False)
+
+    def finish(self) -> str:
+        if self._hold:
+            self._emit(self._hold)
+            self._hold = ""
+        self._set_style(None)
+        self.writer.write("\n")
+        self.writer.flush()
+        return self._seen
+
+    def _consume(self, delta: str, *, final: bool) -> None:
+        buffer = self._hold + delta
+        self._hold = ""
+        while buffer:
+            found = first_marker(buffer)
+            if found is None:
+                keep = 0 if final else holdback_prefix_length(buffer)
+                if keep:
+                    self._emit(buffer[:-keep])
+                    self._hold = buffer[-keep:]
+                else:
+                    self._emit(buffer)
+                return
+            at, marker, kind = found
+            if at:
+                self._emit(buffer[:at])
+            self._switch(kind)
+            buffer = buffer[at + len(marker) :]
+        if final:
+            return
+
+    def _switch(self, kind: MarkerKind) -> None:
+        nxt = mode_after_marker(kind, self._mode)
+        if nxt != self._mode and nxt == "answer":
+            self._emit_plain("\n")
+        self._mode = nxt
+        if nxt == "think" and not self._emitted_think_label:
+            self._emit_plain("\n")
+            self._set_style("think")
+            self._emit_plain("reasoning  ")
+            self._emitted_think_label = True
+        elif nxt == "tool_call" and not self._emitted_tool_label:
+            self._emit_plain("\n")
+            self._set_style("tool_call")
+            self._emit_plain("tool_call  ")
+            self._emitted_tool_label = True
+        elif nxt == "answer":
+            self._set_style("answer")
+
+    def _emit(self, text: str) -> None:
+        if not text:
+            return
+        self._set_style(self._mode)
+        self.writer.write(text)
+        self.writer.flush()
+
+    def _emit_plain(self, text: str) -> None:
+        if not text:
+            return
+        self._set_style(None)
+        self.writer.write(text)
+        self.writer.flush()
+
+    def _set_style(self, kind: SegmentKind | None) -> None:
+        if not self.color:
+            self._style = kind
+            return
+        if kind == self._style:
+            return
+        self.writer.write(ANSI["reset"])
+        if kind in {"think", "tool_call"}:
+            self.writer.write(ANSI[kind])
+        self._style = kind
+
+
+def prompt_tool_results(
+    calls: Sequence[ParsedToolCall],
+    *,
+    read_line: Callable[[str], str],
+    writer: TextIO | None = None,
+    execute_web_search: bool = True,
+) -> list[str]:
+    out = writer if writer is not None else sys.stdout
+    results: list[str] = []
+    auto_search = sum(
+        1
+        for call in calls
+        if execute_web_search and is_web_search_tool_name(call.name)
+    )
+    auto_math = sum(1 for call in calls if is_auto_math_tool_name(call.name))
+    if auto_search:
+        out.write(
+            f"\n{auto_search} web-search call(s) will run automatically "
+            f"(Tavily JSON, sent as {TOOL_RESPONSE_TOKEN}).\n"
+        )
+        out.flush()
+    if auto_math:
+        out.write(
+            f"\n{auto_math} python/calculator call(s) will run automatically "
+            f"(sent as {TOOL_RESPONSE_TOKEN}).\n"
+        )
+        out.flush()
+    manual = len(calls) - auto_search - auto_math
+    if manual:
+        out.write(
+            f"\n{manual} tool call(s). Paste each observation; "
+            f"it is sent as {TOOL_RESPONSE_TOKEN}.\n"
+        )
+        out.flush()
+    for index, call in enumerate(calls, start=1):
+        out.write(
+            f"\n[{index}/{len(calls)}] "
+            f"{format_smoltalk_tool_call(call.name, call.arguments)}\n"
+        )
+        out.flush()
+        if execute_web_search and is_web_search_tool_name(call.name):
+            query = query_from_arguments(call.arguments)
+            out.write(f"searching {query!r}…\n")
+            out.flush()
+            results.append(search_web(query))
+            continue
+        if is_auto_math_tool_name(call.name):
+            out.write("running python…\n")
+            out.flush()
+            try:
+                results.append(execute_math_tool(call.name, call.arguments))
+            except Exception as error:
+                results.append(f"Error: {error}")
+            continue
+        results.append(read_line(f"result[{call.name}]> "))
+    return results
+
+
+def append_tool_round(
+    messages: list[dict[str, Any]],
+    assistant_text: str,
+    raw_results: Sequence[str],
+) -> None:
+    messages.append({"role": "assistant", "content": assistant_text})
+    for raw in raw_results:
+        wrapped = wrap_smoltalk_tool_result(raw)
+        if not wrapped:
+            raise ValueError("Tool result cannot be empty")
+        messages.append({"role": "tool", "content": raw})
+
+
+def resolve_cli_tools(spec: str | None, tools_json: str | None) -> list[dict[str, Any]] | None:
+    tools: list[dict[str, Any]] = []
+    if spec:
+        for name in spec.split(","):
+            key = name.strip().lower()
+            if not key or key in {"none", "off"}:
+                continue
+            if is_web_search_tool_name(key):
+                tools.append(NEMOTRON_WEB_SEARCH_TOOL)
+                continue
+            if key in {"python", "python-exec", "code-interpreter"}:
+                tools.append(NEMOTRON_PYTHON_EXEC_TOOL)
+                continue
+            if key in {"calculator", "calc"}:
+                tools.append(CALCULATOR_TOOL)
+                continue
+            raise ValueError(
+                f"Unknown built-in tool {name!r}. "
+                "Use web-search, python, calculator, or --tools-json."
+            )
+    if tools_json:
+        payload = json.loads(tools_json)
+        if isinstance(payload, dict):
+            tools.append(payload)
+        elif isinstance(payload, list):
+            tools.extend(payload)
+        else:
+            raise ValueError("tools JSON must be an object or array")
+    return tools or None
+
+
+@dataclass
+class ChatLoopState:
+    messages: list[dict[str, Any]] = field(default_factory=list)
+    system: str = ""
+    enable_thinking: bool = True
+    tools: list[dict[str, Any]] | None = None
+
+    def reset(self) -> None:
+        self.messages = []
+        reset_default_python_session()
+        if self.system.strip():
+            self.messages.append({"role": "system", "content": self.system.strip()})
+
+    def add_user(self, text: str) -> None:
+        self.messages.append({"role": "user", "content": text})
+
+
+def apply_slash_command(state: ChatLoopState, text: str) -> str | None:
+    """Return a status string if `text` is a slash command, else None."""
+
+    command = text.strip()
+    lowered = command.lower()
+    if lowered in {"/exit", "/quit"}:
+        return "exit"
+    if lowered == "/reset":
+        state.reset()
+        return "history cleared"
+    if lowered == "/think":
+        if state.messages:
+            return "thinking can only be changed on a fresh conversation (/reset first)"
+        state.enable_thinking = True
+        return "thinking on — next assistant turn is prefixed with <|think|>"
+    if lowered in {"/no_think", "/nothink"}:
+        if state.messages:
+            return "thinking can only be changed on a fresh conversation (/reset first)"
+        state.enable_thinking = False
+        return "thinking off — next assistant turn is prefixed with <|no_think|>"
+    if lowered.startswith("/system"):
+        rest = command[len("/system") :].strip()
+        state.system = rest
+        state.reset()
+        return "system prompt updated" if rest else "system prompt cleared"
+    return None

tiny_gdn/code_exec.py

@@ -0,0 +1,498 @@
+"""Restricted stateful Python for Nemotron Cascade-2 math tool calls.
+
+Tercet-R saw `stateful_python_code_exec` on `math/math_tool.jsonl` during
+the Cascade-2 SFT stage. Observations are the Jupyter-style stdout / last
+expression value, not a JSON envelope.
+
+The worker process only allows a math-oriented import whitelist and rejects
+dunder access, so a tool call cannot read the filesystem or start a shell.
+"""
+
+from __future__ import annotations
+
+import ast
+import io
+import json
+import math
+import os
+import subprocess
+import sys
+import threading
+from collections.abc import Mapping
+from pathlib import Path
+from queue import Empty, Queue
+from typing import Any
+
+
+NEMOTRON_PYTHON_EXEC_TOOL: dict[str, Any] = {
+    "type": "function",
+    "function": {
+        "name": "stateful_python_code_exec",
+        "description": (
+            "Call this function to execute Python code in a stateful Jupyter "
+            "notebook environment. Python will respond with the output of the "
+            "execution or time out after 120.0 seconds."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "code": {
+                    "type": "string",
+                    "description": "Code to execute",
+                }
+            },
+            "required": ["code"],
+        },
+    },
+}
+
+CALCULATOR_TOOL: dict[str, Any] = {
+    "type": "function",
+    "function": {
+        "name": "calculator",
+        "description": "Evaluate a math expression and return the numeric result.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "expression": {
+                    "type": "string",
+                    "description": "Math expression, for example 12.5 * (3 + 4)",
+                }
+            },
+            "required": ["expression"],
+        },
+    },
+}
+
+PYTHON_EXEC_TOOL_NAMES = frozenset(
+    {
+        "stateful_python_code_exec",
+        "python_code_exec",
+        "code_interpreter",
+        "python_exec",
+        "python",
+    }
+)
+CALCULATOR_TOOL_NAMES = frozenset({"calculator", "calc"})
+DEFAULT_EXEC_TIMEOUT_SECONDS = 30.0
+_SRC_ROOT = Path(__file__).resolve().parents[1]
+
+_session_lock = threading.Lock()
+_default_session: StatefulPythonSession | None = None
+
+
+def _normalize_tool_name(name: str) -> str:
+    return name.strip().lower().replace("-", "_")
+
+
+def is_python_exec_tool_name(name: str) -> bool:
+    return _normalize_tool_name(name) in PYTHON_EXEC_TOOL_NAMES
+
+
+def is_calculator_tool_name(name: str) -> bool:
+    return _normalize_tool_name(name) in CALCULATOR_TOOL_NAMES
+
+
+def is_auto_math_tool_name(name: str) -> bool:
+    return is_python_exec_tool_name(name) or is_calculator_tool_name(name)
+
+
+def code_from_arguments(arguments: Mapping[str, Any] | None) -> str:
+    if not arguments:
+        return ""
+    for key in ("code", "expression", "expr", "source"):
+        raw = arguments.get(key)
+        if isinstance(raw, str) and raw.strip():
+            return raw
+        if raw is not None and key != "code":
+            return str(raw)
+    return ""
+
+
+class CodeExecError(ValueError):
+    """Rejected or failed tool code."""
+
+
+class StatefulPythonSession:
+    """One long-lived restricted interpreter, matching the training tool."""
+
+    def __init__(self, *, timeout_seconds: float = DEFAULT_EXEC_TIMEOUT_SECONDS) -> None:
+        self.timeout_seconds = timeout_seconds
+        self._lock = threading.Lock()
+        self._process: subprocess.Popen[str] | None = None
+
+    def close(self) -> None:
+        with self._lock:
+            self._kill_locked()
+
+    def reset(self) -> None:
+        with self._lock:
+            self._kill_locked()
+
+    def run(self, code: str) -> str:
+        text = code.strip()
+        if not text:
+            raise CodeExecError("code is empty")
+        with self._lock:
+            return self._run_locked(text)
+
+    def _run_locked(self, code: str) -> str:
+        process = self._ensure_process_locked()
+        try:
+            process.stdin.write(json.dumps({"code": code}, ensure_ascii=False) + "\n")
+            process.stdin.flush()
+            line = self._readline_locked(process, self.timeout_seconds)
+        except (BrokenPipeError, OSError) as error:
+            self._kill_locked()
+            raise CodeExecError(f"python worker died: {error}") from error
+        if not line:
+            stderr = ""
+            if process.stderr is not None:
+                try:
+                    stderr = process.stderr.read()
+                except OSError:
+                    stderr = ""
+            self._kill_locked()
+            detail = stderr.strip() or "python worker closed stdout"
+            raise CodeExecError(detail)
+        try:
+            payload = json.loads(line)
+        except json.JSONDecodeError as error:
+            self._kill_locked()
+            raise CodeExecError(f"python worker returned invalid JSON: {line!r}") from error
+        if not isinstance(payload, dict):
+            raise CodeExecError("python worker returned a non-object")
+        output = str(payload.get("output") or "")
+        if payload.get("ok") is True:
+            return output if output else "None"
+        raise CodeExecError(output or "execution failed")
+
+    def _readline_locked(self, process: subprocess.Popen[str], timeout: float) -> str:
+        if process.stdout is None:
+            raise CodeExecError("python worker has no stdout")
+        lines: Queue[str] = Queue()
+
+        def _read() -> None:
+            lines.put(process.stdout.readline() if process.stdout is not None else "")
+
+        reader = threading.Thread(target=_read, daemon=True)
+        reader.start()
+        reader.join(timeout)
+        if reader.is_alive():
+            self._kill_locked()
+            raise CodeExecError(f"timed out after {timeout:g}s")
+        try:
+            return lines.get_nowait()
+        except Empty:
+            return ""
+
+    def _ensure_process_locked(self) -> subprocess.Popen[str]:
+        process = self._process
+        if process is not None and process.poll() is None:
+            return process
+        env = os.environ.copy()
+        pythonpath = env.get("PYTHONPATH", "")
+        env["PYTHONPATH"] = (
+            str(_SRC_ROOT) if not pythonpath else f"{_SRC_ROOT}{os.pathsep}{pythonpath}"
+        )
+        self._process = subprocess.Popen(
+            [sys.executable, "-m", "tiny_gdn.code_exec"],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            bufsize=1,
+            env=env,
+        )
+        if self._process.stdin is None or self._process.stdout is None:
+            self._kill_locked()
+            raise CodeExecError("failed to start python worker")
+        return self._process
+
+    def _kill_locked(self) -> None:
+        process = self._process
+        self._process = None
+        if process is None:
+            return
+        try:
+            process.kill()
+            process.wait(timeout=2)
+        except (OSError, subprocess.TimeoutExpired):
+            pass
+
+
+def default_python_session() -> StatefulPythonSession:
+    global _default_session
+    with _session_lock:
+        if _default_session is None:
+            _default_session = StatefulPythonSession()
+        return _default_session
+
+
+def reset_default_python_session() -> None:
+    global _default_session
+    with _session_lock:
+        session = _default_session
+        _default_session = None
+    if session is not None:
+        session.close()
+
+
+def execute_calculator(expression: str) -> str:
+    tree = ast.parse(expression, mode="eval")
+    _reject_unsafe_ast(tree)
+    value = eval(  # noqa: S307 — AST already rejected unsafe nodes
+        compile(tree, "<calculator>", "eval"),
+        {"__builtins__": {}},
+        _calculator_namespace(),
+    )
+    return _format_result(value)
+
+
+def execute_math_tool(name: str, arguments: Mapping[str, Any] | None) -> str:
+    source = code_from_arguments(arguments)
+    if is_calculator_tool_name(name) and "\n" not in source and ";" not in source:
+        try:
+            return execute_calculator(source)
+        except (SyntaxError, CodeExecError, TypeError, ValueError, ZeroDivisionError) as error:
+            raise CodeExecError(str(error)) from error
+    return default_python_session().run(source)
+
+
+def _calculator_namespace() -> dict[str, Any]:
+    names = {
+        key: getattr(math, key)
+        for key in dir(math)
+        if not key.startswith("_")
+    }
+    names.update(
+        {
+            "abs": abs,
+            "min": min,
+            "max": max,
+            "round": round,
+            "pow": pow,
+            "pi": math.pi,
+            "e": math.e,
+            "tau": math.tau,
+            "inf": math.inf,
+        }
+    )
+    return names
+
+
+def _reject_unsafe_ast(tree: ast.AST) -> None:
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.ClassDef, ast.AsyncFunctionDef)):
+            raise CodeExecError("class definitions are blocked")
+        if isinstance(node, ast.Attribute) and node.attr.startswith("_"):
+            raise CodeExecError("dunder attribute access is blocked")
+        if isinstance(node, ast.Name) and node.id.startswith("_"):
+            raise CodeExecError("dunder names are blocked")
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
+            if node.func.id in _BANNED_CALLS:
+                raise CodeExecError(f"{node.func.id}() is blocked")
+
+
+_BANNED_CALLS = frozenset(
+    {
+        "eval",
+        "exec",
+        "compile",
+        "open",
+        "input",
+        "breakpoint",
+        "getattr",
+        "setattr",
+        "delattr",
+        "globals",
+        "locals",
+        "vars",
+        "dir",
+        "help",
+        "__import__",
+        "memoryview",
+        "exit",
+        "quit",
+    }
+)
+
+
+def _format_result(value: Any) -> str:
+    if value is None:
+        return "None"
+    if isinstance(value, bool):
+        return str(value)
+    if isinstance(value, float):
+        if value.is_integer() and abs(value) < 1e15:
+            return str(int(value))
+        return format(value, ".12g")
+    return str(value)
+
+
+_ALLOWED_IMPORT_ROOTS = frozenset(
+    {
+        "cmath",
+        "collections",
+        "copy",
+        "decimal",
+        "fractions",
+        "functools",
+        "itertools",
+        "json",
+        "math",
+        "mpmath",
+        "numbers",
+        "numpy",
+        "operator",
+        "re",
+        "statistics",
+        "string",
+        "sympy",
+        "textwrap",
+        "unicodedata",
+    }
+)
+
+_WORKER_NAMESPACE: dict[str, Any] | None = None
+
+
+def _allowed_import(
+    name: str,
+    globals: dict[str, Any] | None = None,
+    locals: dict[str, Any] | None = None,
+    fromlist: tuple[str, ...] = (),
+    level: int = 0,
+) -> Any:
+    root = name.split(".")[0]
+    if root not in _ALLOWED_IMPORT_ROOTS:
+        raise ImportError(f"import of {name!r} is blocked")
+    return __import__(name, globals, locals, fromlist, level)
+
+
+def _worker_namespace() -> dict[str, Any]:
+    global _WORKER_NAMESPACE
+    if _WORKER_NAMESPACE is None:
+        builtins = {
+            "abs": abs,
+            "all": all,
+            "any": any,
+            "bin": bin,
+            "bool": bool,
+            "bytes": bytes,
+            "chr": chr,
+            "complex": complex,
+            "dict": dict,
+            "divmod": divmod,
+            "enumerate": enumerate,
+            "filter": filter,
+            "float": float,
+            "format": format,
+            "frozenset": frozenset,
+            "hex": hex,
+            "int": int,
+            "isinstance": isinstance,
+            "issubclass": issubclass,
+            "iter": iter,
+            "len": len,
+            "list": list,
+            "map": map,
+            "max": max,
+            "min": min,
+            "next": next,
+            "oct": oct,
+            "ord": ord,
+            "pow": pow,
+            "print": print,
+            "range": range,
+            "repr": repr,
+            "reversed": reversed,
+            "round": round,
+            "set": set,
+            "slice": slice,
+            "sorted": sorted,
+            "str": str,
+            "sum": sum,
+            "tuple": tuple,
+            "zip": zip,
+            "True": True,
+            "False": False,
+            "None": None,
+            "__import__": _allowed_import,
+        }
+        _WORKER_NAMESPACE = {
+            "__builtins__": builtins,
+            "__name__": "__tool__",
+            "math": math,
+        }
+    return _WORKER_NAMESPACE
+
+
+def run_cell(code: str) -> str:
+    tree = ast.parse(code)
+    _reject_unsafe_import_roots(tree)
+    _reject_unsafe_ast(tree)
+    namespace = _worker_namespace()
+    buffer = io.StringIO()
+    previous = sys.stdout
+    sys.stdout = buffer
+    try:
+        if tree.body and isinstance(tree.body[-1], ast.Expr):
+            body = tree.body[:-1]
+            last = tree.body[-1]
+            if body:
+                exec(  # noqa: S102
+                    compile(ast.Module(body, type_ignores=[]), "<tool>", "exec"),
+                    namespace,
+                    namespace,
+                )
+            value = eval(  # noqa: S307
+                compile(ast.Expression(last.value), "<tool>", "eval"),
+                namespace,
+                namespace,
+            )
+            printed = buffer.getvalue()
+            if value is None:
+                return printed if printed else "None"
+            rendered = _format_result(value)
+            return f"{printed}{rendered}" if printed else rendered
+        exec(compile(tree, "<tool>", "exec"), namespace, namespace)  # noqa: S102
+        printed = buffer.getvalue()
+        return printed if printed else "None"
+    finally:
+        sys.stdout = previous
+
+
+def _reject_unsafe_import_roots(tree: ast.AST) -> None:
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            for alias in node.names:
+                root = alias.name.split(".")[0]
+                if root not in _ALLOWED_IMPORT_ROOTS:
+                    raise CodeExecError(f"import of {alias.name!r} is blocked")
+        elif isinstance(node, ast.ImportFrom):
+            root = (node.module or "").split(".")[0]
+            if root not in _ALLOWED_IMPORT_ROOTS:
+                raise CodeExecError(f"import of {node.module!r} is blocked")
+
+
+def _worker_loop() -> None:
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            request = json.loads(line)
+            code = str(request.get("code") or "")
+            output = run_cell(code)
+            sys.stdout.write(json.dumps({"ok": True, "output": output}, ensure_ascii=False) + "\n")
+        except Exception as error:
+            detail = str(error) or type(error).__name__
+            sys.stdout.write(
+                json.dumps({"ok": False, "output": detail}, ensure_ascii=False) + "\n"
+            )
+        sys.stdout.flush()
+
+
+if __name__ == "__main__":
+    _worker_loop()

tiny_gdn/config.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, fields
+from pathlib import Path
+from typing import Any
+
+LINEAR_LAYER_TYPES = frozenset({"gdn2", "kda", "diag_kdn"})
+DIAG_KDN_HEAD_DIMS = (64, 128)
+
+
+@dataclass(frozen=True)
+class TinyGDNConfig:
+    architecture: str = "TinyGDNForCausalLM"
+    model_type: str = "tiny_gdn"
+
+    vocab_size: int = 49_152
+    # Deep-thin sizing is deliberate: controlled sub-billion studies find
+    # depth materially more valuable than width around the 125M-150M scale.
+    hidden_size: int = 512
+    intermediate_size: int = 1_472
+    num_hidden_layers: int = 32
+
+    num_attention_heads: int = 4
+    num_key_value_heads: int = 1
+    attention_head_dim: int = 128
+    full_attention_interval: int = 4
+    attention_dropout: float = 0.0
+    partial_rotary_factor: float = 0.5
+    rope_theta: float = 1_000_000.0
+
+    linear_num_heads: int = 4
+    linear_num_value_heads: int = 4
+    linear_head_dim: int = 128
+    linear_expand_v: float = 1.0
+    linear_conv_kernel_dim: int = 4
+    allow_negative_eigenvalues: bool = False
+    linear_layer_type: str = "gdn2"
+    attention_layer_type: str = "full_attention"
+    mlp_activation: str = "swiglu"
+    use_block_attn_res: bool = False
+    kda_safe_gate: bool = True
+    kda_lower_bound: float = -5.0
+    mla_q_lora_rank: int | None = 512
+    mla_kv_lora_rank: int = 256
+    mla_qk_nope_head_dim: int = 128
+    mla_v_head_dim: int = 128
+    situ_gate_cap: float = 4.0
+    situ_up_cap: float = 25.0
+
+    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 LINEAR_LAYER_TYPES:
+            raise ValueError(
+                f"linear_layer_type must be one of {sorted(LINEAR_LAYER_TYPES)}"
+            )
+        if self.linear_layer_type == "diag_kdn":
+            if self.linear_head_dim not in DIAG_KDN_HEAD_DIMS:
+                raise ValueError(
+                    f"diag_kdn needs linear_head_dim in {DIAG_KDN_HEAD_DIMS}"
+                )
+            if self.linear_num_value_heads != self.linear_num_heads:
+                raise ValueError("diag_kdn needs linear_num_value_heads == linear_num_heads")
+            if self.linear_expand_v != 1.0:
+                raise ValueError("diag_kdn needs linear_expand_v == 1.0")
+            if self.allow_negative_eigenvalues:
+                raise ValueError("diag_kdn derives its gain from the covariance; negative eigenvalues are unsupported")
+        if self.attention_layer_type not in {"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)
+
+
+def haiku_kdn_config(**overrides: Any) -> TinyGDNConfig:
+    """Haiku with Diagonal KDN in place of KDA in the 27 linear layers.
+
+    Opt-in: at equal tokens it trailed KDA in Haiku pilots and trains ~23%
+    slower at full depth (docs/architecture.md, "Haiku v2").
+    """
+
+    return haiku_config(**{"linear_layer_type": "diag_kdn", **overrides})
+
+
+def senryu_config(**overrides: Any) -> TinyGDNConfig:
+    """800M Senryu: 3×Gated DeltaNet-2 + 1×Gated-MLA over 44 layers, AttnRes, SiTU-GLU.
+
+    Haiku's width, heads, tokenizer, and residual design with GDN-2's
+    channel-wise erase and write gates in the 33 linear layers. A rank-512
+    query LoRA costs as many parameters as a full 1024×1024 query, so queries
+    are full rank, and the KV latent doubles to 512. The feed-forward narrows
+    to 3392 to pay for the depth. Next-token only: MTP has no pilot win at
+    this scale (docs/architecture.md, "Senryu").
+    """
+
+    return haiku_config(
+        **{
+            "architecture": "SenryuForCausalLM",
+            "model_type": "senryu",
+            "intermediate_size": 3392,
+            "num_hidden_layers": 44,
+            "linear_layer_type": "gdn2",
+            "mla_q_lora_rank": None,
+            "mla_kv_lora_rank": 512,
+            "mtp_num_heads": 0,
+            "mtp_loss_weight": 0.0,
+            **overrides,
+        }
+    )

tiny_gdn/detokenize.py

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

tiny_gdn/fused_ops.py

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

tiny_gdn/haiku_layers.py

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

tiny_gdn/model.py

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

tiny_gdn/nn_common.py

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

tiny_gdn/smoltalk_chat.py

@@ -0,0 +1,676 @@
+"""SmolTalk2 / SmolLM3 chat-template contract for Tercet-R.
+
+The SFT subset stores the toggle and tools in `chat_template_kwargs`:
+
+- `enable_thinking`: True → assistant turns receive a zero-loss
+  `<|think|>\\n` control prefix; assistant targets may emit `<think>` naturally
+- `enable_thinking`: False → assistant turns receive a zero-loss
+  `<|no_think|>\\n` control prefix
+- `custom_instructions`: the system prompt, used verbatim. Only the SmolLM3
+  template directives `/think`, `/no_think`, and `/system_override` are
+  removed because they are control markers, not prompt text
+- `xml_tools` / `python_tools`: appended to the system prompt as a Hermes
+  `<tools>` block only when tools are present and the prompt does not already
+  carry one
+- `tool_format` (inference only): "auto" and "json" render Hermes JSON
+  `<tool_call>` blocks, the layout the served models were trained on. "xml"
+  is Cascade's OpenHands layout and is used only when a request asks for it
+
+There is no default identity, metadata header, or wrapper: a conversation
+without a system prompt and without tools has no system turn at all.
+
+Tool results (`tool` / `function` roles, or text still wrapped in Nemotron
+`<tool_response>` XML) are rendered as a user turn where every result is
+prefixed with the `<|tool_response|>` special token.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass, replace
+from typing import Any, Literal
+
+from tiny_gdn.tools import (
+    json_ready,
+    TOOL_RESPONSE_CLOSE,
+    TOOL_RESPONSE_OPEN,
+    TOOLS_CLOSE,
+    TOOLS_OPEN,
+    assistant_message_content,
+    coerce_tools,
+    format_tool_calls_xml,
+    loads_json,
+    message_text,
+    parse_arguments_payload,
+    render_available_tools,
+    strip_leaked_chatml,
+)
+
+
+ReasoningMode = Literal["/think", "/no_think"]
+ToolFormat = Literal["json", "xml"]
+ToolFormatChoice = Literal["auto", "json", "xml"]
+
+NO_THINK_TOKEN = "<|no_think|>"
+NO_THINK_PREFIX = f"{NO_THINK_TOKEN}\n"
+THINK_CONTROL_TOKEN = "<|think|>"
+THINK_CONTROL_PREFIX = f"{THINK_CONTROL_TOKEN}\n"
+TOOL_RESPONSE_TOKEN = "<|tool_response|>"
+TOOL_RESPONSE_PREFIX = f"{TOOL_RESPONSE_TOKEN}\n"
+THINK_OPEN = "<think>"
+THINK_CLOSE = "</think>"
+TOOL_CALL_OPEN = "<tool_call>"
+TOOL_CALL_CLOSE = "</tool_call>"
+SYSTEM_OVERRIDE_MARK = "/system_override"
+THINK_MARK = "/think"
+NO_THINK_MARK = "/no_think"
+
+# SmolTalk2 SFT used this markdown web_search schema. Tercet-R stage-3
+# packs Nemotron Agentic-v2 search instead (`web-search` + Tavily JSON);
+# see tiny_gdn.web_search.NEMOTRON_WEB_SEARCH_TOOL.
+SMOLTALK_WEB_SEARCH_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "web_search",
+        "description": (
+            "Performs a web search for a query and returns a string of the "
+            "top search results formatted as markdown with titles, links, "
+            "and descriptions."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "The search query to perform.",
+                }
+            },
+            "required": ["query"],
+        },
+    },
+}
+XML_TOOLS_PREAMBLE = (
+    "You may call one or more functions to assist with the user query.\n"
+    "You are provided with function signatures within <tools></tools> XML "
+    "tags:\n<tools>\n"
+)
+XML_TOOLS_EPILOGUE = (
+    "</tools>\n\n"
+    "For 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>"
+)
+PYTHON_TOOLS_PREAMBLE = (
+    "When you send a message containing Python code between '<code>' and "
+    "'</code>' tags, it will be executed in a stateful Jupyter notebook "
+    "environment, and you will then be given the output to continued "
+    "reasoning in an agentic loop.\n\n"
+    "You can use the following tools in your python code like regular "
+    "functions:\n<tools>\n"
+)
+PYTHON_TOOLS_EPILOGUE = (
+    "</tools>\n\n"
+    "The state persists between code executions: so variables that you "
+    "define in one step are still available thereafter."
+)
+_EMPTY_THINK_BLOCK = re.compile(
+    re.escape(THINK_OPEN) + r"\s*" + re.escape(THINK_CLOSE)
+)
+
+
+@dataclass(frozen=True)
+class ChatTemplateKwargs:
+    custom_instructions: str = ""
+    enable_thinking: bool = True
+    xml_tools: tuple[str, ...] = ()
+    python_tools: tuple[str, ...] = ()
+    tool_format: ToolFormat = "json"
+
+
+def coerce_tool_strings(raw_tools: Any) -> tuple[str, ...]:
+    if raw_tools is None:
+        return ()
+    if isinstance(raw_tools, str):
+        text = raw_tools.strip()
+        return (text,) if text else ()
+    if isinstance(raw_tools, dict):
+        raw_tools = [raw_tools]
+    if not isinstance(raw_tools, list):
+        return ()
+    rendered: list[str] = []
+    for item in raw_tools:
+        if isinstance(item, str):
+            text = item.strip()
+            if text:
+                rendered.append(text)
+            continue
+        if isinstance(item, dict):
+            rendered.append(json.dumps(item, ensure_ascii=False))
+    return tuple(rendered)
+
+
+def parse_chat_template_kwargs(raw: Any) -> ChatTemplateKwargs:
+    if raw is None:
+        return ChatTemplateKwargs()
+    if isinstance(raw, str):
+        text = raw.strip()
+        if not text:
+            return ChatTemplateKwargs()
+        raw = json.loads(text)
+    if not isinstance(raw, dict):
+        return ChatTemplateKwargs()
+    enable = raw.get("enable_thinking")
+    if enable is None:
+        enable_thinking = True
+    elif isinstance(enable, bool):
+        enable_thinking = enable
+    else:
+        enable_thinking = str(enable).strip().lower() in {"1", "true", "yes", "on"}
+    custom = raw.get("custom_instructions")
+    return ChatTemplateKwargs(
+        custom_instructions=custom if isinstance(custom, str) else "",
+        enable_thinking=enable_thinking,
+        xml_tools=coerce_tool_strings(raw.get("xml_tools")),
+        python_tools=coerce_tool_strings(raw.get("python_tools")),
+    )
+
+
+def resolve_reasoning_mode(
+    enable_thinking: bool,
+    custom_instructions: str,
+) -> ReasoningMode:
+    if NO_THINK_MARK in custom_instructions:
+        return "/no_think"
+    if THINK_MARK in custom_instructions:
+        return "/think"
+    return "/think" if enable_thinking else "/no_think"
+
+
+def strip_template_markers(custom_instructions: str) -> str:
+    return (
+        custom_instructions.replace(NO_THINK_MARK, "")
+        .replace(THINK_MARK, "")
+        .replace(SYSTEM_OVERRIDE_MARK, "")
+        .strip()
+    )
+
+
+def strip_tools_wrapper(item: str) -> str:
+    """Drop a `<tools>` / `</tools>` wrapper some sources bake into each entry."""
+
+    text = item.strip()
+    if text.startswith(TOOLS_OPEN):
+        text = text[len(TOOLS_OPEN) :]
+    if text.endswith(TOOLS_CLOSE):
+        text = text[: -len(TOOLS_CLOSE)]
+    return text.strip()
+
+
+def render_xml_tools_block(xml_tools: tuple[str, ...]) -> str:
+    entries = [strip_tools_wrapper(item) for item in xml_tools]
+    entries = [item for item in entries if item]
+    if not entries:
+        return ""
+    body = "".join(f"{item}\n" for item in entries)
+    return f"{XML_TOOLS_PREAMBLE}{body}{XML_TOOLS_EPILOGUE}"
+
+
+def render_cascade_tools_block(xml_tools: tuple[str, ...]) -> str:
+    """The same tools in Nemotron Cascade-2's layout: XML schemas, `<function=…>` calls."""
+
+    entries = [strip_tools_wrapper(item) for item in xml_tools]
+    return render_available_tools([loads_json(item) for item in entries if item])
+
+
+def render_python_tools_block(python_tools: tuple[str, ...]) -> str:
+    entries = [strip_tools_wrapper(item) for item in python_tools]
+    entries = [item for item in entries if item]
+    if not entries:
+        return ""
+    body = "".join(f"{item}\n" for item in entries)
+    return f"{PYTHON_TOOLS_PREAMBLE}{body}{PYTHON_TOOLS_EPILOGUE}"
+
+
+def unwrap_tool_result(content: str) -> str:
+    """Return the bare tool output, without ChatML leaks or `<tool_response>` XML."""
+
+    text = strip_leaked_chatml(content).strip()
+    if text.startswith(TOOL_RESPONSE_TOKEN):
+        text = text[len(TOOL_RESPONSE_TOKEN) :].strip()
+    if TOOL_RESPONSE_OPEN in text:
+        chunks: list[str] = []
+        cursor = 0
+        while True:
+            start = text.find(TOOL_RESPONSE_OPEN, cursor)
+            if start < 0:
+                tail = text[cursor:].strip()
+                if tail:
+                    chunks.append(tail)
+                break
+            head = text[cursor:start].strip()
+            if head:
+                chunks.append(head)
+            start += len(TOOL_RESPONSE_OPEN)
+            end = text.find(TOOL_RESPONSE_CLOSE, start)
+            if end < 0:
+                inner = text[start:].strip()
+                cursor = len(text)
+            else:
+                inner = text[start:end].strip()
+                cursor = end + len(TOOL_RESPONSE_CLOSE)
+            if inner:
+                chunks.append(inner)
+            if cursor >= len(text):
+                break
+        text = "\n".join(chunks)
+    return text
+
+
+def wrap_smoltalk_tool_result(content: str) -> str:
+    text = unwrap_tool_result(content)
+    if not text:
+        return ""
+    return f"{TOOL_RESPONSE_PREFIX}{text}"
+
+
+def strip_empty_think_blocks(content: str) -> str:
+    """Drop `<think>` spans whose body is empty or whitespace-only.
+
+    Nemotron and similar traces sometimes emit `<think></think>` or
+    `<think>\\n</think>` with no reasoning. Those are not think turns.
+    """
+
+    stripped = _EMPTY_THINK_BLOCK.sub("", content)
+    if stripped == content:
+        return content
+    return stripped.lstrip("\r\n")
+
+
+def rewrite_tool_result_content(content: str) -> str:
+    stripped = content.lstrip()
+    if stripped.startswith(TOOL_RESPONSE_TOKEN) and TOOL_RESPONSE_OPEN not in stripped:
+        return content
+    text = strip_leaked_chatml(content).strip()
+    if TOOL_RESPONSE_OPEN in text:
+        chunks: list[str] = []
+        cursor = 0
+        while True:
+            start = text.find(TOOL_RESPONSE_OPEN, cursor)
+            if start < 0:
+                break
+            start += len(TOOL_RESPONSE_OPEN)
+            end = text.find(TOOL_RESPONSE_CLOSE, start)
+            inner = (text[start:end] if end >= 0 else text[start:]).strip()
+            if inner:
+                chunks.append(wrap_smoltalk_tool_result(inner))
+            if end < 0:
+                break
+            cursor = end + len(TOOL_RESPONSE_CLOSE)
+        if chunks:
+            return "\n".join(chunks)
+    return wrap_smoltalk_tool_result(text)
+
+
+def apply_reasoning_and_tool_controls(messages: list[dict[str, str]]) -> list[dict[str, str]]:
+    """Attach Haiku/Tercet-R control tokens to any ChatML conversation.
+
+    Assistant turns without a `masked_prefix` get `<|think|>` when the
+    content already contains a `<think>` block, otherwise `<|no_think|>`.
+    Empty or whitespace-only `<think>` spans are stripped and treated as
+    no-think, including when a think prefix was already attached.
+    `tool` roles and Nemotron `<tool_response>` user turns become user
+    turns prefixed with the `<|tool_response|>` special.
+    Existing prefixes and already-rewritten tool results are left alone
+    unless an empty think block was removed and no real think remains.
+    """
+
+    prepared: list[dict[str, str]] = []
+    pending: list[str] = []
+
+    def flush_tools() -> None:
+        if pending:
+            prepared.append({"role": "user", "content": "\n".join(pending)})
+            pending.clear()
+
+    for message in messages:
+        role = message["role"]
+        content = message["content"]
+        if role == "tool":
+            rewritten = rewrite_tool_result_content(content)
+            if rewritten:
+                pending.append(rewritten)
+            continue
+        if role == "user" and (
+            content.lstrip().startswith(TOOL_RESPONSE_TOKEN)
+            or TOOL_RESPONSE_OPEN in content
+        ):
+            rewritten = rewrite_tool_result_content(content)
+            if rewritten:
+                pending.append(rewritten)
+            continue
+        flush_tools()
+        if role == "assistant":
+            rendered = dict(message)
+            original = content
+            content = strip_empty_think_blocks(content)
+            rendered["content"] = content
+            has_think = THINK_OPEN in content
+            existing_prefix = rendered.get("masked_prefix")
+            if not existing_prefix:
+                rendered["masked_prefix"] = (
+                    THINK_CONTROL_PREFIX if has_think else NO_THINK_PREFIX
+                )
+            elif content != original and not has_think:
+                rendered["masked_prefix"] = NO_THINK_PREFIX
+            prepared.append(rendered)
+            continue
+        prepared.append(dict(message))
+    flush_tools()
+    return prepared
+
+
+def openai_tools_as_xml_strings(raw_tools: Any) -> tuple[str, ...]:
+    if raw_tools is None:
+        return ()
+    rendered: list[str] = []
+    for tool in coerce_tools(raw_tools):
+        rendered.append(
+            json.dumps(
+                json_ready({"type": "function", "function": tool}),
+                ensure_ascii=False,
+                default=str,
+            )
+        )
+    return tuple(rendered)
+
+
+def format_smoltalk_tool_call(name: str, arguments: dict[str, Any]) -> str:
+    payload = {"name": name, "arguments": json_ready(arguments)}
+    return (
+        f"{TOOL_CALL_OPEN}\n"
+        f"{json.dumps(payload, ensure_ascii=False, default=str)}\n"
+        f"{TOOL_CALL_CLOSE}"
+    )
+
+
+def format_smoltalk_tool_calls(raw_tool_calls: Any) -> str:
+    if not raw_tool_calls:
+        return ""
+    if not isinstance(raw_tool_calls, list):
+        raise ValueError("tool_calls must be a list")
+    chunks: list[str] = []
+    for raw_call in raw_tool_calls:
+        if not isinstance(raw_call, dict):
+            raise ValueError("Each tool_call must be an object")
+        payload = (
+            raw_call.get("function")
+            if isinstance(raw_call.get("function"), dict)
+            else raw_call
+        )
+        if not isinstance(payload, dict):
+            raise ValueError("tool_call is missing a function object")
+        name = payload.get("name")
+        if not isinstance(name, str) or not name.strip():
+            raise ValueError("tool_call is missing a function name")
+        arguments = parse_arguments_payload(payload.get("arguments"))
+        chunks.append(format_smoltalk_tool_call(name.strip(), arguments))
+    return "\n".join(chunks)
+
+
+def assistant_smoltalk_content(message: dict[str, Any], *, tool_format: ToolFormat = "json") -> str:
+    content = assistant_message_content(
+        {**message, "tool_calls": None}
+    )
+    if TOOL_CALL_OPEN in content:
+        return content
+    if tool_format == "xml":
+        tool_xml = format_tool_calls_xml(message.get("tool_calls")).rstrip("\n")
+    else:
+        tool_xml = format_smoltalk_tool_calls(message.get("tool_calls"))
+    if not tool_xml:
+        return content
+    if content.strip():
+        return f"{content.rstrip()}\n{tool_xml}"
+    return tool_xml
+
+
+def render_system_prompt(kwargs: ChatTemplateKwargs) -> str:
+    """Render the system turn verbatim; empty string means no system turn."""
+
+    custom = strip_leaked_chatml(strip_template_markers(kwargs.custom_instructions))
+    tool_blocks: list[str] = []
+    if TOOLS_OPEN not in custom:
+        if kwargs.tool_format == "xml":
+            xml_block = render_cascade_tools_block(kwargs.xml_tools)
+        else:
+            xml_block = render_xml_tools_block(kwargs.xml_tools)
+        if xml_block:
+            tool_blocks.append(xml_block)
+        python_block = render_python_tools_block(kwargs.python_tools)
+        if python_block:
+            tool_blocks.append(python_block)
+    sections = [custom] if custom else []
+    sections.extend(tool_blocks)
+    return "\n\n".join(sections).strip()
+
+
+def merge_kwargs_with_system_message(
+    messages: list[dict[str, Any]],
+    kwargs: ChatTemplateKwargs,
+) -> tuple[list[dict[str, Any]], ChatTemplateKwargs]:
+    if not messages:
+        return messages, kwargs
+    first = messages[0]
+    if first.get("role") != "system":
+        return messages, kwargs
+    system_text = message_text(first.get("content"))
+    if system_text.strip():
+        merged = (
+            system_text
+            if not kwargs.custom_instructions.strip()
+            else f"{kwargs.custom_instructions.rstrip()}\n\n{system_text.strip()}"
+        )
+        kwargs = replace(kwargs, custom_instructions=merged)
+    return messages[1:], kwargs
+
+
+def kwargs_from_inference_request(
+    *,
+    enable_thinking: bool,
+    custom_instructions: str = "",
+    xml_tools: Any = None,
+    python_tools: Any = None,
+    tools: Any = None,
+    tool_format: ToolFormatChoice = "json",
+) -> ChatTemplateKwargs:
+    """``tool_format="auto"`` is the Hermes JSON layout every served model was trained on.
+
+    The Cascade XML layout is only used when a request asks for ``xml`` explicitly.
+    """
+
+    xml = coerce_tool_strings(xml_tools)
+    if not xml:
+        xml = openai_tools_as_xml_strings(tools)
+    if tool_format == "auto":
+        tool_format = "json"
+    return ChatTemplateKwargs(
+        custom_instructions=custom_instructions,
+        enable_thinking=enable_thinking,
+        xml_tools=xml,
+        python_tools=coerce_tool_strings(python_tools),
+        tool_format=tool_format,
+    )
+
+
+def _normalize_role(raw_role: Any, index: int) -> str:
+    if not isinstance(raw_role, str):
+        raise ValueError(f"Unsupported chat role at index {index}: {raw_role!r}")
+    role = raw_role.strip().lower()
+    if role == "human":
+        return "user"
+    if role == "gpt":
+        return "assistant"
+    if role == "function":
+        return "tool"
+    if role not in {"system", "user", "assistant", "tool"}:
+        raise ValueError(f"Unsupported chat role at index {index}: {raw_role!r}")
+    return role
+
+
+def materialize_smoltalk_messages(
+    raw_messages: list[dict[str, Any]],
+    kwargs: ChatTemplateKwargs,
+) -> list[dict[str, str]]:
+    if not raw_messages:
+        raise ValueError("Chat history cannot be empty")
+    remaining, resolved = merge_kwargs_with_system_message(raw_messages, kwargs)
+    for index, raw_message in enumerate(remaining):
+        if not isinstance(raw_message, dict):
+            raise ValueError(f"Unsupported chat message at index {index}")
+        if _normalize_role(raw_message.get("role"), index) != "system":
+            continue
+        extra = strip_leaked_chatml(message_text(raw_message.get("content")))
+        if extra:
+            resolved = replace(
+                resolved,
+                custom_instructions=(
+                    f"{resolved.custom_instructions.rstrip()}\n\n{extra}".strip()
+                ),
+            )
+    mode = resolve_reasoning_mode(
+        resolved.enable_thinking,
+        resolved.custom_instructions,
+    )
+    masked_prefix = THINK_CONTROL_PREFIX if mode == "/think" else NO_THINK_PREFIX
+    prepared: list[dict[str, str]] = []
+    system_prompt = render_system_prompt(resolved)
+    if system_prompt:
+        prepared.append({"role": "system", "content": system_prompt})
+    pending_tool_results: list[str] = []
+    for index, raw_message in enumerate(remaining):
+        role = _normalize_role(raw_message.get("role"), index)
+        if role == "system":
+            continue
+        if role == "assistant":
+            content = assistant_smoltalk_content(raw_message, tool_format=resolved.tool_format)
+        else:
+            content = message_text(raw_message.get("content"))
+        if role == "tool":
+            content = wrap_smoltalk_tool_result(content)
+            if not content:
+                raise ValueError(f"Chat content at index {index} must be non-empty")
+            pending_tool_results.append(content)
+            continue
+        if pending_tool_results:
+            prepared.append(
+                {"role": "user", "content": "\n".join(pending_tool_results)}
+            )
+            pending_tool_results = []
+        if not content.strip():
+            raise ValueError(f"Chat content at index {index} must be non-empty")
+        prepared_message = {"role": role, "content": content}
+        if role == "assistant":
+            prepared_message["masked_prefix"] = masked_prefix
+        prepared.append(prepared_message)
+    if pending_tool_results:
+        prepared.append({"role": "user", "content": "\n".join(pending_tool_results)})
+    if not any(message["role"] != "system" for message in prepared):
+        raise ValueError("Chat history cannot be empty")
+    return prepared
+
+
+def format_think_block(reasoning: str, content: str) -> str:
+    """SmolTalk2 layout: `<think>\\n…\\n</think>\\n\\nanswer`."""
+
+    body = reasoning.strip()
+    if content:
+        return f"{THINK_OPEN}\n{body}\n{THINK_CLOSE}\n\n{content}"
+    return f"{THINK_OPEN}\n{body}\n{THINK_CLOSE}\n\n"
+
+
+def render_assistant_turn(
+    content: str,
+    *,
+    reasoning: str | None,
+    tool_calls: list[dict[str, Any]] | None,
+) -> dict[str, str]:
+    """Render one assistant turn with a per-turn thinking control prefix.
+
+    A turn with reasoning gets `<|think|>` and a `<think>` block; a turn
+    without reasoning gets `<|no_think|>`. Tool calls become SmolTalk JSON
+    `<tool_call>` blocks unless the content already inlines them.
+    """
+
+    text = strip_empty_think_blocks(content.strip())
+    if tool_calls and TOOL_CALL_OPEN not in text:
+        calls = "\n".join(
+            format_smoltalk_tool_call(call["name"], call["arguments"]) for call in tool_calls
+        )
+        text = f"{text}\n{calls}" if text else calls
+    if reasoning and reasoning.strip():
+        return {
+            "role": "assistant",
+            "content": format_think_block(reasoning, text),
+            "masked_prefix": THINK_CONTROL_PREFIX,
+        }
+    return {"role": "assistant", "content": text, "masked_prefix": NO_THINK_PREFIX}
+
+
+def materialize_canonical_conversation(
+    *,
+    system: str | None,
+    tools: Any,
+    messages: list[dict[str, Any]],
+) -> list[dict[str, str]]:
+    """Render a stage-3 canonical conversation into ChatML-ready turns.
+
+    `messages` entries are `{"role", "content", "reasoning"?, "tool_calls"?}`
+    with roles user / assistant / tool. The system prompt is used verbatim;
+    tools are appended only when the prompt lacks a `<tools>` block.
+    """
+
+    kwargs = ChatTemplateKwargs(
+        custom_instructions=system or "",
+        xml_tools=openai_tools_as_xml_strings(tools) if tools else (),
+    )
+    prepared: list[dict[str, str]] = []
+    system_prompt = render_system_prompt(kwargs)
+    if system_prompt:
+        prepared.append({"role": "system", "content": system_prompt})
+    pending_tool_results: list[str] = []
+    for index, message in enumerate(messages):
+        role = message["role"]
+        if role == "tool":
+            wrapped = wrap_smoltalk_tool_result(message_text(message.get("content")))
+            if not wrapped:
+                raise ValueError(f"Tool result at index {index} is empty")
+            pending_tool_results.append(wrapped)
+            continue
+        if pending_tool_results:
+            prepared.append({"role": "user", "content": "\n".join(pending_tool_results)})
+            pending_tool_results = []
+        if role == "user":
+            content = message_text(message.get("content"))
+            if not content.strip():
+                raise ValueError(f"User turn at index {index} is empty")
+            prepared.append({"role": "user", "content": content})
+            continue
+        if role == "assistant":
+            rendered = render_assistant_turn(
+                message_text(message.get("content")),
+                reasoning=message.get("reasoning"),
+                tool_calls=message.get("tool_calls"),
+            )
+            if not rendered["content"].strip():
+                raise ValueError(f"Assistant turn at index {index} is empty")
+            prepared.append(rendered)
+            continue
+        raise ValueError(f"Unsupported role at index {index}: {role!r}")
+    if pending_tool_results:
+        raise ValueError("Conversation ends with an unanswered tool result")
+    if not prepared or prepared[-1]["role"] != "assistant":
+        raise ValueError("Conversation must end with an assistant turn")
+    return prepared

tiny_gdn/tools.py

@@ -0,0 +1,624 @@
+"""Nemotron Cascade-2 tool protocol (XML in ChatML content).
+
+Tercet-R is trained on `nvidia/Nemotron-Cascade-2-SFT-Data`, which inlines
+available tools, calls, and results as `<tools>` / `<tool_call>` /
+`<tool_response>` in message text. That is not OpenAI `tool_calls` JSON and
+not this tokenizer's unused `<|tool_call|>` / `<|tool_response|>` specials.
+
+This module matches NVIDIA's Cascade-2 chat template:
+https://huggingface.co/nvidia/Nemotron-Cascade-2-30B-A3B/blob/main/chat_template.jinja
+"""
+
+from __future__ import annotations
+
+import ast
+import json
+import re
+import sys
+from dataclasses import dataclass
+from typing import Any
+
+
+TOOL_CALL_OPEN = "<tool_call>"
+TOOL_CALL_CLOSE = "</tool_call>"
+TOOL_RESPONSE_OPEN = "<tool_response>"
+TOOL_RESPONSE_CLOSE = "</tool_response>"
+TOOLS_OPEN = "<tools>"
+TOOLS_CLOSE = "</tools>"
+THINK_OPEN = "<think>"
+THINK_CLOSE = "</think>"
+
+CHATML_START_RE = re.compile(r"^<\|im_start\|>[A-Za-z]+\n")
+CHATML_END_RE = re.compile(r"\n?<\|im_end\|>\s*$")
+TOOL_CALL_BLOCK_RE = re.compile(
+    rf"{re.escape(TOOL_CALL_OPEN)}(.*?){re.escape(TOOL_CALL_CLOSE)}",
+    re.DOTALL,
+)
+FUNCTION_BLOCK_RE = re.compile(
+    r"<function=([^>\s]+)>(.*?)</function>",
+    re.DOTALL,
+)
+PARAMETER_BLOCK_RE = re.compile(
+    r"<parameter=([^>\s]+)>\n?(.*?)\n?</parameter>",
+    re.DOTALL,
+)
+FUNCTION_NAME_RE = re.compile(r"<function=([^>\s]+)")
+
+INFERENCE_ROLE_MAP = {
+    "system": "system",
+    "user": "user",
+    "human": "user",
+    "assistant": "assistant",
+    "gpt": "assistant",
+    "tool": "tool",
+    "function": "tool",
+}
+
+TOOLS_PREAMBLE = "# Tools\n\nYou have access to the following functions:\n\n"
+TOOL_CALL_INSTRUCTIONS = (
+    "\n\nIf you choose to call a function ONLY reply in the following format "
+    "with NO suffix:\n\n"
+    "<tool_call>\n"
+    "<function=example_function_name>\n"
+    "<parameter=example_parameter_1>\n"
+    "value_1\n"
+    "</parameter>\n"
+    "<parameter=example_parameter_2>\n"
+    "This is the value for the second parameter\n"
+    "that can span\n"
+    "multiple lines\n"
+    "</parameter>\n"
+    "</function>\n"
+    "</tool_call>\n\n"
+    "<IMPORTANT>\n"
+    "Reminder:\n"
+    "- Function calls MUST follow the specified format: an inner "
+    "<function=...></function> block must be nested within "
+    "<tool_call></tool_call> XML tags\n"
+    "- Required parameters MUST be specified\n"
+    "- You may provide optional reasoning for your function call in natural "
+    "language BEFORE the function call, but NOT after\n"
+    "- If there is no function call available, answer the question like "
+    "normal with your current knowledge and do not tell the user about "
+    "function calls\n"
+    "</IMPORTANT>"
+)
+
+
+@dataclass(frozen=True)
+class ParsedToolCall:
+    name: str
+    arguments: dict[str, Any]
+    raw: str
+
+
+def message_text(raw_content: Any) -> str:
+    if raw_content is None:
+        return ""
+    if isinstance(raw_content, str):
+        return raw_content
+    if isinstance(raw_content, list):
+        parts: list[str] = []
+        for item in raw_content:
+            if isinstance(item, str):
+                parts.append(item)
+                continue
+            if not isinstance(item, dict):
+                continue
+            part_type = item.get("type")
+            if part_type in {None, "text", "input_text", "output_text"}:
+                text = item.get("text")
+                if isinstance(text, str):
+                    parts.append(text)
+        return "".join(parts)
+    return str(raw_content)
+
+
+def strip_leaked_chatml(content: str) -> str:
+    text = content.strip()
+    while True:
+        match = CHATML_START_RE.match(text)
+        if match is None:
+            break
+        text = text[match.end() :]
+    text = CHATML_END_RE.sub("", text)
+    return text.strip()
+
+
+def has_tools_block(content: str) -> bool:
+    return TOOLS_OPEN in content
+
+
+def wrap_tool_response(content: str) -> str:
+    text = strip_leaked_chatml(content)
+    if TOOL_RESPONSE_OPEN in text:
+        return text
+    return f"{TOOL_RESPONSE_OPEN}\n\n\n{text}\n{TOOL_RESPONSE_CLOSE}"
+
+
+def _xml_value(value: Any) -> str:
+    if isinstance(value, dict) or (
+        isinstance(value, (list, tuple)) and not isinstance(value, (str, bytes))
+    ):
+        return json.dumps(value, ensure_ascii=False)
+    if value is True or value is False or value is None:
+        return str(value)
+    return str(value)
+
+
+def _render_extra_keys(payload: dict[str, Any], handled: set[str]) -> str:
+    chunks: list[str] = []
+    for key, value in payload.items():
+        if key in handled:
+            continue
+        chunks.append(f"\n<{key}>{_xml_value(value)}</{key}>")
+    return "".join(chunks)
+
+
+def _unwrap_tool(raw_tool: Any) -> dict[str, Any]:
+    if not isinstance(raw_tool, dict):
+        raise ValueError("Each tool must be an object")
+    if isinstance(raw_tool.get("function"), dict):
+        tool = dict(raw_tool["function"])
+    else:
+        tool = dict(raw_tool)
+    name = tool.get("name")
+    if not isinstance(name, str) or not name.strip():
+        raise ValueError("Tool is missing a function name")
+    tool["name"] = name.strip()
+    return tool
+
+
+def coerce_tools(raw_tools: Any) -> list[dict[str, Any]]:
+    if raw_tools is None:
+        return []
+    if isinstance(raw_tools, str):
+        text = raw_tools.strip()
+        if not text:
+            return []
+        raw_tools = loads_json(text)
+    if isinstance(raw_tools, dict):
+        raw_tools = [raw_tools]
+    if not isinstance(raw_tools, list):
+        raise ValueError("tools must be a list of function specs")
+    return [_unwrap_tool(item) for item in raw_tools]
+
+
+def render_function_schema(tool: dict[str, Any]) -> str:
+    chunks = [f"\n<function>\n<name>{tool['name']}</name>"]
+    description = tool.get("description")
+    if isinstance(description, str) and description.strip():
+        chunks.append(f"\n<description>{description.strip()}</description>")
+    chunks.append("\n<parameters>")
+    parameters = tool.get("parameters")
+    properties: dict[str, Any] = {}
+    if isinstance(parameters, dict):
+        raw_properties = parameters.get("properties")
+        if isinstance(raw_properties, dict):
+            properties = raw_properties
+        for param_name, raw_fields in properties.items():
+            fields = raw_fields if isinstance(raw_fields, dict) else {}
+            chunks.append("\n<parameter>")
+            chunks.append(f"\n<name>{param_name}</name>")
+            if "type" in fields:
+                chunks.append(f"\n<type>{_xml_value(fields['type'])}</type>")
+            if isinstance(fields.get("description"), str) and fields["description"].strip():
+                chunks.append(
+                    f"\n<description>{fields['description'].strip()}</description>"
+                )
+            if "enum" in fields:
+                chunks.append(f"\n<enum>{_xml_value(fields['enum'])}</enum>")
+            chunks.append(
+                _render_extra_keys(
+                    fields,
+                    {"name", "type", "description", "enum"},
+                )
+            )
+            chunks.append("\n</parameter>")
+        chunks.append(
+            _render_extra_keys(parameters, {"type", "properties", "required"})
+        )
+        if "required" in parameters:
+            chunks.append(f"\n<required>{_xml_value(parameters['required'])}</required>")
+    chunks.append("\n</parameters>")
+    chunks.append(
+        _render_extra_keys(
+            tool,
+            {"type", "name", "description", "parameters"},
+        )
+    )
+    chunks.append("\n</function>")
+    return "".join(chunks)
+
+
+def render_available_tools(raw_tools: Any) -> str:
+    tools = coerce_tools(raw_tools)
+    if not tools:
+        return ""
+    body = "".join(render_function_schema(tool) for tool in tools)
+    return (
+        f"{TOOLS_PREAMBLE}{TOOLS_OPEN}{body}\n{TOOLS_CLOSE}"
+        f"{TOOL_CALL_INSTRUCTIONS}"
+    )
+
+
+# Cascade's OpenHands conversations used XML schemas and `<function=…>` calls.
+# Those conversations are not in the served models' training data; `auto` is JSON.
+XML_LAYOUT_TOOL_NAMES = frozenset({"execute_bash", "str_replace_editor"})
+
+
+def prefers_xml_tool_calls(raw_tools: Any) -> bool:
+    """Whether a tool list holds names Cascade rendered only in the XML layout.
+
+    Served models were not trained on that data, so inference ``auto`` does not
+    use this. It remains for callers that still want to detect the old tool set.
+    """
+
+    try:
+        tools = coerce_tools(raw_tools)
+    except ValueError:
+        return False
+    return any(tool["name"].lower().replace("-", "_") in XML_LAYOUT_TOOL_NAMES for tool in tools)
+
+
+def inject_available_tools(system_content: str, raw_tools: Any) -> str:
+    block = render_available_tools(raw_tools)
+    if not block:
+        return system_content
+    if has_tools_block(system_content):
+        return system_content
+    if not system_content.strip():
+        return block
+    return f"{system_content.rstrip()}\n\n{block}"
+
+
+def parse_json_int(text: str) -> int | str:
+    """Keep oversized JSON integers as strings.
+
+    Python 3.12 refuses to convert integers longer than
+    ``sys.get_int_max_str_digits()`` (default 4300). Tool-call payloads
+    sometimes embed hashes or blobs as bare JSON numbers; those must not
+    abort packing.
+    """
+
+    digits = text.lstrip("+-")
+    limit = sys.get_int_max_str_digits()
+    if limit and len(digits) > limit:
+        return text
+    return int(text)
+
+
+def loads_json(text: str) -> Any:
+    return json.loads(text, parse_int=parse_json_int)
+
+
+# Models sometimes close a JSON string with `".}` instead of `"}`.
+_TRAILING_STRING_PERIOD_RE = re.compile(r'"\s*\.(?=\s*[}\],])')
+_TRAILING_COMMA_RE = re.compile(r",\s*(?=[}\]])")
+_JSON_NAME_RE = re.compile(r'"name"\s*:\s*"([^"]+)"')
+_JSON_QUERY_RE = re.compile(r'"query"\s*:\s*"([^"]*)"')
+
+
+def repair_jsonish(text: str) -> str:
+    repaired = _TRAILING_STRING_PERIOD_RE.sub('"', text.strip())
+    return _TRAILING_COMMA_RE.sub("", repaired)
+
+
+def loads_jsonish(text: str) -> Any:
+    """JSON with common model slips repaired; Python literals (single quotes, True) also parse."""
+
+    stripped = text.strip()
+    if not stripped:
+        raise json.JSONDecodeError("Empty JSON", text, 0)
+    last_error: Exception | None = None
+    for candidate in dict.fromkeys((stripped, repair_jsonish(stripped))):
+        try:
+            return loads_json(candidate)
+        except (json.JSONDecodeError, ValueError, RecursionError) as error:
+            last_error = error
+    try:
+        return json_ready(ast.literal_eval(stripped))
+    except (ValueError, TypeError, SyntaxError, MemoryError, RecursionError):
+        pass
+    try:
+        return loads_json(stripped.replace("'", '"'))
+    except (json.JSONDecodeError, ValueError, RecursionError):
+        if last_error is not None:
+            raise last_error from None
+        raise
+
+
+def extract_json_tool_fields(text: str) -> tuple[str, dict[str, Any]] | None:
+    name_match = _JSON_NAME_RE.search(text)
+    if name_match is None:
+        return None
+    arguments: dict[str, Any] = {}
+    query_match = _JSON_QUERY_RE.search(text)
+    if query_match is not None:
+        arguments["query"] = query_match.group(1)
+    return name_match.group(1).strip(), arguments
+
+
+def _json_ready_int(value: int) -> int | str:
+    try:
+        json.dumps(value)
+    except ValueError:
+        previous = sys.get_int_max_str_digits()
+        sys.set_int_max_str_digits(0)
+        try:
+            return str(value)
+        finally:
+            sys.set_int_max_str_digits(previous)
+    return value
+
+
+def json_ready(value: Any) -> Any:
+    """Coerce Python literals into JSON-serialisable values.
+
+    Some SFT sources store tool arguments as Python literals. `ast.literal_eval`
+    turns `{1, 2}` into a `set` and `...` into `Ellipsis`, which later
+    `json.dumps` calls reject. Oversized ints are stored as decimal strings.
+    """
+
+    if value is None or isinstance(value, (bool, float, str)):
+        return value
+    if isinstance(value, int):
+        return _json_ready_int(value)
+    if value is Ellipsis:
+        return None
+    if isinstance(value, dict):
+        return {str(key): json_ready(item) for key, item in value.items()}
+    if isinstance(value, (set, frozenset)):
+        items = [json_ready(item) for item in value]
+        try:
+            return sorted(
+                items,
+                key=lambda item: json.dumps(item, sort_keys=True, default=str),
+            )
+        except TypeError:
+            return items
+    if isinstance(value, tuple):
+        return [json_ready(item) for item in value]
+    if isinstance(value, list):
+        return [json_ready(item) for item in value]
+    if isinstance(value, bytes):
+        try:
+            return value.decode("utf-8")
+        except UnicodeDecodeError:
+            return list(value)
+    try:
+        json.dumps(value)
+        return value
+    except (TypeError, ValueError):
+        return str(value)
+
+
+def parse_argument_value(raw: str) -> Any:
+    text = raw.strip()
+    if not text:
+        return ""
+    try:
+        return loads_json(text)
+    except (json.JSONDecodeError, ValueError, RecursionError):
+        return text
+
+
+def parse_arguments_payload(raw: Any) -> dict[str, Any]:
+    if raw is None:
+        return {}
+    ready = json_ready(raw)
+    if isinstance(ready, dict):
+        return ready
+    if isinstance(ready, str):
+        text = ready.strip()
+        if not text:
+            return {}
+        try:
+            loaded = loads_json(text)
+        except (json.JSONDecodeError, ValueError, RecursionError):
+            return {"value": ready}
+        loaded = json_ready(loaded)
+        if isinstance(loaded, dict):
+            return loaded
+        return {"value": loaded}
+    return {"value": ready}
+
+
+def format_tool_call_xml(name: str, arguments: dict[str, Any]) -> str:
+    chunks = [f"{TOOL_CALL_OPEN}\n<function={name}>\n"]
+    for key, value in arguments.items():
+        chunks.append(f"<parameter={key}>\n{_xml_value(value)}\n</parameter>\n")
+    chunks.append(f"</function>\n{TOOL_CALL_CLOSE}\n")
+    return "".join(chunks)
+
+
+def format_tool_calls_xml(raw_tool_calls: Any) -> str:
+    if not raw_tool_calls:
+        return ""
+    if not isinstance(raw_tool_calls, list):
+        raise ValueError("tool_calls must be a list")
+    chunks: list[str] = []
+    for raw_call in raw_tool_calls:
+        if not isinstance(raw_call, dict):
+            raise ValueError("Each tool_call must be an object")
+        payload = raw_call.get("function") if isinstance(raw_call.get("function"), dict) else raw_call
+        if not isinstance(payload, dict):
+            raise ValueError("tool_call is missing a function object")
+        name = payload.get("name")
+        if not isinstance(name, str) or not name.strip():
+            raise ValueError("tool_call is missing a function name")
+        arguments = parse_arguments_payload(payload.get("arguments"))
+        chunks.append(format_tool_call_xml(name.strip(), arguments))
+    return "".join(chunks)
+
+
+def _json_tool_call(payload: Any, raw: str) -> ParsedToolCall | None:
+    if not isinstance(payload, dict):
+        return None
+    nested = payload.get("function")
+    source = nested if isinstance(nested, dict) else payload
+    name = source.get("name")
+    if not isinstance(name, str) or not name.strip():
+        name = payload.get("name")
+    if not isinstance(name, str) or not name.strip():
+        return None
+    arguments = parse_arguments_payload(
+        source.get("arguments", source.get("parameters", payload.get("arguments")))
+    )
+    return ParsedToolCall(name=name.strip(), arguments=arguments, raw=raw)
+
+
+def parse_json_tool_calls(
+    inner: str, raw: str = "", *, strict: bool = False
+) -> list[ParsedToolCall] | None:
+    """Calls in the JSON body of one `<tool_call>` block, or None when it holds none.
+
+    The body is one call object or a list of them (``[]`` is an explicit "no call").
+    Unless ``strict``, a body that does not parse still yields the call name and
+    query it mentions.
+    """
+
+    text = inner.strip()
+    if not text:
+        return None
+    try:
+        payload = loads_jsonish(text)
+    except (json.JSONDecodeError, ValueError, RecursionError):
+        payload = None
+    if payload == []:
+        return []
+    items = payload if isinstance(payload, list) else [payload]
+    calls = [call for item in items if (call := _json_tool_call(item, raw)) is not None]
+    if strict:
+        return calls if calls and len(calls) == len(items) else None
+    if calls:
+        return calls
+    extracted = extract_json_tool_fields(repair_jsonish(text))
+    if extracted is None:
+        return None
+    name, arguments = extracted
+    return [ParsedToolCall(name=name, arguments=arguments, raw=raw)]
+
+
+def parse_tool_call_blocks(
+    text: str, *, strict: bool = False
+) -> list[list[ParsedToolCall] | None]:
+    """The calls of each `<tool_call>` block in ``text``; None for a block that does not parse."""
+
+    blocks: list[list[ParsedToolCall] | None] = []
+    for block in TOOL_CALL_BLOCK_RE.finditer(text):
+        inner = block.group(1)
+        raw = block.group(0).strip()
+        calls: list[ParsedToolCall] = []
+        for function in FUNCTION_BLOCK_RE.finditer(inner):
+            arguments: dict[str, Any] = {}
+            for parameter in PARAMETER_BLOCK_RE.finditer(function.group(2)):
+                arguments[parameter.group(1).strip()] = parse_argument_value(
+                    parameter.group(2)
+                )
+            calls.append(
+                ParsedToolCall(
+                    name=function.group(1).strip(),
+                    arguments=json_ready(arguments),
+                    raw=raw,
+                )
+            )
+        blocks.append(calls if calls else parse_json_tool_calls(inner, raw, strict=strict))
+    return blocks
+
+
+def parse_tool_calls(text: str) -> list[ParsedToolCall]:
+    return [call for calls in parse_tool_call_blocks(text) if calls for call in calls]
+
+
+def openai_tool_calls_from_text(text: str) -> list[dict[str, Any]]:
+    encoded: list[dict[str, Any]] = []
+    for index, call in enumerate(parse_tool_calls(text)):
+        encoded.append(
+            {
+                "id": f"call_{index}_{call.name}",
+                "type": "function",
+                "function": {
+                    "name": call.name,
+                    "arguments": json.dumps(call.arguments, ensure_ascii=False),
+                },
+            }
+        )
+    return encoded
+
+
+def assistant_message_content(message: dict[str, Any]) -> str:
+    reasoning = message.get("reasoning_content")
+    content = message_text(message.get("content"))
+    if isinstance(reasoning, str) and reasoning.strip():
+        content = f"{THINK_OPEN}\n{reasoning.strip()}\n{THINK_CLOSE}\n{content}"
+    tool_xml = format_tool_calls_xml(message.get("tool_calls"))
+    if tool_xml:
+        if content.strip():
+            return f"{content.rstrip()}\n{tool_xml}"
+        return tool_xml
+    return content
+
+
+def _flush_tool_group(
+    group: list[str],
+    messages: list[dict[str, str]],
+) -> None:
+    if not group:
+        return
+    messages.append({"role": "user", "content": "\n".join(group)})
+    group.clear()
+
+
+def prepare_inference_messages(
+    raw_messages: list[dict[str, Any]],
+    *,
+    tools: Any | None = None,
+) -> list[dict[str, str]]:
+    if not raw_messages:
+        raise ValueError("Chat history cannot be empty")
+
+    prepared: list[dict[str, str]] = []
+    pending_tool_results: list[str] = []
+    for index, raw_message in enumerate(raw_messages):
+        if not isinstance(raw_message, dict):
+            raise ValueError(f"Unsupported chat message at index {index}")
+        raw_role = raw_message.get("role")
+        if not isinstance(raw_role, str):
+            raise ValueError(f"Unsupported chat role at index {index}: {raw_role!r}")
+        role = INFERENCE_ROLE_MAP.get(raw_role.strip().lower())
+        if role is None:
+            raise ValueError(f"Unsupported chat role at index {index}: {raw_role!r}")
+        if role == "assistant":
+            content = assistant_message_content(raw_message)
+        else:
+            content = message_text(raw_message.get("content"))
+        if role == "tool":
+            content = wrap_tool_response(content)
+            if not content.strip():
+                raise ValueError(f"Chat content at index {index} must be non-empty")
+            pending_tool_results.append(content)
+            continue
+        _flush_tool_group(pending_tool_results, prepared)
+        if not content.strip():
+            if role == "system":
+                continue
+            raise ValueError(f"Chat content at index {index} must be non-empty")
+        if role == "system":
+            content = strip_leaked_chatml(content)
+        prepared.append({"role": role, "content": content})
+    _flush_tool_group(pending_tool_results, prepared)
+
+    tools_block = render_available_tools(tools)
+    if tools_block:
+        if prepared and prepared[0]["role"] == "system":
+            prepared[0] = {
+                "role": "system",
+                "content": inject_available_tools(prepared[0]["content"], tools),
+            }
+        else:
+            prepared.insert(0, {"role": "system", "content": tools_block})
+
+    if not prepared:
+        raise ValueError("Chat history cannot be empty")
+    return prepared

tiny_gdn/web_search.py

@@ -0,0 +1,475 @@
+"""Execute Nemotron Agentic-v2 `web-search` the way Tercet-R stage-3 SFT saw it.
+
+`nvidia/Nemotron-SFT-Agentic-v2` `data/search.jsonl` (packed as
+`agentic_v2_search`) declares a single tool named `web-search` whose
+observations are compact Tavily Search JSON:
+
+    {"query", "follow_up_questions", "answer", "images", "results",
+     "response_time", "request_id"}
+
+Each `results[]` item is `{url, title, content, score, raw_content}`
+(about ten hits, with an LLM `answer` when Tavily supplies one). The chat
+template then prefixes that JSON with `<|tool_response|>`.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from collections.abc import Mapping, Sequence
+from html import unescape
+from html.parser import HTMLParser
+from typing import Any
+from urllib.error import HTTPError, URLError
+from urllib.parse import parse_qs, quote, unquote, urlencode, urlparse
+from urllib.request import Request, urlopen
+
+
+NEMOTRON_WEB_SEARCH_TOOL: dict[str, Any] = {
+    "type": "function",
+    "function": {
+        "name": "web-search",
+        "description": "Search the web for a query.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "Search query.",
+                }
+            },
+            "required": ["query"],
+        },
+    },
+}
+
+WEB_SEARCH_TOOL_NAME = "web-search"
+DEFAULT_MAX_RESULTS = 10
+TAVILY_SEARCH_URL = "https://api.tavily.com/search"
+DUCKDUCKGO_HTML_URL = "https://html.duckduckgo.com/html/"
+WIKIPEDIA_API_URL = "https://en.wikipedia.org/w/api.php"
+USER_AGENT = (
+    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
+    "(KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
+)
+REQUEST_TIMEOUT_SECONDS = 30.0
+
+
+def is_web_search_tool_name(name: str) -> bool:
+    return name.strip().lower().replace("_", "-") == WEB_SEARCH_TOOL_NAME
+
+
+def query_from_arguments(arguments: Mapping[str, Any] | None) -> str:
+    if not arguments:
+        return ""
+    raw = arguments.get("query")
+    if isinstance(raw, str):
+        return raw.strip()
+    if raw is None:
+        return ""
+    return str(raw).strip()
+
+
+def dumps_search_payload(payload: Mapping[str, Any]) -> str:
+    return json.dumps(canonical_search_payload(payload), ensure_ascii=False)
+
+
+def canonical_search_payload(payload: Mapping[str, Any]) -> dict[str, Any]:
+    results_raw = payload.get("results")
+    results: list[dict[str, Any]] = []
+    if isinstance(results_raw, Sequence) and not isinstance(results_raw, (str, bytes)):
+        for item in results_raw:
+            if isinstance(item, Mapping):
+                results.append(_canonical_result(item))
+    images_raw = payload.get("images")
+    images: list[Any] = []
+    if isinstance(images_raw, Sequence) and not isinstance(images_raw, (str, bytes)):
+        images = list(images_raw)
+    return {
+        "query": str(payload.get("query") or ""),
+        "follow_up_questions": payload.get("follow_up_questions"),
+        "answer": payload.get("answer"),
+        "images": images,
+        "results": results,
+        "response_time": payload.get("response_time"),
+        "request_id": payload.get("request_id"),
+    }
+
+
+def _canonical_result(item: Mapping[str, Any]) -> dict[str, Any]:
+    if item.get("error") and not item.get("url") and not item.get("title"):
+        return {"error": str(item["error"])}
+    result: dict[str, Any] = {
+        "url": str(item.get("url") or ""),
+        "title": str(item.get("title") or ""),
+        "content": str(item.get("content") or ""),
+        "score": _as_score(item.get("score")),
+        "raw_content": item.get("raw_content"),
+    }
+    if item.get("error"):
+        result["error"] = str(item["error"])
+    return result
+
+
+def _as_score(value: Any) -> float | None:
+    if value is None or value == "":
+        return None
+    try:
+        return float(value)
+    except (TypeError, ValueError):
+        return None
+
+
+def search_web(
+    query: str,
+    *,
+    max_results: int = DEFAULT_MAX_RESULTS,
+    api_key: str | None = None,
+) -> str:
+    """Return a Tavily-shaped JSON string for one `web-search` observation."""
+
+    cleaned = query.strip()
+    started = time.perf_counter()
+    if not cleaned:
+        return dumps_search_payload(
+            _empty_payload(
+                query,
+                error="Search query was empty",
+                response_time=time.perf_counter() - started,
+            )
+        )
+    key = (api_key if api_key is not None else os.environ.get("TAVILY_API_KEY", "")).strip()
+    if key:
+        try:
+            return dumps_search_payload(
+                _search_tavily(cleaned, api_key=key, max_results=max_results)
+            )
+        except Exception as error:
+            fallback_error = f"Tavily search failed: {error}"
+            try:
+                return dumps_search_payload(
+                    _search_without_tavily(
+                        cleaned,
+                        max_results=max_results,
+                        fallback_error=fallback_error,
+                    )
+                )
+            except Exception:
+                return dumps_search_payload(
+                    _empty_payload(
+                        cleaned,
+                        error=fallback_error,
+                        response_time=time.perf_counter() - started,
+                    )
+                )
+    try:
+        return dumps_search_payload(
+            _search_without_tavily(cleaned, max_results=max_results)
+        )
+    except Exception as error:
+        return dumps_search_payload(
+            _empty_payload(
+                cleaned,
+                error=str(error),
+                response_time=time.perf_counter() - started,
+            )
+        )
+
+
+def _empty_payload(
+    query: str,
+    *,
+    error: str,
+    response_time: float | None,
+) -> dict[str, Any]:
+    return {
+        "query": query,
+        "follow_up_questions": None,
+        "answer": None,
+        "images": [],
+        "results": [{"error": error}],
+        "response_time": response_time,
+        "request_id": None,
+    }
+
+
+def _search_tavily(
+    query: str,
+    *,
+    api_key: str,
+    max_results: int,
+) -> dict[str, Any]:
+    body = json.dumps(
+        {
+            "query": query,
+            "max_results": max(1, min(int(max_results), 20)),
+            "include_answer": True,
+            "include_raw_content": False,
+            "include_images": False,
+            "search_depth": "basic",
+        }
+    ).encode("utf-8")
+    payload = _read_json(
+        TAVILY_SEARCH_URL,
+        data=body,
+        headers={
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+        },
+        method="POST",
+    )
+    if not isinstance(payload, dict):
+        raise RuntimeError("Tavily returned a non-object payload")
+    payload.setdefault("query", query)
+    return payload
+
+
+def _search_without_tavily(
+    query: str,
+    *,
+    max_results: int,
+    fallback_error: str | None = None,
+) -> dict[str, Any]:
+    started = time.perf_counter()
+    results: list[dict[str, Any]] = []
+    errors: list[str] = []
+    if fallback_error:
+        errors.append(fallback_error)
+    try:
+        results.extend(_search_duckduckgo_html(query, max_results=max_results))
+    except Exception as error:
+        errors.append(f"DuckDuckGo search failed: {error}")
+    if len(results) < max_results:
+        try:
+            wiki = _search_wikipedia(query, max_results=max_results)
+            seen = {item["url"] for item in results}
+            for item in wiki:
+                if item["url"] in seen:
+                    continue
+                results.append(item)
+                if len(results) >= max_results:
+                    break
+        except Exception as error:
+            errors.append(f"Wikipedia search failed: {error}")
+    elapsed = time.perf_counter() - started
+    if not results:
+        message = "; ".join(errors) if errors else "Search returned no results"
+        return _empty_payload(query, error=message, response_time=elapsed)
+    clipped = results[: max(1, min(int(max_results), 20))]
+    for index, item in enumerate(clipped):
+        if item.get("score") is None:
+            item["score"] = round(max(0.05, 1.0 - (index * 0.07)), 8)
+        item.setdefault("raw_content", None)
+    answer = clipped[0].get("content") if clipped else None
+    if isinstance(answer, str) and not answer.strip():
+        answer = None
+    return {
+        "query": query,
+        "follow_up_questions": None,
+        "answer": answer,
+        "images": [],
+        "results": clipped,
+        "response_time": elapsed,
+        "request_id": None,
+    }
+
+
+def _search_duckduckgo_html(query: str, *, max_results: int) -> list[dict[str, Any]]:
+    encoded = urlencode({"q": query, "kl": "us-en"})
+    html = _read_text(
+        DUCKDUCKGO_HTML_URL,
+        data=encoded.encode("utf-8"),
+        headers={
+            "Content-Type": "application/x-www-form-urlencoded",
+            "Accept": "text/html",
+        },
+        method="POST",
+    )
+    parser = _DuckDuckGoParser()
+    parser.feed(html)
+    parser.close()
+    results: list[dict[str, Any]] = []
+    seen: set[str] = set()
+    for item in parser.results:
+        url = item["url"]
+        if not url or url in seen:
+            continue
+        seen.add(url)
+        results.append(
+            {
+                "url": url,
+                "title": item["title"],
+                "content": item["content"],
+                "score": None,
+                "raw_content": None,
+            }
+        )
+        if len(results) >= max_results:
+            break
+    return results
+
+
+def _search_wikipedia(query: str, *, max_results: int) -> list[dict[str, Any]]:
+    params = urlencode(
+        {
+            "action": "query",
+            "list": "search",
+            "srsearch": query,
+            "srlimit": max(1, min(int(max_results), 10)),
+            "srprop": "snippet",
+            "format": "json",
+            "utf8": "1",
+        }
+    )
+    payload = _read_json(
+        f"{WIKIPEDIA_API_URL}?{params}",
+        headers={"Accept": "application/json"},
+        method="GET",
+    )
+    if not isinstance(payload, dict):
+        return []
+    query_block = payload.get("query")
+    hits = query_block.get("search") if isinstance(query_block, dict) else None
+    if not isinstance(hits, list):
+        return []
+    results: list[dict[str, Any]] = []
+    for hit in hits:
+        if not isinstance(hit, dict):
+            continue
+        title = str(hit.get("title") or "").strip()
+        if not title:
+            continue
+        snippet = unescape(_strip_tags(str(hit.get("snippet") or ""))).strip()
+        results.append(
+            {
+                "url": f"https://en.wikipedia.org/wiki/{quote(title.replace(' ', '_'))}",
+                "title": title,
+                "content": snippet,
+                "score": None,
+                "raw_content": None,
+            }
+        )
+    return results
+
+
+def unwrap_ddg_url(href: str) -> str:
+    raw = href.strip()
+    if raw.startswith("//"):
+        raw = f"https:{raw}"
+    parsed = urlparse(raw)
+    host = parsed.netloc.lower()
+    if "duckduckgo.com" in host:
+        values = parse_qs(parsed.query).get("uddg") or []
+        if values:
+            return unquote(values[0])
+    return raw
+
+
+def _strip_tags(value: str) -> str:
+    return _TagStripper.strip(value)
+
+
+class _TagStripper(HTMLParser):
+    def __init__(self) -> None:
+        super().__init__()
+        self.chunks: list[str] = []
+
+    def handle_data(self, data: str) -> None:
+        self.chunks.append(data)
+
+    @classmethod
+    def strip(cls, value: str) -> str:
+        parser = cls()
+        parser.feed(value)
+        parser.close()
+        return "".join(parser.chunks)
+
+
+class _DuckDuckGoParser(HTMLParser):
+    def __init__(self) -> None:
+        super().__init__()
+        self.results: list[dict[str, str]] = []
+        self._in_title = False
+        self._in_snippet = False
+        self._href = ""
+        self._buffer: list[str] = []
+        self._pending: dict[str, str] | None = None
+
+    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
+        classes = {item for item in (dict(attrs).get("class") or "").split() if item}
+        if tag == "a" and "result__a" in classes:
+            self._flush_pending()
+            self._in_title = True
+            self._href = dict(attrs).get("href") or ""
+            self._buffer = []
+            return
+        if tag in {"a", "div"} and "result__snippet" in classes:
+            self._in_snippet = True
+            self._buffer = []
+
+    def handle_endtag(self, tag: str) -> None:
+        if tag == "a" and self._in_title:
+            title = unescape("".join(self._buffer)).strip()
+            url = unwrap_ddg_url(self._href)
+            self._pending = {"title": title, "url": url, "content": ""}
+            self._in_title = False
+            self._buffer = []
+            return
+        if self._in_snippet and tag in {"a", "div"}:
+            snippet = unescape("".join(self._buffer)).strip()
+            if self._pending is not None:
+                self._pending["content"] = snippet
+                self.results.append(self._pending)
+                self._pending = None
+            self._in_snippet = False
+            self._buffer = []
+
+    def handle_data(self, data: str) -> None:
+        if self._in_title or self._in_snippet:
+            self._buffer.append(data)
+
+    def close(self) -> None:
+        self._flush_pending()
+        super().close()
+
+    def _flush_pending(self) -> None:
+        if self._pending is None:
+            return
+        self.results.append(self._pending)
+        self._pending = None
+
+
+def _read_json(
+    url: str,
+    *,
+    data: bytes | None = None,
+    headers: Mapping[str, str],
+    method: str,
+) -> Any:
+    text = _read_text(url, data=data, headers=headers, method=method)
+    return json.loads(text)
+
+
+def _read_text(
+    url: str,
+    *,
+    data: bytes | None = None,
+    headers: Mapping[str, str],
+    method: str,
+) -> str:
+    merged = {"User-Agent": USER_AGENT, **dict(headers)}
+    request = Request(url, data=data, headers=merged, method=method)
+    try:
+        with urlopen(request, timeout=REQUEST_TIMEOUT_SECONDS) as response:
+            raw = response.read()
+            charset = response.headers.get_content_charset() or "utf-8"
+    except HTTPError as error:
+        detail = error.read().decode("utf-8", errors="replace")
+        raise RuntimeError(
+            f"{error.code} {error.reason} from {url}: {detail[:500]}"
+        ) from error
+    except URLError as error:
+        raise RuntimeError(f"Could not reach {url}: {error.reason}") from error
+    return raw.decode(charset, errors="replace")

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,31 @@
+{
+  "checkpoint": "checkpoint-00000704",
+  "ema": {
+    "elapsed_seconds": 54.355851609958336,
+    "logps_chosen": -601.3653792912514,
+    "logps_rejected": -652.0577934662236,
+    "loss": 0.8922630808546166,
+    "nll_chosen_per_token": 0.5581774069041227,
+    "pairs": 903,
+    "response_tokens": 2363563,
+    "rewards_accuracy": 0.7895902547065338,
+    "rewards_chosen": -1.6784689598041251,
+    "rewards_margin": 1.3302929308987403,
+    "rewards_rejected": -3.008761755519795
+  },
+  "normal": {
+    "elapsed_seconds": 53.85813916684128,
+    "logps_chosen": -601.1099114064231,
+    "logps_rejected": -651.828419158361,
+    "loss": 0.8913962870077173,
+    "nll_chosen_per_token": 0.5579402858352829,
+    "pairs": 903,
+    "response_tokens": 2363563,
+    "rewards_accuracy": 0.7918050941306756,
+    "rewards_chosen": -1.6656929430111433,
+    "rewards_margin": 1.331598123442691,
+    "rewards_rejected": -2.997290931270764
+  },
+  "optimizer_step": 704,
+  "type": "preference_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',
+]