Add dedicated Kimi Linear GGUF support - #619
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-minimax-01-gguf
branch
from
August 25, 2026 18:50
3f5fdf1 to
d4957ef
Compare
Implement the pinned Kimi Linear hybrid KDA/MLA and dense/MoE graph with an explicit heterogeneous cache ABI. Add strict GGUF metadata and tensor closure validation, compatible quantization handling, execution coverage, and support documentation while deferring generic OGA packaging until its state schema can represent the model exactly. 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 state across padded prefills, route native K/V-B tensors through explicit reshape-aware requantization, and register the public Kimi config and coverage policy. 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-linear-gguf
branch
from
August 25, 2026 19:04
eb61cca to
9bf181f
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds first-class support for the Kimi Linear architecture across Mobius’ model/config/task stack and the GGUF import pipeline, enabling pinned-format GGUF → ONNX builds with strict tensor/metadata closure validation while keeping ORT GenAI runtime packaging explicitly deferred due to heterogeneous state ABI limitations (#605).
Changes:
- Add Kimi Linear config extraction, model graph (KDA + NoPE-MLA + dense/MoE schedule), and a dedicated heterogeneous-state task/ABI.
- Extend GGUF integration with Kimi Linear architecture registration, exact tensor-name mapping, tensor post-processing, strict contract validation, and quantized import handling (including reshape + requantize paths).
- Add broad test coverage for graph building, config extraction, tensor processing/mapping, GGUF import/runtime state behavior, CLI build, and docs/support census updates.
Reviewed changes
Copilot reviewed 24 out of 24 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/model_coverage_test.py | Adds Kimi Linear entry to model coverage registry text. |
| tests/build_graph_test.py | Asserts Kimi Linear per-layer present-state outputs for graph builds. |
| tests/_test_configs.py | Adds tiny Kimi Linear test config for graph construction tests. |
| src/mobius/tasks/_kimi_linear.py | Introduces a dedicated Kimi Linear task with heterogeneous state inputs/outputs and function registration. |
| src/mobius/tasks/init.py | Exports/registers the new Kimi Linear task and task name. |
| src/mobius/models/kimi_linear.py | Adds Kimi Linear model implementation (decoder stack + MoE + weight preprocessing). |
| src/mobius/models/kimi_linear_test.py | Adds unit tests for Kimi Linear weight preprocessing and routing behavior. |
| src/mobius/models/init.py | Exports the new Kimi Linear model. |
| src/mobius/integrations/gguf/_tensor_processors.py | Adds Kimi Linear tensor post-processing for GGUF → HF tensor layout restoration. |
| src/mobius/integrations/gguf/_tensor_processors_test.py | Tests Kimi Linear tensor post-processing and malformed-input rejection. |
| src/mobius/integrations/gguf/_tensor_mapping.py | Adds Kimi Linear recurrent suffix closure and GGUF→HF tensor-name mapping. |
| src/mobius/integrations/gguf/_tensor_mapping_test.py | Tests Kimi Linear tensor mapping correctness and rejection of wrong suffixes. |
| src/mobius/integrations/gguf/_config_mapping.py | Adds GGUF→config extraction and postprocessing for Kimi Linear plus hybrid-layout derivation support. |
| src/mobius/integrations/gguf/_builder.py | Adds strict Kimi Linear GGUF contract validation and quantized reshape/import handling. |
| src/mobius/integrations/gguf/_builder_test.py | Adds end-to-end Kimi Linear GGUF build/runtime/state/roundtrip/CLI and contract-failure tests. |
| src/mobius/integrations/gguf/_arch_registry.py | Marks kimi-linear as importable with concrete config/tensor/graph support (runtime deferred). |
| src/mobius/integrations/gguf/_arch_registry_test.py | Updates architecture support census expectations and adds Kimi Linear spec assertions. |
| src/mobius/components/_kimi_linear.py | Adds Kimi Linear KDA recurrence + MLA attention components used by the model graph. |
| src/mobius/components/init.py | Exports the new Kimi Linear components. |
| src/mobius/_registry.py | Registers kimi_linear model_type and adds a default HF model id entry. |
| src/mobius/_configs/_base.py | Adds KimiLinearConfig config extraction/validation from Transformers configs. |
| src/mobius/_configs/_base_test.py | Adds unit tests for Kimi Linear config extraction and profile rejection. |
| src/mobius/_configs/init.py | Exports KimiLinearConfig. |
| docs/api/build_from_gguf.md | Updates GGUF support census and documents Kimi Linear’s supported/deferred status. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+135
to
+139
| model.metadata_props["mobius.runtime_support"] = ( | ||
| "Deferred: released generic OGA decoder cache schemas do not represent " | ||
| "Kimi Linear's heterogeneous state ABI; tracked by " | ||
| "https://github.com/onnxruntime/mobius#605" | ||
| ) |
Comment on lines
+1550
to
+1565
| num_heads = int(vars(linear).get("num_heads", base.num_attention_heads)) | ||
| head_dim = int(linear.head_dim) | ||
| if num_heads != base.num_attention_heads: | ||
| raise ValueError("Kimi Linear KDA and MLA head counts must match") | ||
| fields = _shallow_fields(base) | ||
| fields.update( | ||
| model_type="kimi_linear", | ||
| layer_types=[ | ||
| "kimi_linear_attention" if i + 1 in set(kda_layers) else "full_attention" | ||
| for i in range(base.num_hidden_layers) | ||
| ], | ||
| linear_num_key_heads=num_heads, | ||
| linear_num_value_heads=num_heads, | ||
| linear_key_head_dim=head_dim, | ||
| linear_value_head_dim=head_dim, | ||
| linear_conv_kernel_dim=int(linear.short_conv_kernel_size), |
Comment on lines
+800
to
+824
| if ( | ||
| min( | ||
| layers, | ||
| hidden, | ||
| dense_intermediate, | ||
| heads, | ||
| qk_dim, | ||
| value_dim, | ||
| kv_rank, | ||
| extra_dim, | ||
| kda_dim, | ||
| conv, | ||
| experts, | ||
| top_k, | ||
| expert_intermediate, | ||
| ) | ||
| <= 0 | ||
| or qk_dim <= extra_dim | ||
| or top_k > experts | ||
| or shared != 1 | ||
| or dense_layers != 1 | ||
| or not math.isfinite(routed_scale) | ||
| or routed_scale <= 0 | ||
| or not math.isfinite(epsilon) | ||
| or epsilon <= 0 |
justinchuby
added a commit
that referenced
this pull request
Aug 25, 2026
## Summary - add a dedicated Kimi-K3 config, model, and heterogeneous-state task with exact KDA/NoPE gated-MLA scheduling, AttnRes mixing, SiTU latent MoE, shared experts, and an untied output head - import pinned llama.cpp `kimi-k3` GGUF metadata and tensors with strict metadata/tensor/shape/storage closure and malformed-input rejection - preserve lossless quantization routes under #609: separate rank-3 MLA projections fail closed unless explicitly dequantized, while fused Q4_0 KV-B is split by exact packed-row reordering across weights, scales, and zero-points - preserve valid KDA convolution history across padding, validate kernel/state/task contracts, and cover replay, reorder, float/quantized import, roundtrip, CLI, and negative paths - keep generic OGA runtime packaging truthfully deferred under #605 because released cache schemas cannot represent the heterogeneous state ABI ## Reconstruction 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-diff` reports all four replayed commits as patch-identical (`=`). ## Validation - Kimi-K3 model and GGUF tests: 30 passed - focused builder/build-graph Kimi-K3 checks: 4 passed, 1900 deselected - affected generic ORT config/runtime/E2E tests: 250 passed - affected OGA metadata tests: 94 passed, 1 skipped - broad non-integration suite: 7765 passed, 56 skipped, 1 subtest passed - repository lintrunner passed - final GPT-5.6 Sol medium review reported no findings ## Real-checkpoint note The pinned `yujiepan/kimi-k3-tiny-random@a5c86ee03f07f7b141508b0108304a1447fbb345` checkpoint was assessed, but macOS cannot run its required `fla-core` kernels 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. --------- Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 25, 2026
## Summary - clean up still-valid typing, documentation, test robustness, error-message, maintainability, and performance findings left by the GGUF PR stack - preserve the behavioral fixes merged in #625-#632 - hash reused multi-GB GGUF sources once, at the final pre-publication integrity gate, while retaining cheap identity checks around staging Exact base: `2db9d33debdc254d879a51b14434c9a81c230f4f` Exact head: `4541ed2bc9d2ab4227484510b4a85b2d9113eb25` ## Reconstructed original 17-item low-priority tranche The persisted audit retained only the totals, so this list was reconstructed from the live threads and current source. All but the already-fixed #596 comment are addressed in this PR. | PR | Comment | Disposition | |---|---:|---| | #550 | 3837144656 | Implemented: correct tuple return annotation | | #552 | 3837389183 | Implemented: multichannel waveform shape docs | | #559 | 3837716261 | Implemented: name-based cache assertions | | #573 | 3854308350 | Implemented: one final GGUF hash, with integrity regression coverage | | #574 | 3854345521 | Implemented: `TensorRole \| None` typing | | #574 | 3854345597 | Implemented: removed obsolete verdict filtering | | #577 | 3854472859 | Implemented: documented SSM sequence length | | #578 | 3843112613 | Implemented: documented F64 passthrough | | #579 | 3854518332 | Implemented: generalized fused-projection error | | #580 | 3854576300 | Implemented: removed brittle node counts | | #583 | 3854681386 | Implemented: documented conditional draft outputs | | #587 | 3854816872 | Implemented: corrected MTP output contract docs | | #596 | 3846110059 | Already fixed on base: unambiguous GQA bias comment | | #600 | 3855174281 | Implemented: metadata-count-only MTP error | | #607 | 3855776048 | Implemented: stable route-field assertions | | #607 | 3855776086 | Implemented: public tensor iterator | | #609 | 3856486136 | Implemented: fail-closed LM-head comment | ## Current unresolved-thread disposition This covers all 48 Copilot threads returned by the reproducible #600-#630 query. The one human #623 thread is excluded. | PR | Comment | Current-main disposition and evidence | |---|---:|---| | #600 | 3855174177 | Already fixed by #629: package cycle and reserved-sidecar validation | | #600 | 3855174238 | Already fixed by #629: explicit MTP sidecar naming/loading | | #600 | 3855174281 | Implemented here: error no longer invents an observed block count | | #602 | 3848155961 | Outside exact stack; already fixed: top-level `expert_dtype` is classified before early return | | #602 | 3848155983 | Outside exact stack; still-valid behavioral block-quant validation, unchanged | | #602 | 3848156000 | Outside exact stack; still-valid truncated-read behavioral finding, unchanged | | #602 | 3848156022 | Outside exact stack; still-valid descriptor byte/dtype validation, unchanged | | #602 | 3848156040 | Outside exact stack; still-valid expert-bank payload validation, unchanged | | #603 | 3855249765 | Already fixed by #630: runtime preflight preserves shard sets | | #603 | 3855249840 | Already fixed by #630: success output follows durable runtime publication | | #604 | 3855343082 | Implemented here: graph-only MTP persistence distinguished from runtime rejection | | #604 | 3855343131 | Already fixed by #630: runtime success messages are atomic | | #607 | 3855776001 | Implemented here: missing generation golden skips before provenance read | | #607 | 3855776048 | Implemented here: only stable route fields are asserted | | #607 | 3855776086 | Implemented here: tensor count uses `tensor_items_raw()` | | #608 | 3856079371 | Still-valid behavioral cache-symlink containment finding; unchanged | | #608 | 3856079415 | Still-valid behavioral lowercase-digest validation finding; unchanged | | #609 | 3856486136 | Implemented here: comment matches value-preserving policy | | #610 | 3855541683 | Already fixed by #628: Falcon bias precedence is explicit | | #610 | 3855541761 | Already fixed by #628: CTRL tiny config exercises projection biases | | #611 | 3856595840 | Implemented here: runtime test resolves the distribution providing the module | | #612 | 3855677383 | Already fixed by #625: supported-version endianness detection | | #612 | 3855677427 | Implemented here: shared `INT64_MAX` sentinel | | #612 | 3855677460 | Implemented here: shared PLaMo2 width inference | | #612 | 3855677486 | Implemented here: accepted PLaMo2 activation spellings are explicit | | #613 | 3855845678 | Implemented here: canonical issue URL | | #613 | 3855845757 | Implemented here: Mamba-1 function-registration docs | | #613 | 3855845806 | Implemented here: test expects the canonical issue URL | | #614 | 3855988545 | Implemented here: removed stale Nemotron-H divergence comments | | #614 | 3855988597 | Already fixed by #628: zero-head geometry raises actionable `ValueError` | | #615 | 3856082290 | Already fixed by #626: dense GraniteHybrid bias closure | | #618 | 3856729330 | Implemented here: required routes filter ORT GenAI evidence | | #618 | 3856729409 | Implemented here: env-selected runtime version is authoritative | | #618 | 3856729490 | Stale/N/A: PR-description-only matrix claim; repository workflow claims one pinned version | | #618 | 3856729563 | Implemented here: schema tail restored to normal indentation | | #619 | 3856342484 | Already fixed on base: Kimi Linear uses `/issues/605` | | #619 | 3856342532 | Already fixed by #628: config rejects convolution kernels below 2 | | #619 | 3856342580 | Already fixed by #628: GGUF contract rejects convolution kernels below 2 | | #620 | 3855717041 | Already fixed by #627: tied LM-head-only checkpoints are retained | | #621 | 3856722324 | Already fixed by #628: Kimi-K3 required metadata is complete | | #623 | 3855931210 | N/A to current main: comment belongs to open, unmerged #623 | | #623 | 3855939302 | N/A to current main: comment belongs to open, unmerged #623 | | #623 | 3855939358 | N/A to current main: comment belongs to open, unmerged #623 | | #623 | 3855939394 | N/A to current main: comment belongs to open, unmerged #623 | | #624 | 3856777152 | Implemented here: runtime compatibility reuses the emitted model type | | #625 | 3857049733 | Implemented here: unsupported header reports both endian candidates | | #629 | 3857313182 | Newer post-audit behavioral sidecar-symlink cleanup finding; unchanged | | #629 | 3857313251 | Newer post-audit cross-platform path-safety finding; unchanged | ## Validation - affected GGUF/package/ORT GenAI/model/schema tests: 1,040 passed - broad non-integration suite: 7,851 passed, 56 skipped, 1 subtest passed - generated GGUF docs checks: 7 passed - initialized `lintrunner`; full lint/format passed - GPT-5.6 Sol medium review: one integrity finding fixed; re-review found no significant issues Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
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
Runtime boundary
Generic OGA packaging remains deferred to #605 because the released decoder schema cannot represent mixed MLA KV and KDA convolution/matrix states exactly.
Validation
Stacked on #616 (
justinchuby-minimax-01-gguf).