Skip to content

fix: restore Kimi Linear CUDA inference - #2296

Open
inureyes wants to merge 6 commits into
mainfrom
fix/issue-2292-kimi-cuda-inference
Open

inureyes wants to merge 6 commits into
mainfrom
fix/issue-2292-kimi-cuda-inference

Conversation

@inureyes

@inureyes inureyes commented Oct 13, 2026 •

Copy link
Copy Markdown
Member

Problem

Kimi Linear treated the repository's additive causal mask (0 for 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

  • Cast the causal mask to the MLA positional-score dtype and add it to the scores, matching the mask contract and the Kimi K3/DeepSeek MLA implementations.
  • Add an exact mask-value regression and a loaded three-token MLA prefill regression that requires finite output.
  • Record before/after results for the pinned Kimi Linear checkpoint in bilingual technical reports.

Validation

  • cargo fmt --all -- --check

  • git diff --check

  • Root and surgery CUDA test-fast build 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_cond semantics 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-4bit at 919175d75427f1c10cfc17e7254b53295edddf72: the 21-token prompt returned exactly 4; the 3,411-token prompt returned exactly AMBER-RAVEN-7429 with MLXCEL_PREFILL_CHUNK unset and set to 0. The ordinary MLXCEL_SDPA_DETERMINISTIC=0 short 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

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 inureyes added status:review Under review type:bug Bug fixes, error corrections, or issue resolutions priority:medium Medium priority area:models Model architectures, weights, loading, metadata area:inference Generation, sampling, decoding (incl. speculative, DRY) platform:linux Linux (CUDA / packaging) specific labels Oct 13, 2026
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.
inureyes added a commit that referenced this pull request Oct 14, 2026
Record final CUDA route, combined-tree Kimi inference, and production-source Metal validation for PR #2304. Preserve source and binary provenance, independent #2296 attribution, and unavailable ROCm/M5 and deferred integration coverage.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:inference Generation, sampling, decoding (incl. speculative, DRY) area:models Model architectures, weights, loading, metadata platform:linux Linux (CUDA / packaging) specific priority:medium Medium priority status:review Under review type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix(kimi_linear): diagnose CUDA token-zero degeneration

1 participant