fix: opt Kimi Linear out of chunked prefill

TECHNICAL_REPORTS/2293-kimi-linear-chunked-prefill-20261013.en.md

@@ -0,0 +1,44 @@
+# Technical report: PR #2293 — Kimi Linear chunked-prefill opt-out
+
+**Date**: 2026-10-13
+
+**Status**: Verified with runtime follow-up
+
+**Risk level**: Medium
+
+## Executive summary
+
+Kimi Linear now advertises `supports_chunked_prefill() == false`. This prevents cache-less `LanguageModel` callers from splitting a prompt across no-sequence-ID forwards that each reset the model-owned Delta and MLA state. A tiny model loaded through the production weight constructor verifies the capability through the trait interface.
+
+## Problem statement
+
+`KimiLinearModel::forward_for_sequence` treats a multi-token forward without a sequence ID as a new prompt. That reset keeps two independent cache-less calls isolated, but a no-ID chunking caller would discard the recurrent state built by every earlier chunk. The shared trait default advertised chunking support despite that contract.
+
+Current DirectEngine prefill passes one stable sequence ID through every piece, so its path does not enter the no-ID reset. The change still defines the conservative trait capability and prevents legacy or alternate cache-less callers from selecting an unsafe split.
+
+## Technical decision
+
+The model opts out instead of changing `ModelOwnedSequenceState`. An explicit continuation mechanism would require a cross-family state-contract redesign and must preserve independent-prompt resets. The opt-out is local and reversible, but it can replace per-chunk transient memory bounds with one graph over a long prompt.
+
+## Change summary
+
+- Added the Kimi Linear `supports_chunked_prefill` override with the no-ID state rationale.
+- Added a zero-layer, tied-embedding fixture that loads through `KimiLinearModel::from_weights` and asserts the policy through `&dyn LanguageModel`.
+- Reconfirmed the five sibling reset/default-capability combinations and filed #2285 through #2289.
+
+## Validation
+
+The CUDA `mlxcel` and `mlxcel-surgery` library test graph builds successfully. The focused loaded-model test passes 1/1 under the GPU lock. Formatting and diff checks pass, and independent implementation and security/performance reviews reported no findings.
+
+Real-checkpoint validation used `mlx-community/Kimi-Linear-48B-A3B-Instruct-4bit` revision `919175d75427f1c10cfc17e7254b53295edddf72`. Baseline main `190b83c3` and the candidate built from main `448935b6` plus the frozen policy diff both emitted token ID 0 for all 64 outputs of a 3,411-token prompt under the default setting and `MLXCEL_PREFILL_CHUNK=0`; both also emitted ID 0 for a 21-token chunk-disabled arithmetic prompt. These invalid matching outputs are not an accuracy or numerical-parity pass. The candidate CLI SHA256 was `c0b0319f0b4f0302e9e5dccf02f4f403453ec3dc9aa7e5e1a4ce625f3f060eb7`; diagnosis and exact-answer acceptance are tracked in #2292.
+
+## Follow-up actions
+
+- #2292 must restore normal Kimi Linear CUDA inference and prove exact short and long answers.
+- #2285 through #2289 cover the audited sibling no-ID contracts without widening this change.
+
+## References
+
+- Issue #1749
+- Issues #2285, #2286, #2287, #2288, #2289, and #2292
+- Issue #672 and PR #1741

TECHNICAL_REPORTS/2293-kimi-linear-chunked-prefill-20261013.ko.md

@@ -0,0 +1,44 @@
+# 기술 보고서: PR #2293 — Kimi Linear chunked-prefill opt-out
+
+**작성일**: 2026-10-13
+
+**상태**: 런타임 후속 조치가 있는 검증 완료
+
+**위험도**: 중간
+
+## 요약
+
+Kimi Linear가 이제 `supports_chunked_prefill() == false`를 명시한다. 이에 따라 cache-less `LanguageModel` 호출자가 prompt를 sequence ID 없는 여러 forward로 나눠 각 호출마다 model-owned Delta 및 MLA state를 초기화하는 경로를 선택하지 않는다. production weight constructor로 불러온 작은 모델이 trait interface를 통해 이 capability를 검증한다.
+
+## 1. 문제 정의
+
+`KimiLinearModel::forward_for_sequence`는 sequence ID가 없는 multi-token forward를 새 prompt로 해석한다. 이 초기화는 독립적인 두 cache-less 호출 사이의 state를 격리하지만, no-ID chunking 호출자는 앞 chunk가 만든 recurrent state를 다음 chunk마다 버리게 된다. 기존 shared trait 기본값은 이 계약과 달리 chunking 지원을 선언했다.
+
+현재 DirectEngine prefill은 모든 piece에 동일한 sequence ID를 전달하므로 no-ID 초기화에 들어가지 않는다. 이번 변경은 보수적인 trait capability를 정의하고 legacy 또는 다른 cache-less 호출자가 안전하지 않은 분할을 선택하지 못하게 한다.
+
+## 3. 기술적 선택과 그 이유
+
+`ModelOwnedSequenceState` 계약을 바꾸지 않고 해당 모델이 opt-out하도록 선택했다. 명시적인 continuation mechanism은 여러 모델 family에 걸친 state 계약 재설계가 필요하며 independent-prompt reset도 보존해야 한다. 이 opt-out은 국소적이고 되돌리기 쉽지만, 긴 prompt에서 chunk별 transient memory bound 대신 전체 prompt를 하나의 graph로 처리할 수 있다.
+
+## 7. 변경 요약
+
+- no-ID state 동작의 근거와 함께 Kimi Linear `supports_chunked_prefill` override를 추가했다.
+- zero-layer 및 tied-embedding fixture를 `KimiLinearModel::from_weights`로 불러와 `&dyn LanguageModel`을 통해 정책을 검증한다.
+- 같은 reset/default-capability 조합을 가진 다섯 sibling을 다시 확인하고 #2285부터 #2289까지 등록했다.
+
+## 검증
+
+CUDA `mlxcel` 및 `mlxcel-surgery` library test graph build가 성공했다. GPU lock 아래 focused loaded-model test 1/1이 통과했다. formatting과 diff check가 통과했으며 독립 implementation review와 security/performance review에서 지적 사항이 없었다.
+
+실제 checkpoint 검증에는 `mlx-community/Kimi-Linear-48B-A3B-Instruct-4bit` revision `919175d75427f1c10cfc17e7254b53295edddf72`를 사용했다. baseline main `190b83c3`와 main `448935b6` 및 동결된 policy diff로 만든 candidate는 모두 3,411-token prompt의 64개 출력에서 기본 설정과 `MLXCEL_PREFILL_CHUNK=0` 양쪽 모두 token ID 0을 냈고, 21-token chunk-disabled 산술 prompt에서도 ID 0을 냈다. 동일하게 잘못된 출력은 accuracy 또는 numerical parity 통과가 아니다. candidate CLI SHA256은 `c0b0319f0b4f0302e9e5dccf02f4f403453ec3dc9aa7e5e1a4ce625f3f060eb7`이며 원인 진단과 exact-answer 승인 조건은 #2292에서 추적한다.
+
+## 8. 후속 조치
+
+- #2292에서 Kimi Linear CUDA 정상 inference를 복구하고 short 및 long prompt의 정확한 답을 증명해야 한다.
+- #2285부터 #2289까지는 이 변경의 범위를 넓히지 않고 감사된 sibling no-ID 계약을 다룬다.
+
+## 참고 자료
+
+- Issue #1749
+- Issues #2285, #2286, #2287, #2288, #2289, #2292
+- Issue #672 및 PR #1741

src/models/kimi_linear.rs

@@ -1585,6 +1585,29 @@ impl LanguageModel for KimiLinearModel {
         false // Padding tokens corrupt DeltaCache recurrent state
     }
 
+    /// Prefill through the cache-less [`LanguageModel::forward`] entry point
+    /// runs in one pass for this family.
+    ///
+    /// [`Self::forward_for_sequence`] treats a multi-token forward carrying no
+    /// sequence id as the start of a prompt and resets the model-owned Delta
+    /// and MLA state, which is what makes that entry point safe to call twice
+    /// for two unrelated prompts. A no-ID caller that split one prompt across
+    /// forwards would send exactly that shape once per chunk, so every chunk
+    /// after the first would discard the recurrent state its predecessors
+    /// built. DirectEngine currently avoids that reset by carrying a stable
+    /// sequence id across prefill pieces, but the trait capability also covers
+    /// callers of the cache-less entry point.
+    ///
+    /// Opting out is the conservative half of the fix. It costs the prefill
+    /// memory bound of issue #672 (per-chunk transients instead of one graph
+    /// over the whole prompt). Carrying no-ID state across chunks needs
+    /// `forward_for_sequence` to distinguish "new prompt" from "continuation"
+    /// by something other than the token count, which changes the shared
+    /// `ModelOwnedSequenceState` contract.
+    fn supports_chunked_prefill(&self) -> bool {
+        false
+    }
+
     fn supports_batching(&self) -> bool {
         false // KimiLinear uses internal mixed caches (MLA + DeltaCache), not compatible with per-sequence KV isolation
     }

src/models/kimi_linear_tests.rs

@@ -374,6 +374,36 @@ fn mixed_kimi_caches() -> Vec<KimiLinearCache> {
     ]
 }
 
+/// KimiLinear resets its model-owned Delta and MLA state whenever a
+/// multi-token forward arrives without a sequence id, which keeps two
+/// unrelated cache-less forwards isolated. A no-ID caller must therefore not
+/// split one prompt across multiple forwards because each piece after the
+/// first would discard the recurrent state built by its predecessors.
+#[test]
+fn kimi_linear_loaded_model_opts_out_of_chunked_prefill() {
+    use mlxcel_core::generate::LanguageModel;
+
+    let mut config = mla_config(16);
+    config.num_hidden_layers = 0;
+
+    let mut weights = WeightMap::new();
+    weights.insert(
+        "model.embed_tokens.weight".to_string(),
+        mlxcel_core::zeros(&[32, 32], mlxcel_core::dtype::FLOAT32),
+    );
+    weights.insert(
+        "model.norm.weight".to_string(),
+        mlxcel_core::zeros(&[32], mlxcel_core::dtype::FLOAT32),
+    );
+
+    let model = KimiLinearModel::from_weights(&weights, &config).expect("tiny model loads");
+    let language_model: &dyn LanguageModel = &model;
+    assert!(
+        !language_model.supports_chunked_prefill(),
+        "a no-ID chunked prefill would discard the recurrent state built by every chunk but the last"
+    );
+}
+
 #[test]
 fn kimi_linear_eos_token_ids_come_from_config_metadata() {
     let config = mla_config(16);