Repository navigation
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.
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
This branch has not been deployed
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.
Problem
Kimi Linear treated the repository's additive causal mask (
0for visible keys and negative infinity for blocked keys) as a boolean condition. This reversed key visibility. In every multi-token prefill, the final query row contains only zeros, so the implementation replaced all of its positional scores with negative infinity. Softmax then produced NaNs, which propagated through later hybrid layers and made greedy decoding emit token ID 0 repeatedly.Change
Validation
cargo fmt --all -- --checkgit diff --checkRoot and surgery CUDA
test-fastbuild graph: passed.Focused final regressions: 3 passed, including the two MLA mask tests and Kimi's loaded-model chunk capability test. Replacing only the mask helper with the former
where_condsemantics made both MLA tests fail as expected.Shared issue contracts: 123 Python checks and 49 WebUI checks passed with exit 0.
Pinned checkpoint
mlx-community/Kimi-Linear-48B-A3B-Instruct-4bitat919175d75427f1c10cfc17e7254b53295edddf72: the 21-token prompt returned exactly4; the 3,411-token prompt returned exactlyAMBER-RAVEN-7429withMLXCEL_PREFILL_CHUNKunset and set to0. The ordinaryMLXCEL_SDPA_DETERMINISTIC=0short and long paths returned the same exact answers.Additional ordinary-path generation (
MLXCEL_SDPA_DETERMINISTIC=0, prefill override unset) produced a coherent 58-token, two-sentence explanation of daytime blue sky and reddish sunsets; this qualitative smoke test used the same pinned checkpoint and tested binary.Closes #2292