kimi-k3: full-size model fixes and the MoonViT-3d vision tower - #44
kimi-k3: full-size model fixes and the MoonViT-3d vision tower#44danielhanchen wants to merge 12 commits into
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>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 50bcff46fd
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| ml.get_key(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent); | ||
| ml.get_key(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size); | ||
| ml.get_key(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta); | ||
| ml.get_key(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta); |
There was a problem hiding this comment.
Add the required Kimi-K3 keys to the architecture fixture
When test-llama-archs -a kimi-k3 constructs its synthetic GGUF, get_gguf_ctx() never emits EXPERT_LATENT_LENGTH, ATTN_RES_BLOCK_SIZE, or either situ parameter, so making these keys mandatory causes model loading to fail before any graph test runs. I built and ran this target and it exits with key not found in model: kimi-k3.expert_latent_length; update the Kimi-K3 fixture with valid values for all newly required keys.
Useful? React with 👍 / 👎.
| // Kimi-K3 stacks three node-heavy structures no other hybrid arch combines: 69 KDA | ||
| // layers whose chunked delta-net emits nodes per (chunk x layer), attention residuals | ||
| // that re-score a depth stack of up to 9 entries at every one of the 93 layers, and a | ||
| // latent MoE. The n_tokens*40 budget below is exhausted during graph_reserve at | ||
| // ubatch 3840 -- measured: ggml_new_tensor_impl aborts inside |
There was a problem hiding this comment.
Condense the graph-budget comment
This seven-line block embeds detailed measurements and implementation history while splitting several sentences across lines; retain only the non-obvious reason Kimi-K3 needs the larger node budget so the comment follows the repository's explicit concise-comment and sentence-formatting conventions.
AGENTS.md reference: AGENTS.md:L74-L76
Useful? React with 👍 / 👎.
chat : add Kimi K3 chat format (reasoning, content, typed tool calls)
bbe0afd to
fe31497
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b7320527a1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if "mm_projector.proj.0." in name: | ||
| name = name.replace(".proj.0.", ".proj.linear_1.") | ||
| elif "mm_projector.proj.2." in name: | ||
| name = name.replace(".proj.2.", ".proj.linear_2.") |
There was a problem hiding this comment.
Map the Kimi-K3 post-normalization tensor
When converting the Kimi-K3 multimodal projector, mm_projector.post_norm.weight reaches the superclass unchanged because this block only renames the two projection layers. The tensor map has no mm_projector.post_norm alias for V_MM_POST_NORM, so map_tensor_name() raises ValueError and prevents producing the mmproj GGUF. Add that alias or rename the tensor to the existing mm.post_norm name here.
Useful? React with 👍 / 👎.
| @ModelBase.register("KimiK3ForConditionalGeneration") | ||
| class KimiK3VisionModel(MmprojModel): |
There was a problem hiding this comment.
Register Kimi-K3 in the multimodal model map
When convert_hf_to_gguf.py --mmproj sees KimiK3ForConditionalGeneration, get_model_class(..., mmproj=True) consults MMPROJ_MODEL_MAP, where this architecture is absent, and raises NotImplementedError before conversion.kimivl is imported. Add a KimiK3ForConditionalGeneration: kimivl entry so this newly registered converter is reachable.
Useful? React with 👍 / 👎.
b732052 to
a0e0d2a
Compare
There was a problem hiding this comment.
💡 Codex Review
Lines 1439 to 1440 in a0e0d2a
When users pass --image-min-tokens or --image-max-tokens, these assignments overwrite the custom limits stored by clip_ctx with the GGUF metadata values, so Kimi-K3 always preprocesses at the model defaults despite the CLI contract in common/arg.cpp:2491-2503. This is especially problematic when a smaller maximum is requested to control vision memory use; apply the custom token values when deriving these pixel limits instead of unconditionally selecting min_pixels and max_pixels.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| n_head = self.hparams_vision["vt_num_attention_heads"] | ||
| qkv_hidden = self.hparams_vision.get("qkv_hidden_size") or self.hparams_vision["vt_hidden_size"] | ||
| assert qkv_hidden % n_head == 0, f"qkv_hidden_size {qkv_hidden} not divisible by {n_head} heads" | ||
| self.gguf_writer.add_vision_head_dim(qkv_hidden // n_head) |
There was a problem hiding this comment.
Define the vision head-dimension writer before calling it
Any Kimi-K3 mmproj conversion that reaches set_gguf_parameters() raises AttributeError here: this commit calls GGUFWriter.add_vision_head_dim(), but a repo-wide search of this commit shows neither that method nor a corresponding Keys.ClipVision.Attention key exists. Add the Python metadata key and writer method for the clip.vision.attention.head_dim value expected by the new C++ loader, otherwise no Kimi-K3 projector can be emitted.
Useful? React with 👍 / 👎.
| @classmethod | ||
| def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: | ||
| name, _ = item | ||
| if not name.startswith(("vision_tower.", "mm_projector.")): |
There was a problem hiding this comment.
Filter the unused temporal-position tensor
For MoonViT checkpoints containing vision_tower.patch_embed.pos_emb.time_weight, this prefix filter admits the tensor even though the Kimi-K3 graph is explicitly image-only and never consumes the temporal term. modify_tensors() then forwards it unchanged, while the repository's MMPROJ tensor map has no alias for this name, so conversion raises ValueError; exclude this video-only parameter before delegating to the base converter.
Useful? React with 👍 / 👎.
Graph node budget, LLM_TYPE for the 93-layer config, four hparams promoted to required, and an E8M0 NaN check during the MXFP4 repack. Details in the PR description. Assisted-by: Claude Code
New kimik3 projector type, its graph builder and the mmproj converter. Also adds an optional clip.%s.attention.head_dim so build_vit stops deriving d_head from n_embd, which is wrong whenever a tower's qkv width differs from n_embd. Assisted-by: Claude Code
a0e0d2a to
efc8bc3
Compare
Inkling and Kimi-K3 both register a new architecture, so they collide in the arch/model/mtmd registries and cannot be merged into the nightly as two independent heads. Both sides are kept: - src/llama-arch.cpp: KIMI_K3 and INKLING in llm_arch_is_hybrid and llm_arch_supports_sm_tensor - src/llama-model.cpp: both llama_model_* constructors and both LLAMA_ROPE_TYPE_NONE cases - tools/mtmd/CMakeLists.txt: build models/kimik3.cpp and models/inkling.cpp Assisted-by: Claude Code # Conflicts: # src/llama-arch.cpp # src/llama-model.cpp # tools/mtmd/CMakeLists.txt
The nightly died in resolve: #40 carried MiniMax-M3 plus Inkling, upstream master now ships its own MiniMax-M3, and the two collided in src/models/models.h and tests/test-backend-ops.cpp. Every build job was skipped. Drop the two MiniMax entries (ggml-org#24523 is closed anyway, so it was already inert) and repin on the four PRs we actually want: ggml-org#24423 DiffusionGemma ggml-org#25731 TML Inkling #44 Kimi-K3 fixes + MoonViT-3d vision tower ggml-org#26185 Kimi-K3 text ggml-org#25731 and #44 both had their conflicts against current master fixed on their own branches, and #44 also carries the Inkling merge: the two archs land in the same arch/model/mtmd registries, so they cannot go in as two independent heads and one of them has to know about the other. ggml-org#26185 is last on purpose. Its head is an ancestor of #44, so by the time the resolver gets there the merge is a no-op; listed before #44 it conflicts on tests/test-llama-archs.cpp. It stays in the set so the release manifest names it. Verified locally against b10173: the four merge in this order with zero unmerged paths, and the merged tree builds with -DGGML_CUDA=OFF -DLLAMA_CURL=OFF -DLLAMA_BUILD_TESTS=ON. Assisted-by: Claude Code
cf67f0d to
12e01bd
Compare
|
Superseded by #48, which is what the nightly pins now. This branch carried the Kimi-K3 text model, the Inkling merge and the vision tower all in one head. #48 is the same two Kimi-K3 commits stacked on ggml-org#26185 instead, so the diff stays reviewable and ggml-org#26185 merges as a no-op and still gets named in the release manifest. Its base branch moved when #48 was restacked, so the diff here no longer means anything (29 files, conflicting). Closing rather than leaving it to rot. |
Base branch
kimi-k3-text-baseis a pin of ggml-org#26185 at06eec9f5, so the diff below is only the changes added on top of that PR.Builds on ggml-org#26185 (Kimi-K3 text) so the full-size 2.72T checkpoint
loads, runs, and handles images. Two commits: text-model fixes, then the vision tower.
Text model
Graph node budget. The shared
n_tokens*40branch ingraph_max_nodesis exhaustedduring
graph_reserveat ubatch 3840, aborting inggml_new_tensor_implinsidellm_build_delta_net_base::build_delta_net_chunking. Kimi-K3 stacks three node-heavystructures no other hybrid arch combines: 69 KDA layers whose chunked delta-net emits nodes
per (chunk x layer), attention residuals that re-score a depth stack of up to 9 entries at
every one of the 93 layers, and a latent MoE. Nodes only size graph metadata, roughly 400 B
of host memory each, so the extra headroom is cheap.
Required hparams.
expert_latent_length,attn_res.block_size,situ_betaandsitu_linear_betawere read withrequired=false. Every Kimi-K3 checkpoint defines allfour, so an absent key means the KV name is wrong. Defaulting silently turns that into wrong
numbers rather than a load error:
situ_betafalls back to 1.0 instead of 4.0, which changesthe activation function, and
n_expert_latentfalls back to 0, which disables the latent-MoEprojections entirely. Both load cleanly and produce garbage.
LLM_TYPE_2_8T_A50B. 2.72T total parameters, roughly 50B active. Previously reported as
"unknown" by llama-server and the model print-out.
E8M0 NaN check. 0xff is the NaN exponent.
ggml_validate_row_dataalready rejects it,but only later during quantization and with an error that does not name the source tensor.
Checking during the repack keeps the tensor name in scope.
Vision
PR 26185 is text-only, so a converted Kimi-K3 has no vision encoder. This adds the image
path: a
kimik3projector type, its graph builder, and the mmproj converter.The tower follows the Kimi-K2.5 one (
clip_graph_kimik25) with four differences: RMSNorm andno biases anywhere,
qkv_hidden_size1536 againstvt_hidden_size1024, bilinear rather thanbicubic position-embedding interpolation, and a patchmergerv2 projector whose RMSNorm sits
after the projection at text-hidden width instead of before it at merged-patch width.
Video is out of scope. At t == 1 the temporal average pool is the identity and the divided
position embedding reduces to its 2D term, so the image path matches the reference exactly;
t > 1 is not supported.
build_vit fix, independent of Kimi-K3
clip_graphderivedd_headasn_embd / n_headand offset the fused K/V views byn_embd.Both are wrong for any tower whose qkv width differs from its embedding width. A new optional
clip.%s.attention.head_dimkey carries the real value and falls back to the old derivationwhen absent, and the K/V offsets now use
n_head * d_head, which is the actual layout of afused qkv tensor.
Checked against every mmproj available locally: for all of them
n_headdividesn_embd, son_head * d_head == n_embdand both expressions are identical. The only files where theydiffer are Kimi-K3's own (n_embd 1024, n_head 12), which supply the new key.
Testing
multi-turn recall and a code-generation turn: all correct.
mtmdbuilds clean with no new warnings.The image path has not been exercised on a real image; it is verified structurally and by
build only.