Kimi K3: integrate coli chat, API, and Web - #684
Closed
ZacharyZcR wants to merge 11 commits into
Closed
Conversation
…alar + AVX2 The native layout of compressed-tensors 'mxfp4-pack-quantized' checkpoints: packed [O,I/2] u8 (low nibble = even column, bit3 = sign) + scales [O,I/32] ue8m0, w = v * 2^(s-127). Computes on the vendor bytes directly - QAT weights (Kimi K3 routed experts) are exactly the trained values, so they are streamed and multiplied as-is, never re-encoded. AVX2 decodes a 16-byte group with one pshufb over the doubled-e2m1 int8 LUT; the 0.5 un-doubling rides the group scale. Scalar and SIMD paths use the same bit-trick exponent decode so both agree on the s=0/255 edge cases. Verified against an independent numpy implementation of the spec (rel err 1.2e-7 = f32 noise) and byte-conventions against compressed-tensors' pack_fp4_to_uint8/unpack_fp4_from_uint8.
… vocabs
Two additions for Kimi K3's tokenizer (tiktoken vocab, fancy-regex pattern):
- kimi Split family (sniffed via \p{Han} in the pre_tokenizer regex): the
o200k case-aware rules plus a leading [\p{Han}]+ rule, Han masked out of
the letter classes (fancy-regex &&[^\p{Han}] intersections), and no '/'
tail in the punctuation rule. A Han codepoint can only match the Han rule,
so Han-free text tokenizes exactly like o200k minus the '/' tail.
- rank-BPE: when model.merges is empty, bpe_piece merges 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 standard max-rank-minimizing reconstruction mis-merges 'newlines'
as new+l+ines), which is why K3's generated tokenizer.json ships no merges.
tests/test_tok_kimi.c + tools/k3_tokenizer.py --ctest cross-check the C
encoder against tiktoken: exact on all 18 corpus cases (Chinese, kanji vs
kana, Korean, CRLF, contractions, emoji, mixed-script boundaries). Existing
families are untouched: merges-bearing tokenizers keep the merges path.
New sibling engine for Kimi K3 (2.8T / 104B active, 93 layers), following the one-engine-per-family pattern (colibri.c untouched). Shares st.h / json.h / tok.h / quant.h. make kimi_k3; ./kimi_k3 <model_dir> "prompt" --ngen N. - 69 KDA layers: SiLU(causal-conv4) projections, L2-normed q/k, per-channel bounded decay gk = -5*sigmoid(exp(A_log[h])*(f_b(f_a x)+dt_bias)), delta rule S=(I-beta k k^T)Diag(e^gk)S + beta k v^T, per-head RMSNorm x sigmoid full-rank output gate. Checkpoint A_log [128] = per-head [96] zero-padded. - 24 gated MLA layers, absorb form, NoPE (nothing is rotated; the 64 extra dims are cached raw), sigmoid output gate before o_proj. - AttnRes: the residual stream is replaced by softmax mixes over block snapshots (layers 0,12,...,84) + the running prefix sum - twice per layer and once before the head, scored against res_norm.w*res_proj.w, raw-entry mixing in fp32. - Stable LatentMoE: sigmoid router + score-correction bias, top-16, raw-score renormalization; latent 7168<->3584 projections around the experts; RMSNorm on the aggregate; 2 fused shared experts; SiTU-GLU (4 tanh(g/4) sigmoid(g) * 25 tanh(u/25)). - Routed experts stream from the safetensors shards as native MXFP4 through a per-layer LRU (K3_EXPERT_GB): one pread per expert (the six tensors are stored back-to-back - measured), loads issued in disk-offset order (experts are NOT id-ordered in the shards - measured), multi-drive shard split via K3_DIRS (st_init_multi). - Non-expert BF16 is quantized at load (K3_BITS int4-g64 / int8 / f32, K3_MLA_BITS, K3_HEAD_BITS) or read pre-quantized from a repacked container (U8 + .qs sidecars, tools/k3_repack.py) - container start measured 30.4s -> 0.6s on two layers. Validation (docs/kimi_k3.md): with injected inputs and f32 weights the engine matches an independent numpy reference (tools/k3_ref.py) on real checkpoint weights across all four layer types to rel-L2 <= 2.2e-6; container-vs-load-time quantization is bit-identical (layer-0 trace equal); tokenizer exact vs tiktoken. Known limits: single-token stepping (sequential prefill), CPU only, no chat template.
- k3_tokenizer.py: tiktoken.model -> tokenizer.json (byte-level vocab, EMPTY merges on purpose -> tok.h rank-BPE, kimi Split regex, 256 added tokens from tokenizer_config). --check replays rank-BPE in python vs tiktoken; --ctest drives tests/test_tok_kimi against tiktoken. 18/18 exact. - k3_ref.py: independent numpy implementation of the K3 stack (KDA recurrence, gated NoPE MLA, AttnRes, LatentMoE with MXFP4 dequant) reading tensors straight from the HF shards; gen/run/cmp subcommands produce and compare the engine's K3_TRACE hidden-state dumps. This is the reference the engine is validated against (rel-L2 <= 2.2e-6 at f32 on real weights). - k3_repack.py: one streaming read of the HF snapshot -> ideal ingestion container, so a disk-to-disk copy IS the conversion. Experts byte-identical in id order (six tensors back-to-back = the engine's slot layout, one pread per expert), big BF16 mats quantized with the ENGINE-EXACT algorithm (int8/row or int4-g64, --bits/--mla-bits/--head-bits) into U8 + .qs f32 sidecars, sensitive smalls passed through, vision dropped, spec-valid safetensors (no holes) + regenerated index, per-shard resume, --verify-full re-reads every expert byte. Kills the engine's 10-15 min load-time quantization; quantized bytes verified bit-identical to the load-time path. ~1.56 TB source -> ~1.50 TB (int8) / ~1.48 TB (int4).
Full-model bring-up on a 62 GiB box exposed three streaming-path issues: 1. Expert loads were sequential buffered preads; with ~55 GiB of resident weights the page cache has no headroom and reclaim throttled reads to ~1.8 GB/s (drive: 7.1 GB/s O_DIRECT, 2.9 buffered — measured). A token's misses now load in parallel into working-set slots and read O_DIRECT through st.h's twin fds (aligned-window + tiny buffered tail; K3_DIRECT=0 reverts to buffered + WILLNEED). Expert reads: 6.3 GB/s. Decode: 26-28 -> 9.2-10.7 s/token on the 93-layer model. 2. The LRU cap floor was topk, silently committing topk*92 slots (~26 GiB) regardless of K3_EXPERT_GB; experts are consumed one at a time, so the floor is now 1 and the budget is honored. 3. Explicit K3_BITS=4 on an int8 container downcasts to int4-g64 at load: ~35 vs ~57 GiB resident (fits next to a desktop session), same decode speed once I/O is fixed; unset K3_BITS keeps container bits (default unchanged). 93-layer sanity (greedy, repacked int8 container): 'The theory of general relativity was developed by' -> int8: 'Einstein when he felt that the theory of special relativity was insufficient...', int4: 'Einstein in 1915, and it'. Init 42 s (int8) / 131 s (downcast pass).
Teacher-forced bit-width comparisons (with --ngen 0): run the same --ids twice under different K3_BITS and diff the logit streams offline. Used for the int8-vs-int4 dense-side quality evaluation on real text.
The float kernel converts every decoded weight nibble to float before the
FMA; at 25.8 GB of packed expert bytes per K3 token that path measured
~4.4 GB/s effective. The doubled e2m1 values are exact int8 (same LUT), so
a 32-group reduces to maddubs(|w|, sign(x,w)) + madd — the dot_i4i8 idiom —
with the group scale folded once per group:
y = sum_g idot_g * 2^(e8-127)*0.5 * xscale_g
Activations quantize per 32-group to int8 (~0.4% group noise, aligned with
the weight groups so scale folding is exact). Unit test vs the float kernel:
rel-L2 3.5e-3 on random data. On the full model: teacher-forced real-text
ppl 2.684 vs 2.694 (float), argmax agreement 98.4%, KL mean 0.003 — free at
text level, 3x faster (expert compute 5.8 -> 1.9 s/token). Scalar fallback
included; I%32!=0 falls back to the float kernel.
Decode 9.4 -> 5.2 s/token (and 3.0 opt-in), on top of the O_DIRECT work: - K3_PIPE (default 1): a small loader-thread pool (K3_LOAD_THREADS, 4) runs the expert preads while the compute loop chews already-loaded experts — expert j's matmuls overlap expert j+1's read. t_eload now reports the UN-hidden wait (2.5 s/token; pure I/O would be 3.8). - K3_IDOT (default 1): expert matmuls via quant.h's new matmul_mxfp4_i8. Real-text gate: ppl 2.684 vs 2.694 float, agreement 98.4% — free. - K3_TOPP (default 0 = off): drop the low-weight tail of the top-16 to cumulative weight p, renormalized (GLM TOPP semantics) — the only lever cutting I/O and compute together. At p=0.7 it keeps ~8.1/16 experts: 3.0 s/token, but ppl 2.811 vs 2.694 (dNLL +0.043) with accuracy unchanged at 73.4% — measurable, so it stays OPT-IN; gate any setting with a K3_LOGITS A/B first. Full-model numbers, greedy sanity text unchanged. make check green.
Prefill now runs layer-major over K3_CHUNK-token chunks (default 32): - every dense matmul batches over the chunk, so each weight matrix streams from RAM once per chunk instead of once per token; - the MoE builds the chunk's expert union per layer and loads each unique expert ONCE (position lists via counting sort), feeding the same pipelined block loader as decode. Measured on real text at C=32: 2.7x dedup — neighbouring tokens share experts far more than QB-flat routing predicts — 9.6 instead of 25.8 GB streamed per token; - the lm_head runs only on the chunk's last token (K3_LOGITS still dumps every position); - sequential state (KDA recurrence, MLA cache append, AttnRes bookkeeping) advances per token inside each layer in the original order, so chunked results are BIT-IDENTICAL to token-at-a-time: the 125-position teacher-forced logit streams at C=32 vs C=1 match exactly (max abs diff 0.0). K3_TRACE forces C=1 (trace rows are token-major by contract). The KDA delta-rule state sweeps get explicit AVX2 (mul+fmadd over the 128-wide state rows, scalar fallback when head_dim%8!=0): attention 0.9 -> ~0.5 s/token; vs the pre-SIMD logits: 100% argmax agreement, ppl 2.6841 unchanged at 4 decimals. Full-model prefill: ~5.3 -> 2.0 s/token (125-token prompt, C=32, int8 dense, K3_LOGITS on — pure generation prefill is slightly better still). Decode unchanged (same C=1 path). make check green.
…hannel
K3 ships its chat renderer as encoding_k3.py: an XTML structure where only
<|open|>/<|close|>/<|sep|>/<|end_of_msg|> are special tokens and tag names /
attributes are ordinary text, encoded as standalone segments. A turn is
<|open|>message role="user"<|sep|>TEXT<|close|>message<|sep|><|end_of_msg|>
and the generation prompt opens the assistant message plus its STRUCTURAL
thinking channel (<|open|>think<|sep|>; K3_THINK=0 opens <response> directly
— in thinking mode the channel exists in every assistant message, which is
the preserved-thinking property).
The C builder mirrors the reference segment-for-segment (segment boundaries
are token boundaries, so per-segment BPE matters) and is verified TOKEN-EXACT
against encoding_k3.py + tiktoken for system/no-system and thinking/
non-thinking prompts. The CLI hides the structure at decode and prints
[think]/[response] banners; <|end_of_msg|> is already the configured eos.
Full-model greedy check ('how many r in strawberry'): the model reasons in
the think channel (spells the word, finds positions 3/8/9), closes it, opens
response, answers '3 times' — native use of the format end to end.
Single system+user turn only for now (no multi-turn CLI, no tool-call
rendering); make check green.
ZacharyZcR
force-pushed
the
feat/kimi-unified-chat
branch
from
July 29, 2026 11:21
a49bddf to
ca3dcc6
Compare
Contributor
|
Thanks @ZacharyZcR — this was exactly the right direction, and #676 has now absorbed it: the branch is rebased onto current dev and carries your integration commit as-is ( Two things we added on top, from the box that holds the full snapshot:
So once #676 lands, this PR is fully covered by it. |
Owner
|
Closing as absorbed — not rejected. @steve-m cherry-picked your integration commit into #676 as-is with authorship preserved ( @ZacharyZcR — thank you: routing Kimi through the gateway instead of templating in @steve-m — thanks for handling the overlap cleanly rather than duplicating it. |
This was referenced Jul 29, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
SERVE=1protocol tokimi_k3coli chat,coli serve, andcoli webDepends on #676. This is intentionally stacked on its exact commits; once #676 lands, this PR should reduce to the integration commit.
Validation
make -C c kimi_k3make -C c check— 224 Python tests passed, 13 skipped, plus all C testsPYTHONPATH=c python3 -m unittest c.tests.test_openai_server— 104 passedencoding_k3.py: 77 token IDs, exact matchgit diff --checkFull-model test boundary
The local validation host has 93 GiB RAM and a 12 GiB RTX 4070, so it cannot load the full Kimi K3 snapshot. Tokenization, protocol, gateway, TUI/Web paths, compilation, and regression tests are covered here; a complete model generation run remains the release gate on a host that can load #676's snapshot.
Current limitations