Upload 4 files

README.md

@@ -1,3 +1,162 @@
----
-license: cc-by-nc-4.0
----
+---
+license: cc-by-nc-4.0
+language:
+  - en
+pipeline_tag: text-generation
+tags:
+  - causal-lm
+  - dense
+  - recurrent
+  - custom-architecture
+  - preview
+---
+
+# Sabaki Preview
+
+Sabaki is an experimental decoder-only language model implemented in PyTorch. It combines recurrent Kimi Delta Attention (KDA), periodic gated rotary attention, and dense SiTU-GLU feed-forward layers.
+
+This post-trained model is a research preview for local experiments, evaluation, and further post-training. Instruction following and factual accuracy can be uneven.
+
+## At a glance
+
+| Property | Value |
+| --- | --- |
+| Model parameters, excluding embeddings | **124,128,144 (124.1M)** |
+| Active parameters per token, excluding embeddings | **124,128,144 (124.1M)**; dense feed-forward layers throughout |
+| Tied embedding matrix, counted separately | 37,748,736 parameters; shared by token input and vocabulary output |
+| Layers | 16: 12 recurrent KDA layers and 4 global-attention layers |
+| Hidden size | 768 |
+| Heads | 12 heads of width 64 |
+| Feed-forward layers | Dense SiTU-GLU, intermediate width 2,048 in every block |
+| Vocabulary | 49,152-token lossless byte-level BPE |
+| Configured context | 16,384 tokens |
+| Implementation | Custom PyTorch model with FP32 weights; BF16 training compute |
+| License | CC BY-NC 4.0 |
+
+The 16,384-token context is a configuration limit, not a measured long-context capability. The pretraining continuation used 2,048-token sequences, and the post-training checkpoints used 4,096-token sequences. Reliable behavior beyond those training lengths has not been established.
+
+## Model design
+
+The model repeats three KDA blocks followed by one gated rotary-attention block:
+
+| Blocks (one-based) | Sequence mixer | Feed-forward path |
+| --- | --- | --- |
+| 1-3, 5-7, 9-11, 13-15 | Kimi Delta Attention | Dense SiTU-GLU, width 2,048 |
+| 4, 8, 12, 16 | Gated causal multi-head attention with RoPE | Dense SiTU-GLU, width 2,048 |
+
+Each block applies RMSNorm before its sequence mixer and before its feed-forward layer, with an ordinary residual addition after each. A final RMSNorm precedes the tied vocabulary projection. RMSNorm epsilon is `1e-6`.
+
+### Recurrent KDA layers
+
+KDA projects queries, keys, and values at the full 768-channel width and applies separate causal depthwise convolutions with kernel size 4. Queries and keys are normalized before the delta-rule recurrence. A rank-64 projection produces input-dependent per-channel decay, and a sigmoid write gate controls state updates. Each head maintains a 64 x 64 recurrent matrix, plus short convolution histories. The recurrent output is normalized, multiplied by a learned sigmoid output gate, and projected back to the residual stream.
+
+These layers carry fixed-size state as the sequence grows. The four global-attention layers still retain a conventional KV cache, so the complete model does not have a constant-size cache.
+
+### Gated rotary attention
+
+The global layers use full-width Q, K, and V projections with 12 query/key/value heads of width 64. RoPE is applied to Q and K with base `10,000`, followed by causal scaled dot-product attention. A learned sigmoid gate modulates the merged heads before the output projection. These are standard full-head attention layers, not latent-compressed MLA.
+
+### Dense SiTU-GLU feed-forward layers
+
+Every block projects the normalized hidden state into two 2,048-wide branches, applies the bounded SiTU gate and up activations, multiplies them, and projects back to width 768. With gate preactivation `g` and up preactivation `u`, the activation is:
+
+```text
+(4 * tanh(g / 4) * sigmoid(g)) * (25 * tanh(u / 25))
+```
+
+There are no routed experts, shared experts, routers, or expert-balancing biases. The dense checkpoints are incompatible with the older MoE architecture.
+
+| Component | Parameters, excluding embeddings |
+| --- | ---: |
+| 12 KDA mixers, including their internal normalization | 36,808,848 |
+| 4 gated rotary-attention mixers | 11,796,480 |
+| 16 dense feed-forward layers | 75,497,472 |
+| Block and final RMSNorm scales | 25,344 |
+| **Total** | **124,128,144** |
+
+## Training
+
+| Stage | Training tokens |
+| --- | --- |
+| Pretraining | 16B tokens |
+| Supervised fine-tuning (SFT) | About 200M assistant-target tokens |
+| Further grounded-chat SFT | About 165K assistant-target tokens |
+
+The completed checkpoints contain FP32 model tensors and no optimizer state. Training used BF16 compute.
+
+## Running with PyTorch
+
+Use Sabaki's custom PyTorch implementation with the matching tokenizer and trained weights. Generation runs the model in evaluation mode with gradients disabled, processes the prompt once, and then uses cached recurrent states and attention keys/values to generate subsequent tokens.
+
+The prompt serializer uses dedicated control-token IDs. A single-turn prompt has this structure:
+
+```text
+<|bos|><|user|>Explain why the sky appears blue.<|user_end|><|assistant|>
+```
+
+Assistant answers use `<|final|>...<|final_end|><|assistant_end|>`, with `<|eos|>` closing a training episode. Use the tokenizer's control-token serialization to construct this prompt. Plain `USER:` / `ASSISTANT:` labels are not the trained protocol.
+
+The PyTorch generation defaults are temperature `0.2`, top-p `0.95`, and a maximum of `128` new tokens. Generation stops at `<|final_end|>` or `<|eos|>`, or when the token budget is exhausted. Keep the prompt and generated continuation within the configured context length.
+
+Always set a generation limit.
+
+## Evaluation
+
+These scores describe **Sabaki after the 200M-token SFT stage**. The later grounded chat stage has a separate evaluation below.
+
+| Benchmark | Questions | Zero-shot normalized accuracy |
+| --- | ---: | ---: |
+| HellaSwag | 10,042 | 36.85% |
+| ARC-Easy | 2,376 | 47.05% |
+| ARC-Challenge | 1,172 | 28.75% |
+| PIQA | 1,838 | 65.83% |
+| ArithMark-3 | 1,000 | 36.00% |
+
+These are local evaluation results, not a verified leaderboard submission.
+
+The later grounded-chat checkpoint has a separate assistant-token evaluation: held-out NLL 2.4488 and perplexity 11.5746 across 50 conversations and 3,703 assistant tokens. This does not replace a five-benchmark evaluation of those weights.
+
+### Intelligence Index comparison
+
+Sabaki scores **20.33** on the Open SLM Leaderboard's Intelligence Index (INT), using the full benchmark evaluation of the 200M token SFT stage. The index adjusts each benchmark for chance performance, averages ARC-Easy and ARC-Challenge into a single ARC component, and weights HellaSwag, ARC, and PIQA at 1.0 each and ArithMark-3 at 0.65.
+
+The following comparison uses the [Open SLM Leaderboard](https://huggingface.co/spaces/AxiomicLabs/Open_SLM_Leaderboard) scores and formula checked on 12 October 2026. Higher is better.
+
+| Model | INT score |
+| --- | ---: |
+| GPT-X3-150M | 33.05 |
+| SmolLM2-135M | 27.13 |
+| MobileLLM-R1-140M-base | 24.64 |
+| GPT-X2-125M | 23.36 |
+| JugnuLM-110M-R2+ | 20.39 |
+| **Sabaki** | **20.33** |
+| Speck2-140M | 20.02 |
+| GPT-X-125M | 19.94 |
+| GTM-v2-chat | 15.36 |
+| OPT-125M | 13.80 |
+| GPT-2 | 13.58 |
+
+Sabaki is numerically close to JugnuLM-110M-R2+, slightly ahead of Speck2-140M and GPT-X-125M, and above GPT-2 and OPT-125M. It trails SmolLM2-135M by 6.80 INT points. Small differences should not be interpreted as established performance advantages.
+
+By score alone, Sabaki would sit 21st among the 210 existing entries with all five benchmark scores checked on that date. This is an unofficial comparison using local PyTorch results; it does not establish leaderboard eligibility or a verified placement. The later grounded-chat checkpoint has no recorded full INT evaluation.
+
+## Limitations
+
+This preview can produce incorrect facts, invented citations, repetitive text, and unreliable instructions. Languages other than English have not been evaluated separately. Protocol tokens for tools and other structured outputs do not establish reliable tool-use capabilities. The model has not received comprehensive safety evaluation or preference optimization.
+
+## License
+
+The model weights are available under the **Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0)** license. Attribution is required, and commercial use is not permitted. Software retains its own license.
+
+## Citation
+
+If this preview helps your research, please cite the model repository. A formal citation can be added with a full release.
+
+```bibtex
+@software{sabaki_preview_2026,
+  author = {{Local-Axiom-AI}},
+  title  = {Sabaki Preview},
+  year   = {2026},
+  note   = {A compact hybrid recurrent and gated rotary-attention language model with dense SiTU feed-forward layers}
+}
+```

Sabaki-Dense.pt

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

infer.py

@@ -0,0 +1,670 @@
+"""Standalone Sabaki Dense inference: Python 3.11+ and PyTorch 2.8+.
+
+Keep this script, a regular Sabaki Dense checkpoint, and tokenizer.json together.
+No repository package, Transformers, Triton, or Flash Linear Attention required.
+
+Examples:
+    python infer.py --checkpoint chat-grounded-500.pt --prompt "Why is the sky blue?"
+    python infer.py --checkpoint unc-chat.pt --prompt "Hello!" --device cpu
+    python infer.py --checkpoint dense-pretrain-sixday.pt --raw --prompt "The sky is"
+
+This reference implementation favors portability over optimized prefill speed.
+Checkpoint format identifiers retain their original spelling for compatibility.
+"""
+from __future__ import annotations
+
+import argparse
+from dataclasses import asdict, dataclass
+import heapq
+import json
+import math
+from pathlib import Path
+import sys
+from typing import Any, Iterable, Sequence
+
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+@dataclass
+class SabakiConfig:
+    """Architecture configuration.
+
+    ``canonical()`` uses dense FFNs and head-width rotary attention. ``tiny()`` keeps
+    the same block types and dataflow but is intended for tests and examples.
+    """
+
+    architecture_version: str = "kurokiba-dense-v1"
+    vocab_size: int = 49_152
+    d_model: int = 768
+    num_layers: int = 16
+    num_heads: int = 12
+    head_dim: int = 64
+    context_length: int = 16_384
+    conv_kernel_size: int = 4
+    kda_chunk_size: int = 64
+    decay_rank: int = 64
+    decay_min_log: float = -5.0
+    rope_theta: float = 10_000.0
+    dense_hidden_dim: int = 2_048
+    rms_norm_eps: float = 1e-6
+    situ_gate_cap: float = 4.0
+    situ_up_cap: float = 25.0
+    initializer_std: float = 0.02
+    init_seed: int = 20_260_729
+
+    def __post_init__(self) -> None:
+        if self.architecture_version != "kurokiba-dense-v1":
+            raise ValueError("unsupported architecture; initialize a new kurokiba-dense-v1 model")
+        if self.d_model != self.num_heads * self.head_dim:
+            raise ValueError("d_model must equal num_heads * head_dim")
+        if min(self.vocab_size, self.d_model, self.num_heads, self.decay_rank, self.conv_kernel_size) < 1:
+            raise ValueError("model, vocabulary, convolution, and decay dimensions must be positive")
+        if self.num_layers < 1:
+            raise ValueError("num_layers must be positive")
+        if self.head_dim < 2 or self.head_dim % 2:
+            raise ValueError("head_dim must be positive and even for RoPE")
+        if not math.isfinite(self.rope_theta) or self.rope_theta <= 1.0:
+            raise ValueError("rope_theta must be finite and greater than one")
+        if self.dense_hidden_dim < 1:
+            raise ValueError("dense_hidden_dim must be positive")
+        if self.context_length < 1:
+            raise ValueError("context_length must be positive")
+        if self.kda_chunk_size < 1:
+            raise ValueError("kda_chunk_size must be positive")
+
+    @classmethod
+    def canonical(cls) -> "SabakiConfig":
+        return cls()
+
+    @classmethod
+    def tiny(cls) -> "SabakiConfig":
+        return cls(
+            vocab_size=512,
+            d_model=64,
+            num_layers=4,
+            num_heads=4,
+            head_dim=16,
+            context_length=256,
+            conv_kernel_size=4,
+            kda_chunk_size=8,
+            decay_rank=16,
+            dense_hidden_dim=128,
+        )
+
+    @property
+    def attention_types(self) -> tuple[str, ...]:
+        # Every fourth layer is global rotary attention; the last layer is always global.
+        return tuple(
+            "attention" if (index + 1) % 4 == 0 or index == self.num_layers - 1 else "kda"
+            for index in range(self.num_layers)
+        )
+
+    @property
+    def residual_std(self) -> float:
+        return self.initializer_std / (2 * self.num_layers) ** 0.5
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, values: dict[str, Any]) -> "SabakiConfig":
+        if values.get("architecture_version") != "kurokiba-dense-v1" or "attnres_group_size" in values:
+            raise ValueError("MoE/legacy config is incompatible with dense-v1; initialize new dense weights")
+        return cls(**values)
+
+    def save(self, path: str | Path) -> None:
+        Path(path).write_text(
+            json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+
+    @classmethod
+    def load(cls, path: str | Path) -> "SabakiConfig":
+        return cls.from_dict(json.loads(Path(path).read_text(encoding="utf-8")))
+
+    def parameter_count_formula(self) -> dict[str, int]:
+        """Return structural parameter counts without allocating the model."""
+        d = self.d_model
+        h = self.num_heads
+        hd = self.head_dim
+
+        kda_per_layer = (
+            3 * d * d
+            + 3 * d * self.conv_kernel_size
+            + self.decay_rank * d
+            + d * self.decay_rank
+            + h * hd
+            + h
+            + h * d
+            + h * hd
+            + 2 * d * d
+        )
+        attention_per_layer = 5 * d * d  # fused QKV, output, output gate
+        dense = self.num_layers * 3 * d * self.dense_hidden_dim
+        kda_layers = self.attention_types.count("kda")
+        attention_layers = self.attention_types.count("attention")
+        norms = (2 * self.num_layers + 1) * d
+        embedding = self.vocab_size * d
+        non_embedding = kda_layers * kda_per_layer + attention_layers * attention_per_layer + dense + norms
+        return {
+            "embedding": embedding,
+            "kda": kda_layers * kda_per_layer,
+            "attention": attention_layers * attention_per_layer,
+            "dense_ffn": dense,
+            "module_and_final_norms": norms,
+            "trainable": embedding + non_embedding,
+            "persistent": embedding + non_embedding,
+            "active_non_embedding": non_embedding,
+            "active_with_embedding": embedding + non_embedding,
+        }
+
+
+CONTROL_TOKENS: tuple[str, ...] = (
+    "<|pad|>",
+    "<|bos|>",
+    "<|eos|>",
+    "<|system|>",
+    "<|system_end|>",
+    "<|developer|>",
+    "<|developer_end|>",
+    "<|user|>",
+    "<|user_end|>",
+    "<|assistant|>",
+    "<|assistant_end|>",
+    "<|tool_schema|>",
+    "<|tool_schema_end|>",
+    "<|tool_call|>",
+    "<|tool_call_end|>",
+    "<|parallel_tool_calls|>",
+    "<|parallel_tool_calls_end|>",
+    "<|tool_result|>",
+    "<|tool_result_end|>",
+    "<|tool_error|>",
+    "<|tool_error_end|>",
+    "<|final|>",
+    "<|final_end|>",
+    "<|memory|>",
+    "<|memory_end|>",
+    "<|confirmation_required|>",
+    "<|confirmation_required_end|>",
+    "<|confirmation_granted|>",
+    "<|confirmation_denied|>",
+    "<|citation|>",
+    "<|citation_end|>",
+    "<|observation|>",
+    "<|observation_end|>",
+    "<|sep|>",
+    "<|json|>",
+    "<|json_end|>",
+    "<|cancelled|>",
+    "<|truncated|>",
+    "<|unknown|>",
+    "<|tool_denied|>",
+    "<|tool_denied_end|>",
+)
+
+RESERVED_TOKEN_COUNT = 256
+BYTE_TOKEN_OFFSET = 256
+BASE_VOCAB_SIZE = RESERVED_TOKEN_COUNT + 256
+CANONICAL_VOCAB_SIZE = 49_152
+CANONICAL_MERGE_COUNT = CANONICAL_VOCAB_SIZE - BASE_VOCAB_SIZE
+
+
+class ByteBPETokenizer:
+    """Lossless UTF-8 byte BPE with an isolated reserved-token namespace.
+
+    Ordinary ``encode`` never recognizes control-token-looking text. Trusted
+    callers must insert controls explicitly with ``control_id`` or
+    ``encode_parts``.
+    """
+
+    def __init__(
+        self,
+        merges: Sequence[tuple[int, int]] = (),
+        *,
+        vocab_size: int = CANONICAL_VOCAB_SIZE,
+    ) -> None:
+        if vocab_size < BASE_VOCAB_SIZE + len(merges):
+            raise ValueError("vocab_size is too small for the supplied merges")
+        self.vocab_size = vocab_size
+        self.merges = tuple((int(left), int(right)) for left, right in merges)
+        self._merge_rank: dict[tuple[int, int], int] = {}
+        self._merge_id: dict[tuple[int, int], int] = {}
+        self._symbols: dict[int, bytes] = {
+            BYTE_TOKEN_OFFSET + value: bytes((value,)) for value in range(256)
+        }
+        for rank, pair in enumerate(self.merges):
+            token_id = BASE_VOCAB_SIZE + rank
+            if pair in self._merge_rank:
+                raise ValueError(f"duplicate merge pair at rank {rank}: {pair}")
+            if pair[0] not in self._symbols or pair[1] not in self._symbols:
+                raise ValueError(f"merge {pair} references an unknown token")
+            self._merge_rank[pair] = rank
+            self._merge_id[pair] = token_id
+            self._symbols[token_id] = self._symbols[pair[0]] + self._symbols[pair[1]]
+        self._control_to_id = {
+            token: index for index, token in enumerate(CONTROL_TOKENS)
+        }
+
+    @property
+    def trained_vocab_size(self) -> int:
+        return BASE_VOCAB_SIZE + len(self.merges)
+
+    def control_id(self, token: str) -> int:
+        try:
+            return self._control_to_id[token]
+        except KeyError as exc:
+            raise KeyError(f"unknown control token: {token}") from exc
+
+    def encode_bytes(self, value: bytes) -> list[int]:
+        tokens = [BYTE_TOKEN_OFFSET + byte for byte in value]
+        if len(tokens) < 2 or not self._merge_rank:
+            return tokens
+
+        length = len(tokens)
+        previous = [index - 1 for index in range(length)]
+        following = [index + 1 for index in range(length)]
+        following[-1] = -1
+        alive = [True] * length
+        candidates: list[tuple[int, int, int, int]] = []
+
+        def add_candidate(left_index: int) -> None:
+            if left_index < 0 or not alive[left_index]:
+                return
+            right_index = following[left_index]
+            if right_index < 0 or not alive[right_index]:
+                return
+            pair = (tokens[left_index], tokens[right_index])
+            rank = self._merge_rank.get(pair)
+            if rank is not None:
+                heapq.heappush(
+                    candidates,
+                    (rank, left_index, pair[0], pair[1]),
+                )
+
+        for index in range(length - 1):
+            add_candidate(index)
+
+        while candidates:
+            rank, left_index, expected_left, expected_right = heapq.heappop(
+                candidates
+            )
+            if not alive[left_index] or tokens[left_index] != expected_left:
+                continue
+            right_index = following[left_index]
+            if (
+                right_index < 0
+                or not alive[right_index]
+                or tokens[right_index] != expected_right
+            ):
+                continue
+
+            tokens[left_index] = BASE_VOCAB_SIZE + rank
+            alive[right_index] = False
+            next_index = following[right_index]
+            following[left_index] = next_index
+            if next_index >= 0:
+                previous[next_index] = left_index
+            add_candidate(previous[left_index])
+            add_candidate(left_index)
+
+        encoded: list[int] = []
+        index = 0
+        while index >= 0:
+            if alive[index]:
+                encoded.append(tokens[index])
+            index = following[index]
+        return encoded
+
+    def encode(self, text: str) -> list[int]:
+        return self.encode_bytes(text.encode("utf-8"))
+
+    def encode_parts(self, parts: Iterable[str | tuple[str, str]]) -> list[int]:
+        """Encode trusted controls and ordinary text without token injection.
+
+        A plain string is always ordinary text. A ``("control", token)`` tuple
+        explicitly inserts one reserved token.
+        """
+        token_ids: list[int] = []
+        for part in parts:
+            if isinstance(part, str):
+                token_ids.extend(self.encode(part))
+            elif len(part) == 2 and part[0] == "control":
+                token_ids.append(self.control_id(part[1]))
+            else:
+                raise ValueError(f"invalid tokenizer part: {part!r}")
+        return token_ids
+
+    def decode_bytes(self, token_ids: Iterable[int], *, skip_control: bool = False) -> bytes:
+        pieces: list[bytes] = []
+        for token_id in token_ids:
+            token_id = int(token_id)
+            if 0 <= token_id < RESERVED_TOKEN_COUNT:
+                if skip_control:
+                    continue
+                if token_id >= len(CONTROL_TOKENS):
+                    raise ValueError(f"reserved token {token_id} has no display string")
+                pieces.append(CONTROL_TOKENS[token_id].encode("ascii"))
+            else:
+                try:
+                    pieces.append(self._symbols[token_id])
+                except KeyError as exc:
+                    raise ValueError(
+                        f"token {token_id} is not defined by this tokenizer"
+                    ) from exc
+        return b"".join(pieces)
+
+    def decode(
+        self,
+        token_ids: Iterable[int],
+        *,
+        errors: str = "strict",
+        skip_control: bool = False,
+    ) -> str:
+        return self.decode_bytes(token_ids, skip_control=skip_control).decode(
+            "utf-8", errors=errors
+        )
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "version": 1,
+            "type": "lossless_byte_bpe",
+            "vocab_size": self.vocab_size,
+            "reserved_token_count": RESERVED_TOKEN_COUNT,
+            "control_tokens": list(CONTROL_TOKENS),
+            "merges": [list(pair) for pair in self.merges],
+        }
+
+    def save(self, path: str | Path) -> None:
+        Path(path).write_text(
+            json.dumps(self.to_dict(), ensure_ascii=False, indent=2) + "\n",
+            encoding="utf-8",
+        )
+
+    @classmethod
+    def load(cls, path: str | Path) -> "ByteBPETokenizer":
+        payload = json.loads(Path(path).read_text(encoding="utf-8"))
+        if payload.get("reserved_token_count") != RESERVED_TOKEN_COUNT:
+            raise ValueError("tokenizer uses an incompatible reserved-token range")
+        if tuple(payload.get("control_tokens", ())) != CONTROL_TOKENS:
+            raise ValueError("tokenizer control-token map is incompatible")
+        return cls(
+            [tuple(pair) for pair in payload["merges"]],
+            vocab_size=int(payload["vocab_size"]),
+        )
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, shape, eps):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(shape))
+        self.eps = eps
+
+    def forward(self, x):
+        return F.rms_norm(x, self.weight.shape, self.weight, self.eps)
+
+
+class CausalDepthwiseConv(nn.Module):
+    def __init__(self, width, kernel):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(width, kernel))
+
+    def forward(self, x, history):
+        history_length = history.shape[-1]
+        extended = torch.cat((history.flip(-1).transpose(1, 2), x), dim=1)
+        result = x * self.weight[:, 0]
+        for tap in range(1, self.weight.shape[1]):
+            start = history_length - tap
+            result = result + extended[:, start:start + x.shape[1]] * self.weight[:, tap]
+        newest = torch.cat((x.flip(1).transpose(1, 2), history), dim=-1)
+        return result, newest[..., :history_length]
+
+
+class KDA(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.config = cfg
+        d, h, hd = cfg.d_model, cfg.num_heads, cfg.head_dim
+        self.q_proj = nn.Linear(d, d, bias=False)
+        self.k_proj = nn.Linear(d, d, bias=False)
+        self.v_proj = nn.Linear(d, d, bias=False)
+        self.q_conv = CausalDepthwiseConv(d, cfg.conv_kernel_size)
+        self.k_conv = CausalDepthwiseConv(d, cfg.conv_kernel_size)
+        self.v_conv = CausalDepthwiseConv(d, cfg.conv_kernel_size)
+        self.alpha_down = nn.Linear(d, cfg.decay_rank, bias=False)
+        self.alpha_up = nn.Linear(cfg.decay_rank, d, bias=False)
+        self.alpha_bias = nn.Parameter(torch.empty(h, hd))
+        self.alpha_log_scale = nn.Parameter(torch.empty(h))
+        self.beta_proj = nn.Linear(d, h, bias=False)
+        self.head_norm = RMSNorm((h, hd), cfg.rms_norm_eps)
+        self.output_gate = nn.Linear(d, d, bias=False)
+        self.out_proj = nn.Linear(d, d, bias=False)
+
+    def forward(self, x, state=None):
+        cfg = self.config
+        batch, length, _ = x.shape
+        if state is None:
+            recurrent = torch.zeros(batch, cfg.num_heads, cfg.head_dim, cfg.head_dim,
+                                    dtype=torch.float32, device=x.device)
+            histories = [x.new_zeros(batch, cfg.d_model, cfg.conv_kernel_size - 1) for _ in range(3)]
+        else:
+            recurrent = state[0].float()
+            histories = [item.to(dtype=x.dtype) for item in state[1:]]
+        mixed, new_histories = [], []
+        for projection, convolution, history in zip(
+            (self.q_proj, self.k_proj, self.v_proj),
+            (self.q_conv, self.k_conv, self.v_conv), histories
+        ):
+            value, history = convolution(projection(x), history)
+            mixed.append(F.silu(value).view(batch, length, cfg.num_heads, cfg.head_dim).float())
+            new_histories.append(history)
+        q, k, v = mixed
+        q = q / torch.sqrt(q.square().sum(-1, keepdim=True) + cfg.rms_norm_eps)
+        k = k / torch.sqrt(k.square().sum(-1, keepdim=True) + cfg.rms_norm_eps)
+        decay = self.alpha_up(self.alpha_down(x)) + self.alpha_bias.flatten()
+        z = decay.view_as(q).float()
+        scale = self.alpha_log_scale.float().exp().view(1, 1, -1, 1)
+        alpha = (cfg.decay_min_log * torch.sigmoid(scale * z)).exp()
+        beta = self.beta_proj(x).sigmoid().float()
+        outputs = []
+        for index in range(length):
+            decayed = alpha[:, index].unsqueeze(-1) * recurrent
+            prediction = torch.matmul(decayed.transpose(-2, -1), k[:, index].unsqueeze(-1)).squeeze(-1)
+            error = v[:, index] - prediction
+            recurrent = decayed + beta[:, index, :, None, None] * (k[:, index].unsqueeze(-1) * error.unsqueeze(-2))
+            outputs.append(torch.matmul(recurrent.transpose(-2, -1), q[:, index].unsqueeze(-1)).squeeze(-1))
+        output = self.head_norm(torch.stack(outputs, dim=1).to(x.dtype)).flatten(2)
+        result = self.out_proj(self.output_gate(x).sigmoid() * output)
+        return result, (recurrent.to(x.dtype), *new_histories)
+
+
+class RotaryAttention(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.config = cfg
+        self.qkv_proj = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
+        self.out_proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        self.output_gate = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        self._rope_tables = None
+
+    def rotary(self, x, offset):
+        cfg = self.config
+        if self._rope_tables is None or self._rope_tables[0].device != x.device:
+            frequencies = cfg.rope_theta ** (-torch.arange(0, cfg.head_dim, 2, device=x.device, dtype=torch.float32) / cfg.head_dim)
+            angles = torch.arange(cfg.context_length, device=x.device, dtype=torch.float32)[:, None] * frequencies[None, :]
+            self._rope_tables = angles.cos(), angles.sin()
+        cos, sin = (table[offset:offset + x.shape[2]][None, None] for table in self._rope_tables)
+        even, odd = x.float()[..., 0::2], x.float()[..., 1::2]
+        return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2).to(x.dtype)
+
+    def forward(self, x, state=None):
+        cfg = self.config
+        batch, length, _ = x.shape
+        q, k, v = self.qkv_proj(x).view(batch, length, 3, cfg.num_heads, cfg.head_dim).unbind(2)
+        q, k, v = (item.transpose(1, 2) for item in (q, k, v))
+        offset = 0 if state is None else state[0].shape[2]
+        q, k = self.rotary(q, offset), self.rotary(k, offset)
+        if state is not None:
+            k, v = torch.cat((state[0], k), dim=2), torch.cat((state[1], v), dim=2)
+        if offset and length > 1:
+            output = torch.cat([
+                F.scaled_dot_product_attention(q[:, :, i:i+1], k[:, :, :offset+i+1], v[:, :, :offset+i+1], dropout_p=0.0)
+                for i in range(length)
+            ], dim=2)
+        else:
+            output = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=offset == 0)
+        output = output.transpose(1, 2).reshape(batch, length, cfg.d_model)
+        return self.out_proj(self.output_gate(x).sigmoid() * output), (k, v)
+
+
+class SiTUGLU(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.gate_up_proj = nn.Linear(cfg.d_model, 2 * cfg.dense_hidden_dim, bias=False)
+        self.out_proj = nn.Linear(cfg.dense_hidden_dim, cfg.d_model, bias=False)
+        self.gate_cap, self.up_cap = cfg.situ_gate_cap, cfg.situ_up_cap
+
+    def forward(self, x):
+        gate_pre, up_pre = self.gate_up_proj(x).chunk(2, dim=-1)
+        gate = self.gate_cap * torch.tanh(gate_pre / self.gate_cap)
+        gate = gate * torch.sigmoid(gate_pre)
+        up = self.up_cap * torch.tanh(up_pre / self.up_cap)
+        return self.out_proj(gate * up)
+
+
+class Block(nn.Module):
+    def __init__(self, cfg, kind):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.d_model, cfg.rms_norm_eps)
+        self.ffn_norm = RMSNorm(cfg.d_model, cfg.rms_norm_eps)
+        self.attn = KDA(cfg) if kind == "kda" else RotaryAttention(cfg)
+        self.ffn = SiTUGLU(cfg)
+
+
+@dataclass
+class Cache:
+    layers: list
+    length: int
+
+
+class Sabaki(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.config = cfg
+        self.token_embedding = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.blocks = nn.ModuleList(Block(cfg, kind) for kind in cfg.attention_types)
+        self.final_norm = RMSNorm(cfg.d_model, cfg.rms_norm_eps)
+
+    def forward(self, ids, cache=None, *, last_logits_only=True):
+        offset = 0 if cache is None else cache.length
+        if ids.ndim != 2 or ids.shape[1] < 1:
+            raise ValueError("Expected nonempty [batch, sequence] token IDs")
+        if offset + ids.shape[1] > self.config.context_length:
+            raise ValueError("Prompt and continuation exceed the context limit")
+        hidden = self.token_embedding(ids)
+        states = []
+        for index, block in enumerate(self.blocks):
+            attention, state = block.attn(block.attn_norm(hidden), None if cache is None else cache.layers[index])
+            hidden = hidden + attention
+            hidden = hidden + block.ffn(block.ffn_norm(hidden))
+            states.append(state)
+        hidden = self.final_norm(hidden)
+        logits = F.linear(hidden[:, -1:] if last_logits_only else hidden, self.token_embedding.weight)
+        return logits, Cache(states, offset + ids.shape[1])
+
+
+def load_model(path, device, dtype):
+    payload = torch.load(path, map_location="cpu", weights_only=True, mmap=True)
+    if payload.get("format") != "kurokiba-dense-v1":
+        raise ValueError("Expected a regular Sabaki Dense checkpoint; Sabaki Max and MoE are unsupported")
+    cfg = SabakiConfig.from_dict(payload["config"])
+    with torch.device("meta"):
+        model = Sabaki(cfg)
+    model.load_state_dict(payload["model"], strict=True, assign=True)
+    return model.to(device=device, dtype=dtype).eval()
+
+
+def sample(logits, temperature, top_p):
+    if temperature == 0:
+        return logits.argmax(-1)
+    probabilities = (logits.float() / temperature).softmax(-1)
+    if top_p < 1:
+        probabilities, indices = probabilities.sort(descending=True, dim=-1)
+        remove = probabilities.cumsum(-1) - probabilities >= top_p
+        probabilities = probabilities.masked_fill(remove, 0)
+        probabilities = probabilities / probabilities.sum(-1, keepdim=True)
+        return indices.gather(-1, torch.multinomial(probabilities, 1)).squeeze(-1)
+    return torch.multinomial(probabilities, 1).squeeze(-1)
+
+
+def main():
+    directory = Path(__file__).resolve().parent
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--checkpoint", type=Path, default=directory / "chat-grounded-500.pt")
+    parser.add_argument("--tokenizer", type=Path, default=directory / "tokenizer.json")
+    parser.add_argument("--prompt", required=True)
+    parser.add_argument("--system")
+    parser.add_argument("--raw", action="store_true", help="plain continuation, without chat role/control tokens")
+    parser.add_argument("--device", default="auto", help="auto, cpu, cuda, or cuda:N; honors CUDA_VISIBLE_DEVICES")
+    parser.add_argument("--dtype", choices=("auto", "float32", "bfloat16", "float16"), default="auto")
+    parser.add_argument("--max-new-tokens", type=int, default=128)
+    parser.add_argument("--temperature", type=float, default=0.2, help="0 selects greedy decoding")
+    parser.add_argument("--top-p", type=float, default=0.95)
+    parser.add_argument("--seed", type=int, default=20260729)
+    parser.add_argument("--threads", type=int, default=4)
+    args = parser.parse_args()
+    if args.max_new_tokens < 1 or args.threads < 1:
+        parser.error("max-new-tokens and threads must be positive")
+    if not math.isfinite(args.temperature) or args.temperature < 0 or not 0 < args.top_p <= 1:
+        parser.error("temperature must be finite and nonnegative; top-p must be in (0, 1]")
+    if args.raw and args.system is not None:
+        parser.error("--system cannot be combined with --raw")
+    for path in (args.checkpoint, args.tokenizer):
+        if not path.is_file():
+            parser.error(f"Missing file: {path}")
+    torch.set_num_threads(args.threads)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if args.device == "auto" else torch.device(args.device)
+    if device.type not in {"cpu", "cuda"}:
+        parser.error("Supported devices are cpu and cuda[:N]")
+    if device.type == "cuda":
+        if not torch.cuda.is_available():
+            parser.error("CUDA is unavailable; use --device cpu")
+        torch.cuda.set_device(device)
+    dtype_name = args.dtype
+    if dtype_name == "auto":
+        dtype_name = "bfloat16" if device.type == "cuda" and torch.cuda.is_bf16_supported() else "float32"
+    model = load_model(args.checkpoint, device, getattr(torch, dtype_name))
+    tokenizer = ByteBPETokenizer.load(args.tokenizer)
+    if tokenizer.vocab_size != model.config.vocab_size or tokenizer.trained_vocab_size != tokenizer.vocab_size:
+        parser.error("A fully defined matching tokenizer is required")
+    if args.raw:
+        prompt = tokenizer.encode(args.prompt)
+    else:
+        parts = [("control", "<|bos|>")]
+        if args.system is not None:
+            parts += [("control", "<|system|>"), args.system, ("control", "<|system_end|>")]
+        parts += [("control", "<|user|>"), args.prompt, ("control", "<|user_end|>"), ("control", "<|assistant|>")]
+        prompt = tokenizer.encode_parts(parts)
+    if not prompt or len(prompt) >= model.config.context_length:
+        parser.error("Prompt must be nonempty and leave space for generation within the context limit")
+    print(f"Sabaki | {device} | {dtype_name} | prompt: {len(prompt)} tokens", file=sys.stderr)
+    limit = min(args.max_new_tokens, model.config.context_length - len(prompt))
+    stop = {tokenizer.control_id(token) for token in ("<|final_end|>", "<|assistant_end|>", "<|eos|>")}
+    torch.manual_seed(args.seed)
+    generated = []
+    with torch.inference_mode():
+        logits, cache = model(torch.tensor([prompt], device=device))
+        for index in range(limit):
+            token = int(sample(logits[:, -1], args.temperature, args.top_p).item())
+            generated.append(token)
+            if token in stop or index + 1 == limit:
+                break
+            logits, cache = model(torch.tensor([[token]], device=device), cache)
+    if not args.raw and tokenizer.control_id("<|final|>") in generated:
+        generated = generated[generated.index(tokenizer.control_id("<|final|>")) + 1:]
+    print(tokenizer.decode(generated, errors="replace", skip_control=True))
+
+
+if __name__ == "__main__":
+    main()