Kimi K3 engine: native-MXFP4 expert streaming, KDA + gated NoPE MLA + AttnRes + LatentMoE (#658)

[steve-m]

This adds support for Kimi K3
(2.8T parameters, 104B active, 93 layers) as a new sibling engine
c/kimi_k3.c, following the one-engine-per-family pattern (olmoe.c,
inkling.c) — colibri.c is untouched, and the shared-header changes are
purely additive (make check stays green).

Closes #658. The engine implements exactly the split discussed there:

• The routed experts (896/layer, top-16; 93% of the checkpoint's bytes) stay
  native MXFP4 — streamed straight from the original safetensors shards
  and computed on as-is via a new matmul_mxfp4 kernel (e2m1 nibbles +
  ue8m0 per-32 scales, scalar + AVX2). QAT weights are never re-encoded;
  there is deliberately no sub-4-bit expert option, since that would be
  double quantization of QAT weights.
• The engine runs directly on the original HF snapshot (96 shards, as
  downloaded — nothing on disk is modified). The BF16 dense side is
  quantized into RAM at load time (K3_BITS=4|8|32, with separate knobs
  for the MLA projections and lm_head), which is the "convert the dense on
  the fly" option from the [Feature]: Kimi K3 Original Weights loading #658 thread. An optional repack
  (tools/k3_repack.py, below) trades a one-time disk copy for a 42 s cold
  start — also with the experts byte-identical.

K3 is architecturally unlike anything the repo currently runs, which is why
it gets its own engine rather than branches in the GLM one:

• 69 KDA layers (Kimi Delta Attention — a gated delta-rule linear
  recurrence with per-channel bounded decay, short causal convolutions and a
  full-rank sigmoid output gate) interleaved 3:1 with 24 gated MLA
  layers; NoPE — no positional encoding anywhere in the model.
• AttnRes: the plain residual stream is replaced by learned softmax
  mixes over block-level snapshots of the stream (layers 0, 12, …, 84),
  applied twice per layer plus once before the head.
• Stable LatentMoE: the routed experts run in a 3584-dim latent space
  behind shared down/up projections, plus two fused full-width shared
  experts, with the SiTU-GLU activation 4·tanh(g/4)·σ(g)·25·tanh(u/25).

Full architecture/format details and the operational manual are in the new
docs/kimi_k3.md.

It runs the full model

First full-93-layer runs on a Ryzen 9 3900X / 62 GiB / single PCIe-4 NVMe,
greedy, prompt "The theory of general relativity was developed by":

config	resident RAM	cold start	prefill	decode	output
dense int8	56.9 GiB	42 s	~2.0 s/tok	~5.2 s/tok	"…Einstein when he felt that the theory of special relativity was insufficient to describe the"
dense int4 (K3_BITS=4)	37.3 GiB	131 s	~2.0 s/tok	~5.2 s/tok	"…Einstein in 1915, and it"
+ K3_TOPP=0.7 (opt-in)	—	—	—	~3.0 s/tok	quality note below

25.8 GB of experts are streamed per token (worst case, and with K3's
Quantile-Balancing-flat routing the observed case): a token's expert misses
are read in parallel with O_DIRECT (K3_DIRECT=0 for buffered),
reaching 6.3 GB/s against the drive's measured 7.1 GB/s ceiling, and the
preads run on loader threads so expert j's matmuls overlap expert j+1's
read (K3_PIPE). The expert matmuls use int8-quantized activations with
integer dots by default (K3_IDOT; measured free at text level — see the
validation table). Per-token decode budget on the test box: un-hidden I/O
wait 2.5 s + expert compute 1.9 s + attention 0.9 s. Multi-drive splits
work via K3_DIRS (no duplication). Decode is expert-bandwidth-bound by
construction, so it scales with disk speed.

Chat works natively (--chat "msg" [--system "..."]): K3's XTML format
(from the checkpoint's own encoding_k3.py) uses four special tokens around
ordinary-text tags, with the assistant's preserved thinking as a
structural <think> channel (K3_THINK=0 for response-only). The C builder
is verified token-exact against encoding_k3.py + tiktoken, and the
full model uses the format natively — greedy "how many r in strawberry":
reasons in the think channel (spells the word, finds positions 3/8/9), opens
the response channel, answers "3 times". The CLI hides the markup and prints
[think]/[response] banners. (Single system+user turn for now; the format
also covers multi-turn and tool calls — see docs/kimi_k3.md.)

Prefill runs in K3_CHUNK-token chunks (default 32), layer-major: every
dense matmul batches over the chunk (each weight matrix streams from RAM
once per chunk instead of once per token), the MoE loads each unique expert
of the chunk once, and the lm_head runs only on the chunk's last token.
Sequential state (KDA recurrence, MLA cache, AttnRes) advances per token
inside each layer in the original order, so chunked results are
bit-identical to token-at-a-time (validation table). One measured
surprise worth recording: within a 32-token window the chunk's expert union
deduplicates 2.7× (9.6 GB streamed/token instead of 25.8) — K3's routing
is Quantile-Balancing-flat in its marginals but has strong temporal
locality across neighbouring tokens.

K3_TOPP=p (default off) keeps routed experts to cumulative gate weight p,
renormalized — at 0.7 it keeps ~8.1 of 16 experts and reaches ~3.0 s/token,
but real-text ppl moves 2.694 → 2.811 (ΔNLL +0.043; next-token accuracy
unchanged at 73.4%). That is a measurable trade, so it ships opt-in,
consistent with the quality-first stance in the #658 thread; any setting is
easy to gate with a K3_LOGITS A/B.

What's in the commits

1. quant.h: MXFP4 matmul kernel. Verified against an independent numpy
   implementation of the compressed-tensors spec: rel err 1.2e-7; AVX2 ≡
   scalar.
2. tok.h: a third pre-tokenizer family for Kimi's tiktoken pattern
   (leading \p{Han}-run rule, Han-masked letter classes), and a
   rank-BPE mode used when model.merges is empty: merge the adjacent
   pair whose concatenation has the lowest vocab id — tiktoken's own
   algorithm, exact by construction. (Merge-list recovery from ranks is
   provably lossy; the commit message has the counterexample.) Existing
   tokenizers keep the merges path bit-for-bit.
3. c/kimi_k3.c + make kimi_k3 + docs/kimi_k3.md: the engine (~950
   lines). Per-layer expert LRU; one pread per expert (the six tensors of
   an expert are stored back-to-back in the shards — measured), issued in
   disk-offset order (experts are not id-ordered on disk — measured).
4. tools/: k3_tokenizer.py (tiktoken.model → tokenizer.json + tiktoken
   cross-checks), k3_ref.py (independent numpy reference of the full
   layer stack), k3_repack.py (optional: one streaming read of the HF
   snapshot → pre-quantized container; experts byte-identical, verified).
5. Streaming/RAM fixes from the full-model bring-up: parallel O_DIRECT
   expert reads (26–28 → 9.2–10.7 s/token), honest K3_EXPERT_GB
   accounting, and an explicit-K3_BITS=4 downcast for int8 containers
   (57 → 37 GiB resident, fits beside a desktop session).
6. K3_LOGITS: per-prefill-position logit dump (the teacher-forced
   bit-width comparison below).
7. quant.h: matmul_mxfp4_i8 — int8-activation MXFP4 matmul. The doubled
   e2m1 values are exact int8, so a 32-group reduces to
   maddubs(|w|, sign(x,w)) (the dot_i4i8 idiom) with per-group scales
   folded exactly. 3x on expert compute (5.8 → 1.9 s/token).
8. K3_PIPE loader threads (overlap preads with compute) + K3_TOPP
   (opt-in expert pruning). Decode 9.4 → 5.2 s/token at defaults.
9. Chunked prefill (K3_CHUNK, bit-identical by construction and by
   measurement) + AVX2 KDA state sweeps (attention 0.9 → ~0.5 s/token,
   ppl unchanged at 4 decimals). Prefill ~5.3 → 2.0 s/token.
10. --chat: the XTML chat format with the preserved-thinking channel,
    token-exact vs the checkpoint's encoding_k3.py.

Validation

Everything was validated on real checkpoint weights, not synthetic
tensors:

what	result
engine (f32) vs independent numpy reference, 4 tokens × 4 layers covering all four layer types (KDA+dense, KDA+MoE ×2, gated-MLA+MoE)	rel-L2 ≤ 2.2e-6 every layer (f32 noise)
MXFP4 kernel vs numpy spec implementation	rel err 1.2e-7; AVX2 ≡ scalar
C tokenizer vs tiktoken (--ctest, 18 adversarial cases: Chinese, kanji-vs-kana, Korean, CRLF, contractions, emoji, mixed-script boundaries)	18/18 exact
repacked container vs source	every expert byte re-read and compared (--verify-full): identical; quantized tensors bit-identical to the engine's load-time quantization
full 93-layer generation	coherent, factually correct (table above)
dense int8 vs int4, teacher-forced real-text logits (124 positions, K3_LOGITS)	ppl 2.69 vs 2.71; next-token accuracy identical 73.4%; disagreements symmetric (2 wins each)
K3_IDOT int8-activation kernel vs float kernel, same protocol	ppl 2.684 vs 2.694; argmax agreement 98.4%; mean KL 0.003 — free
chunked prefill K3_CHUNK=32 vs =1, 125-position logit streams	bit-identical (max abs diff 0.0)
AVX2 KDA sweeps vs scalar, same protocol	100% argmax agreement; ppl 2.6841 unchanged
--chat prompt ids vs encoding_k3.py + tiktoken (±system, ±thinking)	token-exact
make check on this branch	green

An honest caveat on the reference: k3_ref.py is an independent
implementation, but it encodes the same reading of the released
modeling_kimi_linear.py + fla kernels as the C engine, so a shared
misreading of the modeling code would not be caught. Running HF+fla itself
needs triton/GPU and wasn't feasible here; the formulas were taken
line-by-line from the shipped modeling source and the fla
fused_recurrent_kda kernel (including two checkpoint quirks verified
empirically: A_log is stored [128] = per-head [96] zero-padded, and
dt_bias is per-channel [12288] despite the fla docstring saying [HV]).

One measured property reviewers should know: the AttnRes mix logits have
small margins (~0.02), so the architecture leverages weight-quantization
noise unusually strongly (load-time int8 drifts hidden states ~14% over four
layers on synthetic inputs, while a 1e-4 input perturbation does not
amplify at all). This affects only the dense side — the experts are
untouched QAT bytes, per #658.

That comparison has now been run on real text (teacher-forced over a
124-position factual passage, identical explicit ids, full logits per
position via the K3_LOGITS hook, --ngen 0):

• perplexity on the passage: int8 2.69 vs int4 2.71
  (ΔNLL mean +0.006 nats, median +0.0007);
• next-token accuracy against the text: identical 73.4% for both;
• int8-vs-int4 argmax agreement 92.7% (9/124 differ) and the disagreements
  are symmetric: 2 positions where int8 is right and int4 wrong, 2 the
  other way, 5 where both are wrong — no directional quality loss;
• mean KL(int8‖int4) 0.038 nats (p90 0.10): individual distributions do
  drift (the AttnRes signature, max |ΔNLL| 1.4 at one position) but the
  drift does not bias predictions.

So the synthetic hidden-state sensitivity does not surface as text-level
quality loss: dense int4 is a sound default (37 GiB resident), int8 the
belt-and-braces choice when RAM allows — and since both decode at the same
speed (I/O-bound), the choice is purely RAM. Caveats stated plainly: one
English passage, teacher-forced, and no f32 gold reference (the f32 dense
does not fit the 62 GiB test box, so this is a relative comparison —
consistent with the finer int8 grid, int8 remains the reference).

Scope notes / follow-ups

The originally-open WIP items are done and validated above (bits comparison,
chunked prefill + KDA SIMD, chat template). Deliberate bounds that remain,
listed in docs/kimi_k3.md:

• CPU-only; decode is single-token (K3 has no MTP head to speculate with).
• --chat renders a single system+user turn (the format itself covers
  multi-turn and tool calls).
• A bf16-resident dense mode for large-RAM hosts could follow (the "leave
  dense untouched where RAM allows" arm of [Feature]: Kimi K3 Original Weights loading #658; K3_BITS=32 already gives
  an f32 superset of it at 2× the RAM, and the measured int8 delta on real
  text is within noise).

Interaction with other open PRs

colibri.c is untouched. quant.h/tok.h changes are self-contained
appends (rebased cleanly over the merged #654/#655); tok.h's rank-BPE mode
is gated on tokenizer features no existing model uses. No overlap with the
open #666/#667/#671 beyond the shared one-engine-per-family pattern.

Reproducing the numeric validation needs only shards 1–4 of the checkpoint:
k3_ref.py gen/run/cmp + K3_LAYERS=4 K3_BITS=32 K3_X0=… K3_TRACE=… (exact
commands in docs/kimi_k3.md).

[JustVugg]

@steve-m when tests for you are ok i merge it!

[ZacharyZcR]

I built the shared chat/API/Web integration on top of this PR in #684. The K3 multi-turn wire fixture matches Moonshot official encoding_k3.py exactly (77/77 token IDs), and the local C/Python regression suite passes. Full-model generation remains explicitly gated on a host large enough to load the snapshot.