Repository navigation
fix: restore Kimi Linear CUDA inference - #2296
Merged
Merged
Conversation
Kimi Linear receives causal masks as additive 0/-infinity arrays. Interpreting those values as booleans reversed visibility and made the final prefill query softmax an entirely masked row, propagating NaNs into logits and recurrent decode state. Cast each mask to the positional-score dtype and add it to the scores. Cover the exact mask rows and a real three-token MLA prefill whose final output must remain finite.
Document the additive-mask root cause, the controlled old-mask failures, and the pinned checkpoint's before-and-after token outputs in English and Korean. Record the final source and binary identities, exact CUDA build graph, focused regressions, required chunk settings, and ordinary SDPA validation.
inureyes
added
status:review
Add the completed shared contract gate to the bilingual report after all Python and WebUI checks pass.
Bring the Kimi causal-mask repair onto main’s new MLX pin and rebased CUDA, Metal, and ROCm overlays without changing the proven model source. The existing model evidence remains tied to tested source 7c3f697 and binary b2cb19e0; the refreshed dependency graph requires a rebuild and checkpoint validation.
This was referenced Oct 14, 2026
inureyes
added a commit
that referenced
this pull request
Oct 14, 2026
## Summary Advance MLX from `59d600b5` to captured upstream main `243bdf09` (111 commits, 312 files). Reconcile whole-file overlays to retain mlxcel's numerical behavior and runtime controls without overwriting upstream fixes. - Rebase 11 CPU/CUDA/Metal overlays, retain 17 unchanged upstream-backed files and two helpers, and retire two overlays covered by upstream #4514 and #4554. Preserve CUDA chunking, gathered GEMM, dtype controls, and the #1599 Metal `B * M` correction; use upstream #4594 for Xcode 27. - Retarget ROCm while preserving fork `75915908`. Adapt move-only allocation ownership, borrowed-storage donation safety, compiled callbacks, SDPA sinks, gated-delta fallback, floor division, and odd inverse FFT. Add actual-header host ownership regressions; these do not prove HIP completion ordering. - Fix the new CUDA RHS route's device argument and add aligned MXFP4/MXFP8 reference coverage. Retain reviewed macOS filesystem test portability fixes from #2138. ## Validation - Pin/parser (27), overlay tooling (29), ownership, provenance, contract gates, Python tooling (123), and formatting pass. - Earlier CUDA compilation and 94 focused numerical executions pass. After the final overlay retirement, 45 numerical executions and two forced RHS runs pass; fresh PTX proves f16/bf16 × MXFP4/MXFP8 route execution. - [Earlier Metal CI on production source `3aeccc2b`](https://github.com/lablup/mlxcel/actions/runs/38492190049): fmt/Clippy pass; 12,285 tests pass, 390 ignored; 51 FFmpeg video tests pass. - Actual GB10 Kimi 48B-A3B 4-bit inference passes three focused tests, five exact short/long controls, and a coherent 58-token explanation. This uses combined source `53473340` = production `3aeccc2b` + independent #2296 mask fix `1fec677b`; that model fix is not included here. CLI SHA-256: `b52d70388e0420582c4f60f3341d570c2afd736685b37e2c0c6e50ed9eaf5fae`. ROCm compilation/runtime and M5-specific execution remain unavailable. Ordinary PR CI is tracked separately; queued checks are not passes. The requested whole-batch CUDA/Clippy and fresh main Metal integration gate remains deferred until the expanded follow-ups are ready. This PR does not independently close #1599, #2158, or #1986. English/Korean technical reports record provenance and remaining limits. Latest main `59974283` is integrated, including the mlxcelverse layout migration. Resolve three rename/delete or relocation conflicts, update ownership-test paths and ROCm record numbering, and preserve main's ROCm/MoE improvements. Post-merge CUDA compilation, contract/ownership/provenance gates, Python tests, and 112 focused test executions pass, including directly forced RHS kernels. CUDA library suites on the merged production source pass 11,247 tests (170 ignored); full CUDA workspace/all-targets Clippy and CPU-only linking pass on final `4befdbc7`. The CUDA block-float archive audit and six quantization tests pass. Upgrade rustls from 0.23.45 to 0.23.46 to resolve RUSTSEC-2026-0346 through 0349; local and hosted cargo-deny pass. [Final Metal CI at `4befdbc7`](https://github.com/lablup/mlxcel/actions/runs/38583586205) passes fmt, Clippy, 12,315 workspace tests (390 ignored), and 51 FFmpeg video tests. Earlier actual Kimi results retain their original source attribution. Closes #2303
Merge current main and preserve the additive MLA mask correction through the mlxcelverse migration. Record fresh final-source validation rather than relying on the earlier combined-build evidence. Ten Kimi CUDA regressions and six actual 48B checkpoint inference controls pass, including long-context recall and coherent ordinary-SDPA text generation. Formatting and diff checks pass. Refs #2292
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.
Kimi Linear used the additive causal mask as a boolean condition, reversing key visibility and making the final prefill row entirely negative infinity. Softmax produced NaNs and greedy decoding repeatedly emitted token ID 0. Cast the mask to the MLA score dtype and add it to the scores; add exact mask-value and loaded-prefill finite-output regressions.
Integrate main
c0ae26b2, including the mlxcelverse migration, MLX pin243bdf09, and rustls security update. Bilingual reports preserve the original failure evidence and the fresh final-source validation.Validation:
--lib --tests, warnings denied), formatting, and diff checks pass. Earlier negative-control testing of the former mask semantics failed both new MLA regressions as intended.mlx-community/Kimi-Linear-48B-A3B-Instruct-4bit, revision919175d75427f1c10cfc17e7254b53295edddf72: five exact-answer controls pass. The 21-token prompt returns4; the 3,411-token prompt returnsAMBER-RAVEN-7429with default and single-chunk prefill. Ordinary SDPA also passes the short and long cases.2e4f5211f42b21973d097a00ed60a7768d7f05ba; CLI SHA-256:45753012854036e8b44d1adc622a6eba8bc9d15b601e098d340769581cb1a064. Subsequent changes add validation reports only. Earlier shared contracts passed 123 Python and 49 WebUI checks.Review found no remaining blocking implementation, security, or performance issue. No performance claim is made from these correctness runs.
Closes #2292