From 21e02db0b137956cbab631b5ae44b4dc096d8d73 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 13 Oct 2026 23:26:05 +0900 Subject: [PATCH 1/2] fix: opt Kimi Linear out of no-ID chunked prefill Kimi Linear resets model-owned Delta and MLA state when a multi-token forward carries no sequence ID. Advertise the conservative single-pass policy so cache-less trait callers cannot silently discard earlier prompt pieces while preserving the independent-prompt reset contract. Add a tiny model loaded through `KimiLinearModel::from_weights` and assert the capability through `LanguageModel`. The focused CUDA test, formatting, and diff checks pass. Real checkpoint runs remain degenerate with token ID 0 and are tracked separately in #2292. Refs #1749 --- src/models/kimi_linear.rs | 23 +++++++++++++++++++++++ src/models/kimi_linear_tests.rs | 30 ++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/src/models/kimi_linear.rs b/src/models/kimi_linear.rs index 5b054cb87..3a762c380 100644 --- a/src/models/kimi_linear.rs +++ b/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 } diff --git a/src/models/kimi_linear_tests.rs b/src/models/kimi_linear_tests.rs index 4ac864600..5a18afd5c 100644 --- a/src/models/kimi_linear_tests.rs +++ b/src/models/kimi_linear_tests.rs @@ -374,6 +374,36 @@ fn mixed_kimi_caches() -> Vec { ] } +/// 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); From 1193a84a58dc601bcbff33fc15853411e2161709 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 13 Oct 2026 23:32:40 +0900 Subject: [PATCH 2/2] docs: add technical report for PR #2293 --- ...kimi-linear-chunked-prefill-20261013.en.md | 44 +++++++++++++++++++ ...kimi-linear-chunked-prefill-20261013.ko.md | 44 +++++++++++++++++++ 2 files changed, 88 insertions(+) create mode 100644 TECHNICAL_REPORTS/2293-kimi-linear-chunked-prefill-20261013.en.md create mode 100644 TECHNICAL_REPORTS/2293-kimi-linear-chunked-prefill-20261013.ko.md diff --git a/TECHNICAL_REPORTS/2293-kimi-linear-chunked-prefill-20261013.en.md b/TECHNICAL_REPORTS/2293-kimi-linear-chunked-prefill-20261013.en.md new file mode 100644 index 000000000..d7a5bf14a --- /dev/null +++ b/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 diff --git a/TECHNICAL_REPORTS/2293-kimi-linear-chunked-prefill-20261013.ko.md b/TECHNICAL_REPORTS/2293-kimi-linear-chunked-prefill-20261013.ko.md new file mode 100644 index 000000000..8b00ac7a6 --- /dev/null +++ b/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