fix: restore Kimi Linear CUDA inference

TECHNICAL_REPORTS/pr-2296-kimi-linear-mla-prefill.en.md

@@ -0,0 +1,69 @@
+# PR #2296: Restore finite Kimi Linear MLA prefill
+
+**Date**: 2026-10-14
+
+**Status**: Implemented and independently reviewed; focused and pinned-checkpoint validation passed
+
+**Risk**: Medium (causal masking in Kimi Linear's full-attention layers)
+
+## Summary
+
+PR #2296 fixes Kimi Linear inference on CUDA. The model previously emitted token ID 0 repeatedly for both a 21-token arithmetic prompt and a 3,411-token retrieval prompt. Kimi's MLA path copied the upstream boolean-mask operation, but mlxcel's shared `create_causal_mask` returns an additive mask. Treating zero and negative infinity as conditions reversed visibility and made every final prefill query softmax an entirely blocked row.
+
+The MLA path now casts the additive mask to the positional-score dtype and adds it. Two regressions pin the mask values and drive a real three-token Kimi MLA prefill through the affected operation. The pinned 48B-A3B 4-bit checkpoint then answers `4` and `AMBER-RAVEN-7429` exactly under the required chunk settings.
+
+## Root cause
+
+`KimiLinearModel::forward` creates a causal mask with `mlxcel_core::utils::create_causal_mask`. Its contract is zero for a visible key and negative infinity for a blocked key. `KimiMLAAttention::forward` passed that array to `where_cond(mask, pe_scores, negative_infinity)` as if it were boolean.
+
+Zero converts to false while negative infinity converts to true. Visible keys were therefore replaced by negative infinity and future keys were retained. The final query row in any multi-token prefill has no future keys, so its whole mask row is zero and all positional scores became negative infinity. Adding the content scores did not restore finite values. Softmax of the all-negative-infinity row produced NaNs, which flowed through the first MLA layer, later residuals and recurrent state, and the final logits. Greedy argmax then returned token ID 0.
+
+The defect was independent of chunking. The shipped chat template produced the expected token IDs, checkpoint packing matched the loader, and the checkpoint's grouped routing configuration is inert because both group counts are one.
+
+## Change
+
+`apply_mla_causal_mask` preserves the repository mask contract. It casts the mask only when its dtype differs from `pe_scores`, then adds it. This matches the existing Kimi K3 and DeepSeek MLA adaptations and works with the cached `[L, L + offset]` key axis through normal broadcasting.
+
+The patch deliberately leaves KDA RMS normalization unchanged. A newer upstream epsilon formula is a separate porting difference, but it does not explain the deterministic all-token-0 failure. The mask-only candidate restores exact checkpoint answers, so including the unrelated numerical change would obscure the demonstrated cause.
+
+## Regression tests
+
+`kimi_mla_applies_the_repository_causal_mask_as_additive_values` applies a three-by-three causal mask to known scores. It requires rows `[score, -inf, -inf]`, `[score, score, -inf]`, and three visible scores. This catches both mask reversal and a missing mask.
+
+`kimi_mla_multi_token_prefill_keeps_the_last_query_row_finite` loads a small zero-weight `KimiMLAAttention`, performs a three-token prefill, and requires every output to remain finite. A single-token fixture would not reach the causal-mask branch.
+
+Both tests pass with the fix. As a controlled negative check, replacing only the helper with the former `where_cond` semantics made both tests fail with exit 101: the first visible score was negative infinity instead of zero, and the real MLA fixture produced a non-finite final row. The fixed source was then restored byte for byte before the final build.
+
+## Pinned-checkpoint validation
+
+The acceptance checkpoint is `mlx-community/Kimi-Linear-48B-A3B-Instruct-4bit` at revision `919175d75427f1c10cfc17e7254b53295edddf72`, running on an NVIDIA GB10 with CUDA compute capability 12.1. Greedy generation used temperature zero and seed 1749.
+
+The baseline at `190b83c3`, binary SHA-256 `592c33a141ccd7a497732e132ac6da302661e22400e27fbcdb2cb3ca858ae8ff`, emitted token ID 0 for all 64 generated tokens on the long prompt under both default chunking and `MLXCEL_PREFILL_CHUNK=0`. The decoded output was 64 exclamation marks. The 21-token arithmetic prompt also emitted token ID 0 and decoded to `!`. PR #2293's chunk-capability change produced the same invalid values, ruling out the no-ID cache reset as the numerical cause.
+
+The first mask-only candidate used base `c9e062af`, source SHA-256 `55c28b20dfd54398e8f6603149e8affbeb6883904c6ee012b277e12d370b429e`, and binary SHA-256 `2b2402c0316e9c6d937fa32998c360dba5214dcd7169fc2d9153e6af19472817`. It produced:
+
+- Short prompt, 21 input tokens, `MLXCEL_PREFILL_CHUNK=0`: `4`, token ID `[19]`.
+- Long prompt, 3,411 input tokens, default chunking: `AMBER-RAVEN-7429`, token IDs `[2102, 12169, 16134, 8858, 1254, 12, 51413, 24]`.
+- Long prompt, 3,411 input tokens, `MLXCEL_PREFILL_CHUNK=0`: the same exact answer and token IDs.
+
+The final branch merged main commit `398f9ba8`, which contains PR #2293, at head `7c3f697706c55d79a54b099b0728892e10ad2115`. Its binary SHA-256 was `b2cb19e07aa833d90e25c039daf13ac43190bf67ce4cf0074c6e074238c40de0`. The three required runs repeated the exact initial answers and token IDs. Two additional runs set `MLXCEL_SDPA_DETERMINISTIC=0`; the short prompt still returned `4` and the default long prompt still returned `AMBER-RAVEN-7429`, proving the repair does not depend on the comparison harness's deterministic-SDPA setting.
+
+## Verification
+
+- `cargo test --locked -p mlxcel -p mlxcel-surgery --lib --test chat_template_kwargs --profile test-fast --features cuda -j4 --no-run`: passed in 5 minutes 58 seconds.
+- Focused final regressions: 3 passed, 0 failed. They include both new MLA mask tests and PR #2293's loaded-model chunk-capability test.
+- Shared issue contracts: 123 Python checks and 49 WebUI checks passed with exit 0.
+- Formatting and patch whitespace checks passed after the main merge.
+- Independent correctness and security review: clear, with no finding in scope.
+
+## Scope
+
+This PR repairs Kimi Linear MLA masking and closes #2292. PR #2293 separately disables no-ID chunked prefill for Kimi Linear because its model-owned recurrent state cannot safely restart between chunks. The final branch contains both changes, while the tests and before/after evidence keep their causes distinct.
+
+## Final integration with mlxcelverse and the updated MLX pin
+
+Integrated main `adf0e589`, including #2304's MLX `243bdf09` pin, mlxcelverse layout, and rustls 0.23.46, without source conflicts. The actual-inference validation used staged source tree `2e4f5211f42b21973d097a00ed60a7768d7f05ba` and CLI SHA-256 `45753012854036e8b44d1adc622a6eba8bc9d15b601e098d340769581cb1a064`. This is a fresh build of the final combined source, not a reused pre-integration model result.
+
+On NVIDIA GB10, `mlx-community/Kimi-Linear-48B-A3B-Instruct-4bit` at revision `919175d75427f1c10cfc17e7254b53295edddf72` passes all five exact-answer controls: the short prompt returns `4`, and the 3,411-token long prompt returns `AMBER-RAVEN-7429` with default prefill, explicit single-chunk prefill, and ordinary SDPA. The short prompt also passes ordinary SDPA. A separate ordinary-SDPA natural-language smoke produces 58 tokens in two coherent sentences explaining daytime Rayleigh scattering and the longer atmospheric path at sunset. There is no repeated token-0 collapse. No performance claim is made from these correctness runs.
+
+All ten `models::kimi_linear` CUDA tests pass on the rebuilt final source, including the additive-mask, finite-last-query-row, and loaded-model chunk-capability regressions. Formatting and the diff check pass.

TECHNICAL_REPORTS/pr-2296-kimi-linear-mla-prefill.ko.md

@@ -0,0 +1,69 @@
+# PR #2296: Kimi Linear MLA prefill을 유한하게 복원하기
+
+**날짜**: 2026-10-14
+
+**상태**: 구현 및 독립 리뷰 완료, 집중 테스트와 고정 체크포인트 검증 통과
+
+**위험도**: 중간 (Kimi Linear full-attention 계층의 causal masking 변경)
+
+## 요약
+
+PR #2296은 CUDA에서 Kimi Linear 추론을 복구합니다. 기존 모델은 21토큰 산술 프롬프트와 3,411토큰 검색 프롬프트 모두에서 토큰 ID 0을 반복했습니다. Kimi MLA 경로는 upstream의 불리언 mask 연산을 옮겼지만 mlxcel의 공용 `create_causal_mask`는 가산형 mask를 반환합니다. 0과 음의 무한대를 조건으로 해석하면서 가시성이 뒤집혔고, 모든 prefill의 마지막 query가 완전히 차단된 행을 softmax하게 됐습니다.
+
+이제 MLA 경로는 가산형 mask를 positional score dtype으로 캐스팅한 뒤 더합니다. 회귀 테스트 두 개가 mask 값 자체를 고정하고 실제 3토큰 Kimi MLA prefill을 문제 연산까지 통과시킵니다. 고정한 48B-A3B 4비트 체크포인트는 요구한 chunk 설정에서 `4`와 `AMBER-RAVEN-7429`를 정확히 답합니다.
+
+## 근본 원인
+
+`KimiLinearModel::forward`는 `mlxcel_core::utils::create_causal_mask`로 causal mask를 만듭니다. 이 함수의 계약은 볼 수 있는 key에 0, 차단한 key에 음의 무한대를 놓는 것입니다. `KimiMLAAttention::forward`는 이 배열을 불리언처럼 `where_cond(mask, pe_scores, negative_infinity)`에 넘겼습니다.
+
+0은 false로, 음의 무한대는 true로 변환됩니다. 따라서 볼 수 있는 key는 음의 무한대로 바뀌고 미래 key는 남았습니다. 여러 토큰을 prefill할 때 마지막 query 행에는 미래 key가 없으므로 mask 행 전체가 0이고 positional score 전체가 음의 무한대가 됐습니다. Content score를 더해도 유한한 값이 복구되지 않습니다. 전부 음의 무한대인 행의 softmax가 NaN을 만들고, 첫 MLA 계층부터 이후 residual, recurrent state, 최종 logits로 퍼졌습니다. Greedy argmax는 그 결과 토큰 ID 0을 반환했습니다.
+
+이 결함은 chunking과 무관합니다. 배포된 채팅 템플릿은 예상한 token ID를 만들었고, 체크포인트 packing은 loader와 일치했으며, 이 체크포인트는 두 expert group 값이 모두 1이라 grouped routing 차이도 발화하지 않습니다.
+
+## 변경
+
+`apply_mla_causal_mask`가 저장소의 mask 계약을 보존합니다. Mask dtype이 `pe_scores`와 다를 때만 캐스팅하고 두 배열을 더합니다. Kimi K3와 DeepSeek MLA의 기존 변환 방식과 같고, 캐시를 포함한 `[L, L + offset]` key 축도 일반 broadcast로 처리합니다.
+
+KDA RMS 정규화는 의도적으로 바꾸지 않았습니다. 최신 upstream의 epsilon 식은 별도의 포팅 차이지만 결정적인 ID 0 반복을 설명하지 않습니다. Mask 하나만 고친 후보가 체크포인트의 정확 답변을 복구했으므로 관계없는 수치 변경까지 넣으면 입증한 원인이 흐려집니다.
+
+## 회귀 테스트
+
+`kimi_mla_applies_the_repository_causal_mask_as_additive_values`는 알려진 score에 3x3 causal mask를 적용합니다. 결과 행이 `[score, -inf, -inf]`, `[score, score, -inf]`, 세 개의 visible score가 되도록 요구합니다. Mask 반전과 mask 누락을 모두 잡습니다.
+
+`kimi_mla_multi_token_prefill_keeps_the_last_query_row_finite`는 작은 zero-weight `KimiMLAAttention`을 적재하고 3토큰 prefill을 수행해 모든 출력이 유한하도록 요구합니다. 1토큰 fixture는 causal-mask 분기에 들어가지 않으므로 이 결함을 재현할 수 없습니다.
+
+수정 뒤 두 테스트가 통과했습니다. 통제한 음성 대조에서는 helper 하나만 이전 `where_cond` 방식으로 바꾸자 두 테스트가 exit 101로 모두 실패했습니다. 첫 visible score가 0 대신 음의 무한대였고, 실제 MLA fixture의 마지막 행이 유한하지 않았습니다. 그 뒤 수정 소스를 바이트 단위로 복원하고 최종 빌드를 수행했습니다.
+
+## 고정 체크포인트 검증
+
+검증 체크포인트는 `mlx-community/Kimi-Linear-48B-A3B-Instruct-4bit` revision `919175d75427f1c10cfc17e7254b53295edddf72`입니다. CUDA compute capability 12.1인 NVIDIA GB10에서 실행했고, greedy generation은 temperature 0과 seed 1749를 사용했습니다.
+
+`190b83c3` 기준선의 binary SHA-256은 `592c33a141ccd7a497732e132ac6da302661e22400e27fbcdb2cb3ca858ae8ff`입니다. 장문 프롬프트에서 기본 chunking과 `MLXCEL_PREFILL_CHUNK=0` 모두 생성 토큰 64개가 전부 ID 0이었고, decode 결과는 느낌표 64개였습니다. 21토큰 산술 프롬프트도 토큰 ID 0과 `!`를 냈습니다. PR #2293의 chunk capability 변경 뒤에도 같은 잘못된 값이 나와 no-ID cache reset이 수치 결함의 원인은 아님을 확인했습니다.
+
+첫 mask-only 후보는 base `c9e062af`, source SHA-256 `55c28b20dfd54398e8f6603149e8affbeb6883904c6ee012b277e12d370b429e`, binary SHA-256 `2b2402c0316e9c6d937fa32998c360dba5214dcd7169fc2d9153e6af19472817`을 사용했습니다. 결과는 다음과 같습니다.
+
+- 짧은 프롬프트, 입력 21토큰, `MLXCEL_PREFILL_CHUNK=0`: `4`, token ID `[19]`.
+- 장문 프롬프트, 입력 3,411토큰, 기본 chunking: `AMBER-RAVEN-7429`, token ID `[2102, 12169, 16134, 8858, 1254, 12, 51413, 24]`.
+- 장문 프롬프트, 입력 3,411토큰, `MLXCEL_PREFILL_CHUNK=0`: 같은 정확 답변과 token ID.
+
+최종 브랜치는 PR #2293이 들어간 main commit `398f9ba8`을 병합해 head `7c3f697706c55d79a54b099b0728892e10ad2115`가 됐습니다. Binary SHA-256은 `b2cb19e07aa833d90e25c039daf13ac43190bf67ce4cf0074c6e074238c40de0`입니다. 요구한 세 실행은 초기 후보와 같은 정확 답변과 token ID를 다시 냈습니다. `MLXCEL_SDPA_DETERMINISTIC=0`으로 실행한 추가 두 건에서도 짧은 입력은 `4`, 기본 장문 입력은 `AMBER-RAVEN-7429`를 그대로 반환해, 수정이 비교 harness의 deterministic SDPA 설정에 기대지 않음을 확인했습니다.
+
+## 검증
+
+- `cargo test --locked -p mlxcel -p mlxcel-surgery --lib --test chat_template_kwargs --profile test-fast --features cuda -j4 --no-run`이 5분 58초에 통과했습니다.
+- 최종 집중 회귀는 3개 통과, 0개 실패입니다. 새 MLA mask 테스트 두 개와 PR #2293의 loaded-model chunk-capability 테스트를 포함합니다.
+- 공용 이슈 contract는 Python 123개와 WebUI 49개가 exit 0으로 통과했습니다.
+- Main 병합 뒤 포맷과 patch whitespace 검사가 통과했습니다.
+- 독립 correctness 및 security 리뷰에서 이 범위의 추가 지적이 없었습니다.
+
+## 범위
+
+이 PR은 Kimi Linear MLA masking을 복구하고 #2292를 닫습니다. PR #2293은 Kimi Linear의 model-owned recurrent state가 chunk 사이에서 안전하게 재시작할 수 없어 no-ID chunked prefill을 별도로 비활성화합니다. 최종 브랜치는 두 변경을 모두 포함하지만 테스트와 전후 증거는 원인을 분리해 기록합니다.
+
+## mlxcelverse 및 최신 MLX pin과의 최종 통합
+
+#2304의 MLX `243bdf09` pin, mlxcelverse 경로, rustls 0.23.46을 포함한 main `adf0e589`를 소스 충돌 없이 통합했다. 실제 추론 검증의 staged source tree는 `2e4f5211f42b21973d097a00ed60a7768d7f05ba`이며 CLI SHA-256은 `45753012854036e8b44d1adc622a6eba8bc9d15b601e098d340769581cb1a064`다. 통합 이전 결과를 재사용하지 않고 최종 결합 소스를 새로 빌드했다.
+
+NVIDIA GB10에서 revision `919175d75427f1c10cfc17e7254b53295edddf72`의 `mlx-community/Kimi-Linear-48B-A3B-Instruct-4bit`가 정답 비교 5건을 모두 통과했다. 짧은 prompt는 `4`를 반환했고, 3,411-token 긴 prompt는 기본 prefill, 명시적 단일 청크 prefill, 일반 SDPA에서 모두 `AMBER-RAVEN-7429`를 회수했다. 짧은 prompt도 일반 SDPA에서 통과했다. 별도의 일반 SDPA 자연어 검증에서는 낮의 Rayleigh 산란과 일몰의 긴 대기 경로를 설명하는 일관된 두 문장, 58 token을 생성했다. token 0 반복 현상은 없었다. 이 정확성 실행으로 성능을 주장하지 않는다.
+
+최종 소스로 다시 빌드한 `models::kimi_linear` CUDA 테스트 10개가 모두 통과했다. Additive mask, 마지막 query 행의 유한값, 가중치를 적재한 모델의 청크 지원 계약 회귀 테스트를 포함한다. Formatting과 diff 검사도 통과했다.

src/models/kimi_linear.rs

@@ -26,7 +26,6 @@ use crate::models::gated_delta::{gated_delta_update, scaled_fast_rms_norm_no_wei
 use crate::models::switch_layers::SwitchGLU;
 use crate::models::switch_layers::validate_expert_quantization_params;
 use mlxcel_core::cache::{SequenceId, SequenceStateLayout};
-use mlxcel_core::dtype;
 use mlxcel_core::generate::LanguageModel;
 use mlxcel_core::layers::{KVCache, RMSNorm, UnifiedEmbedding, UnifiedLinear};
 use mlxcel_core::utils::{create_causal_mask, silu, stack_arrays};
@@ -493,6 +492,22 @@ struct KimiMLAAttention {
     scale: f32,
 }
 
+/// Apply the repository's additive attention mask to MLA positional scores.
+///
+/// [`create_causal_mask`] produces zero for visible keys and negative infinity
+/// for blocked keys. Keep that representation additive here instead of using
+/// it as a boolean condition: zero would become `false`, reversing the mask
+/// and leaving the final prefill row entirely blocked.
+fn apply_mla_causal_mask(scores: &MlxArray, mask: &MlxArray) -> UniquePtr<MlxArray> {
+    let scores_dtype = mlxcel_core::array_dtype(scores);
+    let mask = if mlxcel_core::array_dtype(mask) == scores_dtype {
+        mlxcel_core::copy(mask)
+    } else {
+        mlxcel_core::astype(mask, scores_dtype)
+    };
+    mlxcel_core::add(scores, &mask)
+}
+
 impl KimiMLAAttention {
     fn forward(
         &self,
@@ -547,10 +562,9 @@ impl KimiMLAAttention {
         let k_pe_t = mlxcel_core::swap_axes(&cached_k_pe, -1, -2);
         let mut pe_scores = mlxcel_core::matmul(&q_pe_scaled, &k_pe_t);
 
-        // Apply causal mask to pe_scores
+        // Apply the additive causal mask to positional scores.
         if let Some(m) = mask {
-            let neg_inf = mlxcel_core::full_f32(&[1], f32::NEG_INFINITY, dtype::FLOAT16);
-            pe_scores = mlxcel_core::where_cond(m, &pe_scores, &neg_inf);
+            pe_scores = apply_mla_causal_mask(&pe_scores, m);
         }
 
         // Attention computation (different paths for generation vs prefill)