Adding kimi-k3 arch test and saver fix - #21
Open
SolshineCode wants to merge 9 commits into
Open
Conversation
Hybrid KDA (linear) + MLA (full) attention as in Kimi-Linear-48B, plus five things that architecture does not have: 1. cross-layer residual attention (attn_res_block_size) 2. latent MoE (routed experts run at n_expert_latent) 3. situ activation (replaces SwiGLU everywhere) 4. MLA output gate (sigmoid gate before o_proj) 5. full-rank KDA gate (single ssm_g instead of ssm_g_a/ssm_g_b) K3's text_config reports KimiLinearForCausalLM - the older 48B architecture - so get_model_architecture routes on the top-level name instead. The KDA decay gate has two forms, selected by linear_attn_config's gate_lower_bound. It is not a clamp: when set it swaps the activation entirely (fla/ops/kda/gate.py), from -exp(A_log)*softplus(x) to lower_bound*sigmoid(exp(A_log)*x). K3 sets it to -5.0; kimi-linear leaves it unset, so that path is unchanged. Cross-layer residuals reuse ggml_dsv4_hc_pre for the weighted sum. That op is CPU + CUDA only, so Metal/Vulkan will fall back per-node until those kernels exist. The routed experts ship as compressed-tensors "mxfp4-pack-quantized". That is bit-compatible with ggml's MXFP4 - same E2M1 code assignment, same E8M0 scale byte, only the nibble positions within a block differ - so they are repacked rather than dequantized, losslessly and without a ~5.5 TB bf16 round-trip. The repack is built lazily because gguf_writer holds every added tensor until the final write. DeepSeek-V4 was already doing the identical bit-shuffling, so it now shares the helper. Verified against Moonshot's own code path (transformers + fla's Triton KDA kernels) on a tiny model exercising every K3-specific feature. Final-position logits vs the fp32 reference: 6.7e-05 rel / corr 1.00000000 for both the chunked and the recurrent delta-net path. MXFP4 blocks dequantize to the source weights with 0.0e+00 error. Assisted-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
- `_res_parts` buffers (kind, tensor) pairs, not bare tensors - `get_tensors` must return an Iterator, matching ModelBase - LazyBase's `func` takes one argument, so pass the expert loaders through `args` instead of the closure - borrowing KimiLinearModel.set_vocab from an unrelated TextModel is deliberate and safe, but not expressible in the signature No behaviour change: the MXFP4 repack still dequantizes to the source weights with 0.0e+00 error and end-to-end logits are unchanged (8.386e-03 rel, corr 0.99996630). Assisted-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Boris Dvorkin <b_dvorkin@niuitmo.ru>
K3's assistant output is an XTML-ish tagged format built by the template's
open_tag/close_tag macros. Two properties break generic parsing:
1. The generation prompt ends with open_tag('think'), so the completion
starts inside the think section with no opening marker in the output
(thinking_forced_open).
2. Only <|open|>/<|close|>/<|sep|>/<|end_of_msg|> are special tokens; tag
names ("think", "response", "message") are ordinary text tokens.
Adds common_chat_params_init_kimi_k3 (PEG_NATIVE) with detection on the
marker trio, reasoning extraction, response unwrapping, and tool-call
parsing of the tools/call/argument tag structure with argument types
taken from the tool schema. Includes the K3 chat template fixture and 9
test-chat cases derived from real generations of the full 2.8T model.
Verified end-to-end against Kimi-K3-Q2_K (GrEarl/Kimi-K3-GGUF) on 8x B200:
content, reasoning_content, streaming deltas, and tool_calls all correct;
finish_reason stop/tool_calls as appropriate.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Per-role message-start markers for token-level span splitting. User and assistant messages carry only the role attribute, so their full opener (through <|sep|>) is used; system and tool messages continue with more attributes (type=/tool=/index=), so those delimiters stop after the role's closing quote. Verified against the K3 tiktoken vocabulary that the closing quote is always a standalone token across all attribute variants, so the token-level prefix match stays exact. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…rser chat : add Kimi K3 chat format (reasoning, content, typed tool calls)
SolshineCode
force-pushed
the
kimi-k3-arch-test
branch
from
July 28, 2026 17:26
7c0d9b6 to
7438888
Compare
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
4 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.
This adds kimi-k3 to test-llama-archs so the arch gets NMSE and GGUF-roundtrip coverage in CI. This is the execution half of the coverage gap 200lz just flagged in ggml-org#26185 and their proposed converter-level schema test covers the other half, with (I believe) no overlap between the two.
The test model is 5 layers, consisting of dense layer 0, KDA at 0-2, MLA at 3 and 4. That's the smallest shape that exercises the hybrid schedule including the consecutive-MLA ending the released checkpoint has, plus attention residuals (block size 2, so two block boundaries), latent MoE (expert_latent_length 64 < n_embd 128, so the latent projections exist and get tested), situ at the released beta values (4.0 / 25.0), and the KDA gate lower bound (-5.0). head_count stays scalar and head_count_kv is the per-layer vector with 0 marking KDA layers, matching what conversion/kimi_k3.py writes for the real model.
Running this red (TDD) first surfaced two roundtrip gaps, fixed here in llama-model-saver.cpp:
The five new K3 KVs (attn_res.block_size, the situ betas, expert_latent_length, kda gate lower bound) weren't emitted by the saver, so a save/load roundtrip silently rebuilt the model without attention residuals and failed the bit-exact logits check.
output_res_score wasn't in the saver's top-level tensor list, so reload died on a missing tensor.
The results on my hardware as reported by Claude: "CPU (2x Xeon E5-2609v2, AVX only) kimi-k3 NMSE 0.00e+00 and roundtrip OK; CUDA sm_52 (2x Tesla M40) NMSE 2.15e-14 and roundtrip OK on both cards; full all-arch sweep is 119 rows OK with no regressions (kimi-linear unchanged at 0.00e+00 / 1.86e-14)."
AI usage disclosure: YES. The code written was with Claude Code Fable 5 on my machine, using the red/green process above. Also I reviewed every hunk and can walk through them further as needed.