Upload LibraMind Mini weights, tokenizer, and inference code

Adds model.safetensors (bf16), tokenizer.json, config.json, model.py, inference.py, chat_format.py, and the model card.

README.md

@@ -1,3 +1,329 @@
 ---
 license: apache-2.0
+language:
+  - en
+library_name: pytorch
+pipeline_tag: text-generation
+tags:
+  - text-generation
+  - conversational
+  - chat
+  - function-calling
+  - roleplay
+  - hybrid
+  - gated-deltanet
+  - linear-attention
+  - muon-optimizer
+  - research
+  - open-weights
+# datasets:
+#   - HuggingFaceFW/fineweb
+#   - HuggingFaceTB/smoltalk
 ---
+
+# LibraMind (565M parameters)
+
+**LibraMind** is a 565M-parameter chat model pretrained from scratch by [alby13](https://github.com/alby13/LibraMind-AI) on a single consumer GPU (NVIDIA RTX 4090). It is a hybrid architecture that interleaves **18 Gated DeltaNet** linear-recurrent layers with **6 Gated Attention** layers (a repeating `DDDA × 6` pattern), with Block Attention Residuals and logit soft-capping, trained with the Muon optimizer. It went through a full five-stage pipeline: pretraining, midtraining, supervised fine-tuning, preference tuning, and RL with verifiable rewards.
+
+
+> **Research model: no safety training.**
+> LibraMind has not undergone safety training or red-teaming. It can produce inaccurate, biased, offensive, or harmful text, and it will not reliably refuse requests. It also makes factual and arithmetic mistakes (see the examples below). It is released as a research artifact, provided "as is" without warranty of any kind, and the authors accept no liability for its outputs or for any use of the model. You are solely responsible for evaluating it and for any safeguards, filtering, and human oversight your use case needs. Do not deploy it in user-facing, high-stakes, or safety-critical settings without your own safety work.
+> This notice is guidance. It does not add to or modify the terms of the Apache-2.0 license.
+
+## Model Summary
+
+| | |
+|---|---|
+| **Developer** | alby13 |
+| **Parameters** | 565M (565.2M) |
+| **Model type** | Chat model |
+| **Architecture** | Hybrid Gated DeltaNet + Gated Attention, decoder-only |
+| **Layers / width** | 24 layers, hidden size 1,024 (`DDDA × 6`: 3 DeltaNet layers then 1 Attention layer, repeated 6 times) |
+| **Vocabulary** | 32,768 (custom byte-level BPE) |
+| **Context length** | 4,096 tokens (pretrained at 2,048, extended during midtraining) |
+| **Weights precision** | bf16 (about 1.1 GB) |
+| **Language** | English |
+| **License** | Apache-2.0 |
+| **Release** | 10/6/2026 |
+
+## Intended use
+
+- Light assistant tasks: casual conversation, rewriting and summarizing short texts, formatted answers such as lists, sections and word limits, and simple function calling.
+- Hobby and research: a fully documented small-model training run (architecture, data, every stage, evaluations).
+- Game NPCs and interactive characters: dialogue driven by a character card in the system prompt, with optional schema-forced JSON for emotions and actions.
+
+## Intended Use
+
+**Intended for:**
+- Research on small hybrid recurrent/attention language models
+- Studying low-resource, single-GPU training pipelines (pretraining through RL) and the Muon optimizer at small scale
+- Experiments with on-device chat, role-play, and structured output (JSON / tool calls), with your own validation of the results (for example, validate JSON and compute numbers in code)
+
+**Not intended for:**
+- High-stakes or safety-critical decisions (medical, legal, financial, etc.)
+- Unsupervised, user-facing deployment: the model has no safety training
+- Use as a source of factual information without independent verification
+
+## How to Use
+
+LibraMind uses a custom architecture, so it doesn't load in transformers, llama.cpp or other GGUF runtimes. It needs this repository's model.py, inference.py and chat_format.py, plus:
+
+PyTorch with CUDA (trained with 2.14)
+flash-linear-attention 0.5.2, with Triton
+tokenizers
+llguidance (optional, for grammar-forced JSON)
+In this repository, the weights are runs/rl1/model_final.pt (float32, 2.26 GB) and the tokenizer is data/chat/sft/tokenizer.json.
+
+```python
+import torch
+from tokenizers import Tokenizer
+
+from chat_format import ChatEncoder
+from inference import GrammarFactory, generate
+from model import LM, ModelConfig
+
+ck = torch.load("model_final.pt", map_location="cpu", weights_only=False)
+model = LM(ModelConfig(**{**ck["model_config"], "grad_ckpt": False}))
+model.load_state_dict(ck["model"])
+model = model.cuda().to(torch.bfloat16).eval()
+enc = ChatEncoder(Tokenizer.from_file("tokenizer.json"))
+
+
+def chat(messages, schema=None, max_new=300):
+    matcher = [GrammarFactory("tokenizer.json", enc.im_end, model.config.vocab_size).json(schema)] if schema else None
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        out = generate(model, [enc.encode_prompt(messages)], max_new, temperature=0.7, top_p=0.9, rep_penalty=1.1,
+                       stop_ids=[enc.im_end], matchers=matcher, vocab=enc.tok.get_vocab_size())[0]
+    return enc.tok.decode(out)
+
+
+print(chat([{"role": "system", "content": "You are Brom, a gruff dwarven blacksmith. Stay in character."},
+            {"role": "user", "content": "Can you fix my sword?"}]))
+```
+
+generate uses cached decoding: Gated DeltaNet recurrent state plus an attention KV cache, replayed as a CUDA graph. It accepts many prompts at once, at about 1,050 tokens/s across 64 parallel conversations. The repository also has a browser chat window (chatui.cmd) and a console chat (chat.cmd).
+
+Recommended settings:
+
+- temperature 0.4, top-p 0.9, repetition penalty 1.1
+- System prompt: 
+````
+"You are LibraMind, a friendly and thoughtful conversation partner made by alby13. Keep replies natural and to the point: a few sentences for casual chat, more detail only when it's needed. Ask a follow-up question when it helps the conversation.
+````
+- temperature 0 for JSON and tool calls
+- always pass stop_ids=[<|im_end|>]
+
+**Requirements:** 
+
+**Memory:** about 1.1 GB for the weights in bf16, plus activations. The recurrent layers use constant-size state, and only the 6 attention layers grow a KV cache with context length.
+
+
+## Chat format
+
+ChatML with a beginning-of-text token. Roles are system, user, assistant and tool:
+
+```
+<|bos|><|im_start|>system
+{system prompt}<|im_end|>
+<|im_start|>user
+{message}<|im_end|>
+<|im_start|>assistant
+{reply}<|im_end|>
+```
+
+The system prompt is optional. Without one, the model behaves as the LibraMind assistant.
+
+
+## Characters and Personas
+
+Put the character card in the system prompt: name, setting, personality, how they speak, and any facts they must get right. That means prices, quest details, names and places. Roleplay training used this pattern:
+
+```
+You are {name}. Stay in character: speak as {name} would, in the first person, and never say you are an AI.
+
+Character profile:
+{profile}
+```
+
+For game actions, force the reply into a JSON schema with GrammarFactory.json(schema). The output is then guaranteed to parse and match the schema. Example schema: examples/npc_schema.json, with fields line, emotion and action. Describe the fields in the system prompt too, so the content makes sense as well as the format.
+
+## Function calling
+
+```
+List the tools in the system prompt in the format the model was trained on:
+
+You can call the following functions:
+<tools>
+[{"type": "function", "function": {"name": "get_weather", "description": "...", "parameters": {...}}}]
+</tools>
+
+To call a function, reply with a JSON object inside <tool_call></tool_call> tags, for example:
+<tool_call>
+{"name": "function_name", "arguments": {"arg": "value"}}
+</tool_call>
+You may make several calls. Function results come back in <tool_response> tags.
+```
+
+GrammarFactory.tool_calls(tools) can force calls to name a listed function, with arguments that match its schema.
+
+
+## Example Outputs
+
+sampling settings used (temperature 0.4, top-p 0.9, repetition penalty 1.1).
+
+**General chat**
+
+> **User:** Hi! Who are you?
+> **LibraMind:** I'm LibraMind, a small AI language model trained from scratch by alby13. How can I help?
+
+**Role-play**
+
+> **System Prompt:** You are Grom, a gruff dwarven blacksmith in the mountain town of Bringham. Stay in character and keep replies short.
+> **User:** Where can I find some armor? I need some that also protects against magic.
+> **LibraMind:** Well, I ain't no scholar, but I've heard tales from travelers. You might find some armor in the blacksmith's forge. Ask the old sage, Thorne. He knows all about such things.
+
+> **System Prompt:** You are Jin, the stern innkeeper of the Gilded Redbird inn in the village of Ashfield. A room costs 15 silver a night and a bowl of stew costs 3 silver. (...)
+> **User:** How much for a room tonight and some supper?
+> **LibraMind:** A room for 15 silver is a good price. A bowl of stew for 3 silver will be 9 silver.
+
+
+
+**Structured output** (schema-forced decoding; the system prompt asks for JSON with `line`, `emotion`, and `action`)
+
+> **System:** You are Grom (...). Reply only with JSON: "line" is what you say, "emotion" is how you feel, and "action" is what you do (open_shop if the player wants to buy or get something repaired).
+> **User:** I'd like to buy a new axe.
+
+```
+Sample 1: {"line": "axe", "emotion": "happy", "action": "open_shop"}
+Sample 2: {"line": "You've come to the right place, friend. I've got just the thing for you. What's your name, and what kind of axe are you looking for?", "emotion": "neutral", "action": "open_shop"}
+```
+
+**Instruction following**
+
+> **User:** Give me three tips for memorizing for studying as a numbered list.
+> **LibraMind:**
+>1. Create a dedicated study space: Having a specific area for studying can help you focus and stay motivated.
+>2. Use mnemonic devices: These are memory techniques that can help you remember information by associating it with something you already know or have in mind.
+>3. Break your study sessions into smaller intervals: This can help maintain your concentration and prevent burnout.
+
+
+## Architecture
+
+The layout follows Qwen3.5-style hybrid models: three linear-attention layers for every full-attention layer. The plain residual stream is replaced by Block Attention Residuals.
+
+| Component | Details |
+|---|---|
+| **Gated DeltaNet** (18 layers) | Constant-memory linear recurrent state using the delta rule, with a short 1D convolution (kernel 4) and output gating. About 10.5M parameters per layer. |
+| **Gated Attention** (6 layers, every 4th) | Grouped-query attention: 16 query heads × 128 dim, 4 shared KV heads, QK-Norm, RoPE (base 10,000), and a sigmoid output gate. About 7.3M parameters per layer. Only these 6 layers keep a KV cache. |
+| **SwiGLU MLP** (all 24 layers) | Hidden width 3,584, about 11.0M parameters per layer. |
+| **Block Attention Residuals** | 8 blocks of 6 sub-layers; each sub-layer attends over summaries of earlier blocks with a learned query. About 50K parameters in total. |
+| **Normalization / bias** | Pre-RMSNorm without learnable scale; no bias terms anywhere. |
+| **Embeddings / head** | Untied: 33.6M input embedding + 33.6M output head. |
+| **Logit soft-cap** | `logits = 15 · tanh(logits / 15)` |
+| **Tokenizer** | Byte-level BPE, 32,768 vocab, GPT-4-style regex splitting, about 4.7 bytes per token. Trained on 2B characters of web text. |
+
+Parameter split: about 498.1M in the transformer body (24 layers) and 67.1M in embeddings and output head.
+
+## Training
+
+| Stage | Data
+|---|---|
+| 1. Pretraining | Enhanced FineWeb text
+| 2. Midtraining | 50% unseen FineWeb, 50% conversations
+| 3. Supervised fine-tuning (SFT)
+| 4. Preference tuning
+| 5. RL with verifiable rewards (GRPO)
+
+
+**Pretraining details:**
+- **Optimizers:** Muon for the ~497M 2D matrix weights (LR 0.02, cautious weight decay 0.07 on a cosine schedule); AdamW (LR 0.001) for embeddings, output head, gates, and Block Attention Residual queries
+- **Schedule:** 40-step warmup, constant, then linear decay to 5% over the last 65% of training
+- **Precision / compute:** bf16, `torch.compile`, gradient clipping at 1.0, about 17.1 GB VRAM on a single RTX 4090
+
+## Evaluation
+
+### Pretrained base model
+
+Zero-shot evaluation with [lighteval](https://github.com/huggingface/lighteval) in cloze form (accuracy normalized by length; Winogrande and LAMBADA use plain accuracy). The SmolLM2-360M column is from its model card, which reports the same style of evaluation. SmolLM2-360M saw about 4 trillion training tokens, roughly 400× more than LibraMind.
+
+| Benchmark | LibraMind base (10.2B tokens) | SmolLM2-360M (4T tokens) |
+|---|---|---|
+| HellaSwag | 55.3 | 54.5 |
+| ARC (Easy / Challenge) | 65.6 / 35.4 (average 50.5) | average 53.0 |
+| PIQA | 73.9 | 71.7 |
+| OpenBookQA | 31.2 | 37.4 |
+| CommonsenseQA | 40.6 | 38.0 |
+| Social IQa | 44.6 | — |
+| Winogrande | 56.6 | 52.5 |
+| MMLU (cloze) | 30.3 | 35.8 |
+| LAMBADA (accuracy / perplexity) | 46.3 / 13.3 | — |
+
+### Chat model, by training stage
+
+- **IFEval:** 541 prompts, using the lm-eval implementation.
+- **Tool calling:** 200 held-out function-calling conversations from the SFT sources. "Call decision" means it correctly chose whether to call a tool. Numbers in parentheses use grammar-forced (constrained) decoding.
+- **Validation loss:** 1,800 held-out conversations.
+
+| | mid | SFT | preference | RL (released) |
+|---|---|---|---|---|
+| IFEval prompt-level strict / loose | 40.7 / 42.3 | 45.8 / 49.7 | 50.5 / 53.8 | 51.4 / 54.5 |
+| IFEval instruction-level strict / loose | 54.1 / 56.2 | 57.6 / 61.0 | 63.0 / 66.0 | 64.0 / 66.7 |
+| Tools: correct call decision | 99.5 | 99.0 | 99.5 | 99.5 |
+| Tools: valid JSON | 99.0 | 98.5 | 99.5 | 99.5 |
+| Tools: right function (grammar-forced) | 87.8 (88.8) | 91.3 (92.9) | 92.4 (92.9) | 92.4 (92.4) |
+| Tools: exact arguments | 71.4 | 80.1 | 81.1 | 80.6 |
+| Validation loss, conversations | 1.035 | 0.950 | 0.966 | 0.968 |
+
+On 40 held-out role-play prompts, the released model never described itself as an AI. Its average reply on general prompts is 156 words (175 after SFT).
+
+**IFEval caveat:** these scores flatter the model's general instruction-following. Its training data includes IFEval-style constraint exercises, and the RL stage rewarded the same 19 instruction checkers that IFEval scores (on different prompts). For reference, SmolLM2-360M-Instruct reports 41.0 on IFEval (the average of its prompt- and instruction-level scores). LibraMind's equivalent strict average is 57.7.
+
+## Limitations and Bias
+
+- Facts: it often states wrong or invented facts confidently. For example, it described the sitcom Last of the Summer Wine as a science-fiction series. Don't use it as a source of information.
+- Math: it fails at arithmetic beyond single digits (see above). Part of the cause is the tokenizer, which splits numbers into irregular 1–3-digit chunks (6340 → 6|34|0, 12452 → 124|52). Do calculations in code, or give the model a calculator tool.
+- Code: generated code is rarely correct.
+- Using provided information: it usually uses facts from the system prompt, but can garble them, especially numbers. In game use, validate anything that matters.
+- Roleplay: it stays in character but can drift from the character's details, hedge ("my expertise lies elsewhere, but…") or invent backstory. Long conversations lose coherence, and the context limit is 4,096 tokens.
+- Language: English only.
+- Safety: there was no dedicated safety training or red-teaming. Any refusal behavior comes only from the fine-tuning and preference data. Like any web-trained model, it can produce biased, offensive or harmful text. Filter outputs before showing them to players or the public.
+
+## Training Data and Licenses
+
+- **FineWeb:** released under the permissive ODC-By 1.0 (Open Data Commons Attribution) license.
+- **SmolTalk, smol-smoltalk, and SmolTalk2:** Apache 2.0
+- **GSM8K and MMLU:** MIT.
+
+## License
+
+The model weights and code are released under the **Apache License 2.0**; see the `LICENSE` file. Apache-2.0 includes a disclaimer of warranty and a limitation of liability (Sections 7 and 8). Training data carries its own terms; see above.
+
+## Acknowledgements
+
+LibraMind builds on public research and open tools:
+
+- Gated DeltaNet (Yang, Kautz and Hatamizadeh, 2024) and the flash-linear-attention library
+- The Qwen3.5 hybrid layout and gated attention
+- Block Attention Residuals from the Kimi team at Moonshot AI
+- The Muon optimizer (Keller Jordan, 2024)
+- Anchored Preference Optimization (D'Oosterlinck et al., 2024)
+- GRPO (DeepSeekMath, 2024)
+- IFEval (Zhou et al., 2023)
+- lighteval
+- The SmolLM2/SmolLM3 data and recipes from Hugging Face, Tulu 3 from Ai2, GSM8K, and MMLU
+
+## Citation
+
+```bibtex
+@misc{alby13_libramind_2026,
+  author    = {alby13},
+  title     = {LibraMind: A 565M Parameter Hybrid Gated DeltaNet-Attention Chat Model},
+  year      = {2026},
+  publisher = {Local AI Research},
+  url       = {https://github.com/alby13/LibraMind-AI}
+}
+```
+
+## Contact
+
+GitHub issues discussions or Hugging Face Discussions tab.

chat_format.py

@@ -0,0 +1,57 @@
+"""Chat template shared by data preparation, fine-tuning and the chat program (ChatML style).
+
+  <|bos|><|im_start|>system\\n{system}<|im_end|>\\n<|im_start|>user\\n{text}<|im_end|>\\n<|im_start|>assistant\\n{reply}<|im_end|>\\n
+
+Messages are encoded piece by piece, so training and inference tokenize identically.
+Loss mask: 1 on assistant reply tokens and the <|im_end|> that closes them (so the model learns to stop).
+"""
+import numpy as np
+from tokenizers import Tokenizer
+
+BOS, IM_START, IM_END = "<|bos|>", "<|im_start|>", "<|im_end|>"
+ROLES = ("system", "user", "assistant", "tool")
+
+
+def chat_tokenizer(base_tokenizer_path):
+    """The pretraining tokenizer plus the chat special tokens (appended as new ids)."""
+    tok = Tokenizer.from_file(str(base_tokenizer_path))
+    tok.add_special_tokens([IM_START, IM_END])
+    return tok
+
+
+class ChatEncoder:
+    def __init__(self, tok):
+        self.tok = tok
+        self.bos, self.im_start, self.im_end = (tok.token_to_id(t) for t in (BOS, IM_START, IM_END))
+        self.newline = tok.encode("\n", add_special_tokens=False).ids
+        self.headers = {r: tok.encode(f"{r}\n", add_special_tokens=False).ids for r in ROLES}
+
+    def encode_conversation(self, messages):
+        """messages: list of {"role", "content"}. Returns (ids uint16, loss mask uint8)."""
+        ids, mask = [self.bos], [0]
+        for m in messages:
+            head = [self.im_start] + self.headers[m["role"]]
+            body = self.tok.encode(m["content"], add_special_tokens=False).ids + [self.im_end]
+            train = 1 if m["role"] == "assistant" else 0
+            ids += head + body + self.newline
+            mask += [0] * len(head) + [train] * len(body) + [0] * len(self.newline)
+        return np.array(ids, dtype=np.uint16), np.array(mask, dtype=np.uint8)
+
+    def encode_prompt(self, messages):
+        """Conversation so far plus the assistant header, ready for generation."""
+        ids, _ = self.encode_conversation(messages)
+        return np.concatenate([ids, [self.im_start] + self.headers["assistant"]]).astype(np.int64)
+
+    def encode_reply(self, text):
+        """An assistant reply as the model would produce it: its tokens followed by <|im_end|>."""
+        return np.array(self.tok.encode(text, add_special_tokens=False).ids + [self.im_end], dtype=np.int64)
+
+
+def decode_conversation(tok, ids):
+    """Turn a tokenized conversation back into [{"role", "content"}] messages."""
+    text = tok.decode([int(i) for i in ids], skip_special_tokens=False).replace(BOS, "")
+    msgs = []
+    for part in text.split(IM_START)[1:]:
+        role, _, body = part.partition("\n")
+        msgs.append({"role": role, "content": body.split(IM_END)[0]})
+    return msgs

config.json

@@ -0,0 +1,18 @@
+{
+  "vocab_size": 32832,
+  "n_layer": 24,
+  "d_model": 1024,
+  "arch": "hybrid",
+  "attn_every": 4,
+  "n_head": 16,
+  "n_kv_head": 4,
+  "head_dim": 128,
+  "gdn_heads": 16,
+  "gdn_head_dim": 128,
+  "ffn_dim": 3584,
+  "attnres_blocks": 8,
+  "rope_theta": 10000.0,
+  "max_seq_len": 4096,
+  "softcap": 15.0,
+  "grad_ckpt": false
+}

inference.py

@@ -0,0 +1,335 @@
+"""Fast generation: cached incremental decoding, batched sampling, optional grammar constraints.
+
+Each sequence's state is small and fixed per token: attention layers keep their keys/values, DeltaNet
+layers keep their recurrent state (flash-linear-attention cache), and Block AttnRes only mixes within
+a token, so it needs no cache. A prompt is processed once (prefill) and each new token then costs one
+cheap step instead of re-running the whole conversation.
+
+  outs = generate(model, [prompt_ids, ...], max_new=200, stop_ids=[enc.im_end])
+  matcher = GrammarFactory(tokenizer_path, im_end_id).json(schema)   # output guaranteed to parse
+"""
+import torch
+import torch.nn.functional as F
+
+from model import DeltaNet, GatedAttention, attnres_mix, inv_rms, rms
+
+
+class _GDNStates:
+    """The minimal cache interface flash-linear-attention layers expect (len / [] / update)."""
+
+    def __init__(self, n):
+        self.layers = [None] * n
+
+    def __len__(self):
+        return len(self.layers)
+
+    def __getitem__(self, i):
+        return self.layers[i]
+
+    def update(self, layer_idx, recurrent_state=None, conv_state=None, **_):
+        cur = self.layers[layer_idx]
+        if cur is None:  # first call (prefill): keep the tensors
+            self.layers[layer_idx] = dict(recurrent_state=recurrent_state, conv_state=conv_state)
+            return
+        # later calls: write into the existing buffers, so a replayed CUDA graph carries the state forward
+        pairs = [(cur["recurrent_state"], recurrent_state)] + list(zip(cur["conv_state"] or (), conv_state or ()))
+        for dst, src in pairs:
+            if src is not None and dst.data_ptr() != src.data_ptr():
+                dst.copy_(src)
+
+
+class Cache:
+    """Decoding state for a batch: attention keys/values (preallocated to `capacity`), DeltaNet states,
+    and each row's current length (rows may have different lengths)."""
+
+    def __init__(self, model, batch, capacity):
+        dev, dt = model.embed.weight.device, torch.bfloat16
+        self.pos = torch.zeros(batch, dtype=torch.long, device=dev)
+        self.max_pos, self.capacity = 0, capacity
+        self.kv = {j: (torch.zeros(batch, s.nkv, capacity, s.hd, device=dev, dtype=dt),
+                       torch.zeros(batch, s.nkv, capacity, s.hd, device=dev, dtype=dt))
+                   for j, s in enumerate(model.sublayers) if isinstance(s, GatedAttention)}
+        self.gdn = _GDNStates(len(model.sublayers))
+
+    @staticmethod
+    def stack(caches, extra):
+        """Combine single-sequence caches of different lengths into one batch with room for `extra` tokens."""
+        out = Cache.__new__(Cache)
+        out.pos = torch.cat([c.pos for c in caches])
+        out.max_pos = max(c.max_pos for c in caches)
+        out.capacity = out.max_pos + extra
+        out.kv = {}
+        for j, (K0, _) in caches[0].kv.items():
+            K = K0.new_zeros(len(caches), K0.size(1), out.capacity, K0.size(3))
+            V = torch.zeros_like(K)
+            for i, c in enumerate(caches):
+                n = c.max_pos
+                K[i, :, :n], V[i, :, :n] = c.kv[j][0][0, :, :n], c.kv[j][1][0, :, :n]
+            out.kv[j] = (K, V)
+        out.gdn = _GDNStates(len(caches[0].gdn))
+        for i, layer in enumerate(caches[0].gdn.layers):
+            if layer is not None:
+                conv = layer["conv_state"]
+                out.gdn.layers[i] = dict(
+                    recurrent_state=torch.cat([c.gdn.layers[i]["recurrent_state"] for c in caches]),
+                    conv_state=None if conv is None else tuple(
+                        torch.cat([c.gdn.layers[i]["conv_state"][k] for c in caches]) for k in range(len(conv))))
+        return out
+
+
+def _attention(sub, x, K, V, pos):
+    """GatedAttention over cached keys/values; x holds t new tokens per row starting at position pos[row]."""
+    B, t, _ = x.shape
+    q, gate = sub.q_proj(x).split(sub.nh * sub.hd, -1)
+    q = q.view(B, t, sub.nh, sub.hd)
+    k, v = sub.kv_proj(x).view(B, t, 2, sub.nkv, sub.hd).unbind(2)
+    positions = pos[:, None] + torch.arange(t, device=x.device)            # [B, t]
+    cos, sin = (b[0, :, 0][positions].unsqueeze(2).to(x.dtype) for b in (sub.cos, sub.sin))
+
+    def rope(z):
+        z1, z2 = z.chunk(2, -1)
+        return torch.cat([z1 * cos - z2 * sin, z1 * sin + z2 * cos], -1)
+
+    q, k = rope(rms(q)), rope(rms(k))
+    rows = torch.arange(B, device=x.device)[:, None].expand(B, t)
+    K[rows, :, positions] = k.to(K.dtype)
+    V[rows, :, positions] = v.to(V.dtype)
+    keys = K.repeat_interleave(sub.nh // sub.nkv, dim=1)
+    vals = V.repeat_interleave(sub.nh // sub.nkv, dim=1)
+    mask = torch.arange(K.size(2), device=x.device)[None, None, :] <= positions[:, :, None]  # causal per row
+    y = F.scaled_dot_product_attention(q.transpose(1, 2), keys.to(q.dtype), vals.to(q.dtype), attn_mask=mask[:, None])
+    y = y.transpose(1, 2).reshape(B, t, sub.nh * sub.hd)
+    return sub.o_proj(y * torch.sigmoid(gate))
+
+
+@torch.no_grad()
+def step(model, idx, cache):
+    """Feed t new tokens per row (idx [B, t]); returns next-token logits [B, vocab] for the last position."""
+    t = idx.size(1)
+    if cache.max_pos + t > cache.capacity:
+        raise ValueError(f"cache full ({cache.capacity} tokens)")
+    x0 = rms(model.embed(idx))
+    blocks, blocks_inv, partial = [x0], [inv_rms(x0)], None
+    for j, sub in enumerate(model.sublayers):
+        srcs = blocks if partial is None else blocks + [partial]
+        inv = blocks_inv if partial is None else blocks_inv + [inv_rms(partial)]
+        h = rms(attnres_mix(srcs, inv, model.attnres_queries[j]))
+        if isinstance(sub, GatedAttention):
+            out = _attention(sub, h, *cache.kv[j], cache.pos)
+        elif isinstance(sub, DeltaNet):
+            sub.gdn.layer_idx = j
+            out = sub.gdn(h, past_key_values=cache.gdn, use_cache=True)[0]
+        else:
+            out = sub(h)
+        out = out.float()
+        partial = out if partial is None else partial + out
+        if (j + 1) % model.block_size == 0:
+            blocks.append(partial)
+            blocks_inv.append(inv_rms(partial))
+            partial = None
+    srcs = blocks if partial is None else blocks + [partial]
+    inv = blocks_inv if partial is None else blocks_inv + [inv_rms(partial)]
+    h = rms(attnres_mix(srcs, inv, model.attnres_queries[-1]))
+    cache.pos += t
+    cache.max_pos += t
+    return model._logits(h[:, -1])
+
+
+def prefill(model, prompt, extra):
+    """Process one prompt (list/array of ids); returns (cache with room for `extra` tokens, last logits [1, V])."""
+    ids = torch.as_tensor(prompt, dtype=torch.long, device=model.embed.weight.device)[None]
+    cache = Cache(model, 1, ids.size(1) + extra)
+    return cache, step(model, ids, cache)
+
+
+def sample(logits, temperature=0.7, top_p=0.9, top_k=0, recent=None, rep_penalty=1.0, bitmask=None, vocab=None):
+    """Pick one token per row. recent: [B, n] recently generated ids (-1 = none) for the repetition penalty."""
+    logits = logits.float()
+    if vocab is not None and vocab < logits.size(1):
+        logits[:, vocab:] = -float("inf")  # padding rows of the embedding are not real tokens
+    if recent is not None and rep_penalty != 1.0:
+        seen = torch.zeros(logits.size(0), logits.size(1) + 1, device=logits.device, dtype=torch.bool)
+        seen.scatter_(1, torch.where(recent < 0, logits.size(1), recent), True)
+        seen = seen[:, :-1]
+        logits = torch.where(seen, torch.where(logits > 0, logits / rep_penalty, logits * rep_penalty), logits)
+    if bitmask is not None:  # llguidance bitmask: bit i of word w allows token 32*w + i
+        bits = (bitmask.to(logits.device)[:, :, None] >> torch.arange(32, device=logits.device)) & 1
+        logits = logits.masked_fill(bits.reshape(bits.size(0), -1)[:, :logits.size(1)] == 0, -float("inf"))
+    if temperature <= 0:
+        return logits.argmax(-1)
+    logits = logits / temperature
+    if top_k:
+        kth = torch.topk(logits, top_k, dim=-1).values[:, -1:]
+        logits = logits.masked_fill(logits < kth, -float("inf"))
+    probs = torch.softmax(logits, -1)
+    if top_p < 1.0:
+        sp, si = probs.sort(-1, descending=True)
+        sp = sp.masked_fill(sp.cumsum(-1) - sp > top_p, 0.0)
+        probs = torch.zeros_like(probs).scatter_(-1, si, sp)
+    return torch.multinomial(probs, 1)[:, 0]
+
+
+@torch.no_grad()
+def generate(model, prompts, max_new=256, temperature=0.7, top_p=0.9, top_k=0, rep_penalty=1.0, stop_ids=(),
+             matchers=None, batch_size=32, vocab=None, on_token=None, cuda_graph=True):
+    """Sample continuations for many prompts, `batch_size` at a time. Returns a list of token-id lists.
+    matchers: optional per-prompt llguidance matchers (None = unconstrained) that force a grammar.
+    on_token(row_index, token_id): optional callback, e.g. for streaming a single prompt."""
+    from llguidance.torch import allocate_token_bitmask, fill_next_token_bitmask
+    results = [None] * len(prompts)
+    stop = torch.tensor(list(stop_ids) or [-1], device=model.embed.weight.device)
+    for start in range(0, len(prompts), batch_size):
+        idx = list(range(start, min(start + batch_size, len(prompts))))
+        pre = [prefill(model, prompts[i], max_new) for i in idx]
+        cache = Cache.stack([c for c, _ in pre], max_new) if len(pre) > 1 else pre[0][0]
+        logits = torch.cat([l for _, l in pre])
+        del pre
+        B = len(idx)
+        rows_m = [matchers[i] if matchers else None for i in idx]
+        bitmask = allocate_token_bitmask(B, logits.size(1)) if any(rows_m) else None
+        out = [[] for _ in range(B)]
+        done = torch.zeros(B, dtype=torch.bool, device=logits.device)
+        recent = torch.full((B, 64), -1, dtype=torch.long, device=logits.device)
+        graph = _DecodeGraph(model, cache, B) if cuda_graph else None
+        for n in range(max_new):
+            if bitmask is not None:
+                bitmask.fill_(-1)  # all tokens allowed ...
+                for r, m in enumerate(rows_m):
+                    if m is not None and not m.is_stopped():
+                        fill_next_token_bitmask(m, bitmask, r)  # ... except where a grammar forbids them
+            nxt = sample(logits, temperature, top_p, top_k, recent, rep_penalty, bitmask, vocab)
+            nxt = torch.where(done, stop[0].clamp(min=0), nxt)
+            toks = nxt.tolist()
+            for r in range(B):
+                if done[r]:
+                    continue
+                if rows_m[r] is not None:
+                    rows_m[r].consume_token(toks[r])
+                if toks[r] in stop_ids:
+                    done[r] = True
+                    continue
+                out[r].append(toks[r])  # keep it even if it completes the grammar (e.g. the final "}")
+                if on_token:
+                    on_token(idx[r], toks[r])
+                if rows_m[r] is not None and rows_m[r].is_stopped():
+                    done[r] = True
+            recent = torch.cat([recent[:, 1:], nxt[:, None]], 1)
+            if bool(done.all()) or n == max_new - 1:
+                break
+            logits = graph.step(nxt) if graph else step(model, nxt[:, None], cache)
+        for r, i in enumerate(idx):
+            results[i] = out[r]
+    return results
+
+
+class _DecodeGraph:
+    """Records the one-token decode step as a CUDA graph and replays it: one launch instead of ~3,500
+    small kernel launches per token (launch overhead, not the GPU, dominates at this model size).
+    The first steps run normally (they also warm up the Triton kernels); if recording fails, it
+    quietly keeps running the normal way."""
+
+    WARMUP = 2
+
+    def __init__(self, model, cache, batch):
+        self.model, self.cache, self.n, self.graph = model, cache, 0, None
+        self.tok = torch.zeros(batch, 1, dtype=torch.long, device=model.embed.weight.device)
+        self.failed = False
+
+    def step(self, nxt):
+        self.n += 1
+        if self.failed or self.n <= self.WARMUP:
+            return step(self.model, nxt[:, None], self.cache)
+        if self.cache.max_pos + 1 > self.cache.capacity:
+            raise ValueError(f"cache full ({self.cache.capacity} tokens)")
+        self.tok.copy_(nxt[:, None])
+        if self.graph is None:
+            try:
+                self.graph = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(self.graph), torch.autocast("cuda", dtype=torch.bfloat16, cache_enabled=False):
+                    self.out = step(self.model, self.tok, self.cache)  # recorded, not run
+                self.cache.max_pos -= 1  # the recording pass bumped it; the replay below is the real step
+            except Exception as e:  # noqa: BLE001 - fall back to plain steps
+                print(f"(CUDA graph unavailable, using normal decoding: {type(e).__name__}: {str(e)[:120]})")
+                self.failed, self.graph = True, None
+                return step(self.model, nxt[:, None], self.cache)
+        self.graph.replay()
+        self.cache.max_pos += 1
+        return self.out
+
+
+class GrammarFactory:
+    """Builds llguidance matchers for our tokenizer: guaranteed-valid JSON or tool calls."""
+
+    def __init__(self, tokenizer_path, stop_id, n_vocab):
+        """n_vocab: the model's (padded) output size, so masks line up with its logits."""
+        import llguidance.hf
+        from transformers import PreTrainedTokenizerFast
+        from tokenizers import Tokenizer
+        hf = PreTrainedTokenizerFast(tokenizer_object=Tokenizer.from_file(str(tokenizer_path)))
+        self.lltok = llguidance.hf.from_tokenizer(hf, n_vocab=n_vocab, eos_token=stop_id)
+
+    def _matcher(self, grammar):
+        from llguidance import LLMatcher
+        m = LLMatcher(self.lltok, grammar)
+        if m.is_error():
+            raise ValueError(m.get_error())
+        return m
+
+    COMPACT = {"whitespace_flexible": False, "item_separator": ", ", "key_separator": ": "}
+
+    def json(self, schema=None):
+        """Any valid JSON value matching `schema` (a dict; None = any JSON object). Compact whitespace, so a
+        reply can't wander off into blank space and run out of tokens before the JSON is closed."""
+        from llguidance import LLMatcher
+        return self._matcher(LLMatcher.grammar_from_json_schema({**(schema or {"type": "object"}), "x-guidance": self.COMPACT}))
+
+    def tool_calls(self, tools):
+        """One or more <tool_call> blocks whose JSON names a listed function with schema-valid arguments.
+        tools: OpenAI-style [{"type": "function", "function": {"name", "parameters"}}, ...]. Parameter lists in
+        the shorthand some datasets use ({"arg": {"type": "str, optional"}}) are converted to JSON Schema. If a
+        schema still can't be compiled, only the function name is enforced."""
+        from llguidance import LLMatcher
+        fns = [t.get("function", t) for t in tools]
+        try:
+            return self._tool_matcher(fns, strict_args=True)
+        except ValueError:
+            return self._tool_matcher(fns, strict_args=False)
+
+    def _tool_matcher(self, fns, strict_args):
+        import json as _json
+        from llguidance import LLMatcher
+        options = []
+        for fn in fns:
+            params = to_json_schema(fn.get("parameters")) if strict_args else {"type": "object"}
+            options.append({"type": "object", "properties": {"name": {"const": fn["name"]}, "arguments": params},
+                            "required": ["name", "arguments"], "additionalProperties": False})
+        schema = _json.dumps({"anyOf": options, "x-guidance": self.COMPACT})
+        lark = (f'start: call ("\\n" call)*\n'
+                f'call: "<tool_call>\\n" body "\\n</tool_call>"\n'
+                f'body: %json {schema}\n')
+        return self._matcher(LLMatcher.grammar_from_lark(lark))
+
+
+_SHORT_TYPES = {"str": "string", "string": "string", "int": "integer", "integer": "integer", "float": "number",
+                "number": "number", "bool": "boolean", "boolean": "boolean", "list": "array", "array": "array",
+                "dict": "object", "object": "object"}
+
+
+def to_json_schema(params):
+    """Tool parameters as JSON Schema, accepting either real JSON Schema or the {"arg": {"type": "str, optional"}}
+    shorthand. Unknown types become 'any value'. Only declared arguments are allowed."""
+    if not params:
+        return {"type": "object", "properties": {}, "additionalProperties": False}
+    if params.get("type") == "object" or "properties" in params:
+        out = dict(params)
+        out.setdefault("additionalProperties", False)
+        return out
+    props, required = {}, []
+    for name, spec in params.items():
+        spec = spec if isinstance(spec, dict) else {}
+        raw = str(spec.get("type", "")).lower()
+        base = raw.split(",")[0].strip().split("[")[0]
+        props[name] = {"type": _SHORT_TYPES[base]} if base in _SHORT_TYPES else {}
+        if "optional" not in raw:
+            required.append(name)
+    return {"type": "object", "properties": props, "required": required, "additionalProperties": False}

model.py

@@ -0,0 +1,263 @@
+"""Hybrid language model: Gated DeltaNet + gated attention (3:1), Block Attention Residuals.
+
+Architecture follows Qwen3.5 (3 Gated DeltaNet layers per gated full-attention layer,
+sigmoid output gate on attention, QK-norm) with Kimi's Block Attention Residuals replacing
+the plain residual stream. arch="dense" swaps every DeltaNet layer for gated attention,
+which gives a like-for-like Transformer baseline.
+"""
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+
+@dataclass
+class ModelConfig:
+    vocab_size: int = 32768
+    n_layer: int = 24
+    d_model: int = 1024
+    arch: str = "hybrid"      # "hybrid": every `attn_every`-th layer is attention, rest DeltaNet; "dense": all attention
+    attn_every: int = 4
+    n_head: int = 16          # attention query heads
+    n_kv_head: int = 4
+    head_dim: int = 128
+    gdn_heads: int = 16
+    gdn_head_dim: int = 128
+    ffn_dim: int = 3584
+    attnres_blocks: int = 8   # Block AttnRes: the 2*n_layer sublayers are split into this many blocks
+    rope_theta: float = 10000.0
+    max_seq_len: int = 2048
+    softcap: float = 15.0
+    grad_ckpt: bool = False   # recompute sublayers in backward to save VRAM
+
+
+PRESETS = {
+    # smoke tests only
+    "tiny": dict(n_layer=4, d_model=256, n_head=4, n_kv_head=2, head_dim=64, gdn_heads=4, gdn_head_dim=64,
+                 ffn_dim=768, attnres_blocks=4),
+    # ~140M non-embedding: architecture bake-off size
+    "small": dict(n_layer=12, d_model=768, n_head=12, n_kv_head=3, head_dim=128, gdn_heads=12, gdn_head_dim=128,
+                  ffn_dim=2688, attnres_blocks=8),
+    # ~500M non-embedding: Qwen3.5-0.8B layout with a 32k vocab
+    "base": dict(n_layer=24, d_model=1024, n_head=16, n_kv_head=4, head_dim=128, gdn_heads=16, gdn_head_dim=128,
+                 ffn_dim=3584, attnres_blocks=8),
+}
+
+
+def rms(x):
+    return F.rms_norm(x, (x.size(-1),))
+
+
+def inv_rms(x, eps=1e-6):
+    return torch.rsqrt(x.float().pow(2).mean(-1) + eps)
+
+
+class GatedAttention(nn.Module):
+    """Causal GQA attention with QK-norm, RoPE and Qwen's elementwise sigmoid output gate."""
+
+    def __init__(self, c: ModelConfig):
+        super().__init__()
+        self.nh, self.nkv, self.hd = c.n_head, c.n_kv_head, c.head_dim
+        self.q_proj = nn.Linear(c.d_model, 2 * self.nh * self.hd, bias=False)  # query and gate
+        self.kv_proj = nn.Linear(c.d_model, 2 * self.nkv * self.hd, bias=False)
+        self.o_proj = nn.Linear(self.nh * self.hd, c.d_model, bias=False)
+        inv_freq = 1.0 / (c.rope_theta ** (torch.arange(0, self.hd, 2).float() / self.hd))
+        freqs = torch.outer(torch.arange(c.max_seq_len).float(), inv_freq)
+        self.register_buffer("cos", freqs.cos()[None, :, None, :], persistent=False)
+        self.register_buffer("sin", freqs.sin()[None, :, None, :], persistent=False)
+
+    def rope(self, x):
+        T = x.size(1)
+        cos, sin = self.cos[:, :T].to(x.dtype), self.sin[:, :T].to(x.dtype)
+        x1, x2 = x.chunk(2, -1)
+        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], -1)
+
+    def forward(self, x):
+        B, T, _ = x.shape
+        q, gate = self.q_proj(x).split(self.nh * self.hd, -1)
+        q = q.view(B, T, self.nh, self.hd)
+        k, v = self.kv_proj(x).view(B, T, 2, self.nkv, self.hd).unbind(2)
+        q, k = self.rope(rms(q)), self.rope(rms(k))
+        if self.nkv != self.nh:
+            k = k.repeat_interleave(self.nh // self.nkv, dim=2)
+            v = v.repeat_interleave(self.nh // self.nkv, dim=2)
+        y = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True)
+        y = y.transpose(1, 2).reshape(B, T, self.nh * self.hd)
+        return self.o_proj(y * torch.sigmoid(gate))
+
+
+class DeltaNet(nn.Module):
+    """Gated DeltaNet (linear attention with a delta-rule memory) from flash-linear-attention."""
+
+    def __init__(self, c: ModelConfig):
+        super().__init__()
+        from fla.layers import GatedDeltaNet
+        self.gdn = GatedDeltaNet(hidden_size=c.d_model, head_dim=c.gdn_head_dim, num_heads=c.gdn_heads,
+                                 expand_v=1.0, mode="chunk", use_gate=True, use_short_conv=True, conv_size=4)
+
+    def forward(self, x):
+        return self.gdn(x)[0]
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, c: ModelConfig):
+        super().__init__()
+        self.w_gate = nn.Linear(c.d_model, c.ffn_dim, bias=False)
+        self.w_up = nn.Linear(c.d_model, c.ffn_dim, bias=False)
+        self.w_down = nn.Linear(c.ffn_dim, c.d_model, bias=False)
+
+    def forward(self, x):
+        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
+
+
+def attnres_mix(sources, source_inv_rms, query):
+    """Block AttnRes: softmax over sources of <query, RMSNorm(source)>, then a weighted sum of sources."""
+    logits = torch.stack([(s @ query.to(s.dtype)).float() * r for s, r in zip(sources, source_inv_rms)])
+    weights = logits.softmax(0)
+    h = weights[0, ..., None] * sources[0]
+    for w, s in zip(weights[1:], sources[1:]):
+        h = h + w[..., None] * s
+    return h
+
+
+class LM(nn.Module):
+    def __init__(self, c: ModelConfig):
+        super().__init__()
+        self.config = c
+        self.embed = nn.Embedding(c.vocab_size, c.d_model)
+        self.sublayers = nn.ModuleList()
+        for i in range(c.n_layer):
+            is_attn = c.arch == "dense" or (i % c.attn_every == c.attn_every - 1)
+            self.sublayers.append(GatedAttention(c) if is_attn else DeltaNet(c))
+            self.sublayers.append(SwiGLU(c))
+        n_sub = len(self.sublayers)
+        self.block_size = math.ceil(n_sub / c.attnres_blocks)
+        # one pseudo-query per sublayer plus one for the final read-out; zero init = uniform average at start
+        self.attnres_queries = nn.Parameter(torch.zeros(n_sub + 1, c.d_model))
+        self.lm_head = nn.Linear(c.d_model, c.vocab_size, bias=False)
+        self.init_weights()
+
+    @torch.no_grad()
+    def init_weights(self):
+        c = self.config
+        nn.init.normal_(self.embed.weight, std=1.0)
+        nn.init.normal_(self.lm_head.weight, std=0.001)
+        s = 3 ** 0.5 * c.d_model ** -0.5
+        for m in self.sublayers:
+            if isinstance(m, GatedAttention):
+                nn.init.uniform_(m.q_proj.weight, -s, s)
+                nn.init.uniform_(m.kv_proj.weight, -s, s)
+                nn.init.zeros_(m.o_proj.weight)
+            elif isinstance(m, SwiGLU):
+                nn.init.uniform_(m.w_gate.weight, -s, s)
+                nn.init.uniform_(m.w_up.weight, -s, s)
+                nn.init.zeros_(m.w_down.weight)
+            else:  # DeltaNet keeps FLA's init for its gates/decay; match the rest to the attention layers
+                g = m.gdn
+                for lin in (g.q_proj, g.k_proj, g.v_proj, g.g_proj):
+                    nn.init.uniform_(lin.weight, -s, s)
+                nn.init.zeros_(g.o_proj.weight)
+
+    def num_params(self, non_embedding=True):
+        n = sum(p.numel() for p in self.parameters())
+        return n - self.embed.weight.numel() - self.lm_head.weight.numel() if non_embedding else n
+
+    def _run(self, sub, x):
+        if self.config.grad_ckpt and self.training:
+            return checkpoint(sub, x, use_reentrant=False)
+        return sub(x)
+
+    def hidden(self, idx):
+        """Final normalized hidden states [B, T, d_model]."""
+        x0 = rms(self.embed(idx))
+        blocks, blocks_inv = [x0], [inv_rms(x0)]   # completed block sums (block 0 = token embedding)
+        partial = None                             # running sum of sublayer outputs in the current block
+        for j, sub in enumerate(self.sublayers):
+            srcs = blocks if partial is None else blocks + [partial]
+            inv = blocks_inv if partial is None else blocks_inv + [inv_rms(partial)]
+            h = attnres_mix(srcs, inv, self.attnres_queries[j])
+            out = self._run(sub, rms(h)).float()
+            partial = out if partial is None else partial + out
+            if (j + 1) % self.block_size == 0:
+                blocks.append(partial)
+                blocks_inv.append(inv_rms(partial))
+                partial = None
+        srcs = blocks if partial is None else blocks + [partial]
+        inv = blocks_inv if partial is None else blocks_inv + [inv_rms(partial)]
+        return rms(attnres_mix(srcs, inv, self.attnres_queries[-1]))
+
+    def forward(self, idx, targets=None):
+        h = self.hidden(idx)
+        if targets is None:
+            return self._logits(h)
+        # Loss in chunks, recomputing each chunk's logits in backward, so the full
+        # (tokens x vocab) fp32 logit tensor never has to sit in VRAM.
+        # Targets of -1 (prompts, padding during fine-tuning) are ignored.
+        h, targets = h.flatten(0, 1), targets.flatten()
+        loss = 0.0
+        for i in range(0, h.size(0), self.LOSS_CHUNK):
+            loss = loss + checkpoint(self._chunk_loss, h[i:i + self.LOSS_CHUNK], targets[i:i + self.LOSS_CHUNK],
+                                     use_reentrant=False)
+        return loss / (targets >= 0).sum().clamp(min=1)
+
+    LOSS_CHUNK = 4096
+
+    def _logits(self, h):
+        logits = self.lm_head(h).float()
+        if self.config.softcap:
+            logits = self.config.softcap * torch.tanh(logits / self.config.softcap)
+        return logits
+
+    def _chunk_loss(self, h, targets):
+        return F.cross_entropy(self._logits(h), targets, reduction="sum", ignore_index=-1)
+
+    def token_logprobs(self, idx, targets):
+        """Log-probability of each target token, [B, T] (0 where the target is -1). Chunked like the loss."""
+        h, t = self.hidden(idx).flatten(0, 1), targets.flatten()
+        parts = [checkpoint(self._chunk_logprobs, h[i:i + self.LOSS_CHUNK], t[i:i + self.LOSS_CHUNK], use_reentrant=False)
+                 for i in range(0, h.size(0), self.LOSS_CHUNK)]
+        return torch.cat(parts).view(targets.shape)
+
+    def _chunk_logprobs(self, h, targets):
+        lp = torch.log_softmax(self._logits(h), -1).gather(1, targets.clamp(min=0)[:, None])[:, 0]
+        return lp * (targets >= 0)
+
+    def param_groups(self):
+        """Split parameters: hidden matrices go to Muon, everything else to AdamW."""
+        muon, other = [], []
+        for name, p in self.sublayers.named_parameters():
+            (muon if p.ndim == 2 and min(p.shape) >= 64 else other).append(p)
+        other.append(self.attnres_queries)
+        return dict(muon=muon, embed=[self.embed.weight], head=[self.lm_head.weight], other=other)
+
+
+def resize_vocab(state_dict, new_rows):
+    """Grow the embedding and output matrices for added tokens; new rows start at the mean of the old ones."""
+    for key in ("embed.weight", "lm_head.weight"):
+        w = state_dict[key]
+        if w.size(0) < new_rows:
+            extra = w.float().mean(0, keepdim=True).expand(new_rows - w.size(0), -1).to(w.dtype)
+            state_dict[key] = torch.cat([w, extra])
+    return state_dict
+
+
+@torch.no_grad()
+def generate(model, idx, max_new_tokens, temperature=0.8, top_k=50):
+    """Simple sampling that re-runs the full context each step (no cache; fine for short samples)."""
+    for _ in range(max_new_tokens):
+        ctx = idx[:, -model.config.max_seq_len:]
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            logits = model(ctx)[:, -1, :].float()
+        if temperature <= 0:
+            nxt = logits.argmax(-1, keepdim=True)
+        else:
+            logits = logits / temperature
+            if top_k:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("inf")
+            nxt = torch.multinomial(F.softmax(logits, -1), 1)
+        idx = torch.cat([idx, nxt], 1)
+    return idx

model.safetensors

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