Add dedicated Kimi-K3 GGUF support - #621
Merged
Merged
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
justinchuby
force-pushed
the
justinchuby-kimi-linear-gguf
branch
from
August 25, 2026 19:04
eb61cca to
9bf181f
Compare
justinchuby
force-pushed
the
justinchuby-kimi-k3-gguf
branch
from
August 25, 2026 19:44
ce26b16 to
dfdd697
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
Adds first-class support for the Kimi-K3 architecture to Mobius’ model/task stack and GGUF importer, enabling exact graph construction + strict GGUF contract validation while keeping generic runtime packaging deferred due to heterogeneous state ABI limitations (#605).
Changes:
- Add
KimiK3Config,KimiK3CausalLMModel, and a dedicated heterogeneous-state taskKimiK3CausalLMTask. - Extend GGUF import pipeline to recognize
kimi-k3, enforce pinned metadata/tensor closure, and support compatible quantized import paths (including fused KV-B splitting). - Add/extend tests and docs to cover build-graph expectations, tensor processors/mapping, and end-to-end GGUF import + runtime behaviors.
Reviewed changes
Copilot reviewed 23 out of 23 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/build_graph_test.py | Extends state-output assertions to cover kimi_k3_attention layers. |
| tests/_test_configs.py | Adds a tiny representative config entry for kimi_k3 graph-build coverage. |
| src/mobius/tasks/_kimi_linear.py | Fixes the #605 issue URL in runtime_support metadata. |
| src/mobius/tasks/_kimi_k3.py | New dedicated heterogeneous-state task defining K3 cache inputs/outputs and metadata. |
| src/mobius/tasks/init.py | Exports/registers KimiK3CausalLMTask under a dedicated task name. |
| src/mobius/models/kimi_k3.py | Implements the Kimi-K3 decoder (KDA + NoPE MLA + AttnRes mixing + SiTU latent MoE) and weight preprocessing. |
| src/mobius/models/kimi_k3_test.py | Unit tests for config extraction, graph contracts, dtype behavior, and weight preprocessing. |
| src/mobius/models/init.py | Exports KimiK3CausalLMModel. |
| src/mobius/integrations/gguf/_tensor_processors.py | Adds K3 tensor post-processing (conv kernel layout, decay inversion, MLA reshapes, residual-score expansion). |
| src/mobius/integrations/gguf/_tensor_processors_test.py | Adds focused tests for K3 tensor processing behavior. |
| src/mobius/integrations/gguf/_tensor_mapping.py | Adds kimi-k3 GGUF→HF tensor mapping + suffix rules. |
| src/mobius/integrations/gguf/_tensor_mapping_test.py | Tests exact mapping coverage + rejection of wrong/legacy suffixes for kimi-k3. |
| src/mobius/integrations/gguf/_kimi_k3_test.py | End-to-end tests for the pinned K3 GGUF contract (float/quantized, replay/reorder, padding invariance, CLI, negative paths). |
| src/mobius/integrations/gguf/_config_mapping.py | Adds K3 GGUF metadata key-map + postprocessor; derives hybrid layout with K3 layer-type tag. |
| src/mobius/integrations/gguf/_builder.py | Validates pinned K3 GGUF tensor/metadata closure; enforces dedicated task; supports fused KV-B quantized splitting. |
| src/mobius/integrations/gguf/_builder_test.py | Extends Kimi Linear GGUF tests (conv validation + generic task override rejection) and updates cohort coverage. |
| src/mobius/integrations/gguf/_arch_registry.py | Promotes kimi-k3 to supported config/tensor/graph with runtime deferred. |
| src/mobius/integrations/gguf/_arch_registry_test.py | Updates supported-architecture census + adds explicit K3 support expectations. |
| src/mobius/_registry.py | Registers kimi_k3 model type and adds a tiny-random reference model ID for testing. |
| src/mobius/_configs/_base.py | Adds new config fields and implements KimiK3Config.from_transformers; tightens Kimi Linear conv-kernel validation. |
| src/mobius/_configs/_base_test.py | Adds regression test for Kimi Linear conv-kernel minimum validation. |
| src/mobius/_configs/init.py | Exports KimiK3Config. |
| docs/api/build_from_gguf.md | Updates supported-architecture census + K3 entry and “remaining hybrid cohort” narrative. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+1156
to
+1179
| required_metadata=( | ||
| "context_length", | ||
| "embedding_length", | ||
| "feed_forward_length", | ||
| "block_count", | ||
| "attention.head_count", | ||
| "attention.head_count_kv", | ||
| "attention.layer_norm_rms_epsilon", | ||
| "attention.q_lora_rank", | ||
| "attention.key_length_mla", | ||
| "attention.value_length_mla", | ||
| "attention.kv_lora_rank", | ||
| "rope.dimension_count", | ||
| "ssm.conv_kernel", | ||
| "kda.head_dim", | ||
| "expert_count", | ||
| "expert_used_count", | ||
| "expert_feed_forward_length", | ||
| "expert_shared_count", | ||
| "expert_gating_func", | ||
| "attn_res.block_size", | ||
| "activation.situ_beta", | ||
| "activation.situ_linear_beta", | ||
| ), |
justinchuby
force-pushed
the
justinchuby-kimi-k3-gguf
branch
from
August 25, 2026 19:53
dfdd697 to
4a797f7
Compare
Implement the pinned Kimi-K3 KDA/NoPE-MLA schedule, attention residual topology, latent MoE, heterogeneous cache ABI, and strict GGUF metadata and tensor validation. Preserve compatible quantized projections, safely split fused MLA weights, and cover numerical parity, runtime state behavior, CLI, roundtrip, and rejection paths. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Preserve valid KDA convolution history across right padding, align latent-MoE routing weights with low-precision expert outputs, close config and GGUF kernel-size validation, and carry the parent Kimi Linear review fixes and model coverage IDs. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Adapt Kimi-K3 quantized import to the fail-closed policy from #609: reject lossy rank-3 MLA reshaping, retain explicit float import, and split fused Q4_0 KV-B with exact packed-row reordering. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Reject generic task overrides that cannot represent heterogeneous Kimi state and keep KDA recurrence masks in FP32 while retaining model-dtype output masking. Add targeted task and mixed-Mul regression coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
justinchuby-kimi-k3-gguf
branch
from
August 25, 2026 20:13
4a797f7 to
4e50827
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
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
kimi-k3GGUF metadata and tensors with strict metadata/tensor/shape/storage closure and malformed-input rejectionReconstruction
Reconstructed after #619 was admin squash-merged. The PR contains only the four-commit Kimi-K3 delta and review follow-ups on live-main base
04b7e3f6f2d9de5d5741aafb8fa6375a18eee693. This preserves #611 graph-driven ORT GenAI decoder configs, #618 ORT GenAI end-to-end CI, #624 generic decoder config migration, and the Kimi Linear, MiniMax, runtime, and fail-closed quantization changes already on main.git range-diffreports all four replayed commits as patch-identical (=).Validation
Real-checkpoint note
The pinned
yujiepan/kimi-k3-tiny-random@a5c86ee03f07f7b141508b0108304a1447fbb345checkpoint was assessed, but macOS cannot run its requiredfla-corekernels and its selective compressed-tensors MXFP4 expert representation is unsupported by the generic HF loader. That format fails explicitly rather than producing an incorrectly quantized graph.