fix(kimi_linear): chunked prefill discards the recurrent state

src/models/kimi_linear.rs の LanguageModel 実装から始め、prefill の機能に関するメソッドを src/models/kimi_k3.rs:2011-2033 と比較します。kimi_linear_tests.rs と既存の kimi_k3_opts_out_of_chunked_prefill テストを読み、一覧にある cargo tests を実行します。LanguageModel のパスが Kimi Linear で chunked prefill を使用しなくなり、回帰テストと検証コマンドが成功すれば完了です。

Problem / Background

KimiLinearModel answers any prompt longer than MLXCEL_PREFILL_CHUNK (default 2048) from its last chunk alone. Everything before the final chunk never reaches the recurrent state, and nothing raises: no error, no warning, no shape mismatch, just a wrong answer that looks like a normal one.

Found while porting Kimi K3 (#1334, PR #1741), which inherited this code shape byte for byte. Chunked prefill itself came from #672 / #674.

Current Behavior

forward_for_sequence at src/models/kimi_linear.rs:1553 opens with a reset keyed on the token count:

let seq_len = mlxcel_core::array_shape(input_ids)[1];
if seq_id.is_none() && seq_len > 1 {
    self.sequence_state.replace_internal(self.make_kimi_caches());
}

That reset is what makes the cache-less LanguageModel::forward entry point (kimi_linear.rs:1607) safe to call for two unrelated prompts in a row: a multi-token forward with no sequence id is read as the start of a new prompt.

LanguageModel::supports_chunked_prefill defaults to true (src/lib/mlxcel-core/src/generate.rs:404) and KimiLinearModel does not override it, so a prompt longer than one chunk is routed through chunked_prefill_last_logits (generate.rs:302, dispatched at generate.rs:1738 and generate.rs:2526). That helper calls model.forward_last_logits(...) once per chunk with no sequence id, and KimiLinearModel does not override forward_last_logits, so each chunk lands on forward and then on the reset above. Every chunk after the first discards the KimiLinearCache::MLA(KVCache) (kv_latent, k_pe) pair and the KimiLinearCache::Delta(KimiDeltaCache) (q/k/v_conv_state, ssm_state) recurrence its predecessors built (kimi_linear.rs:417-451).

The server path is unaffected: its own prefill_chunk_size (default 512, src/server/cli_input.rs:2252) drives forward_for_sequence with a real SequenceId, which never enters the reset branch. This is a CLI and bench defect.

Proposed Solution

Option 1, minimal and immediate. Override supports_chunked_prefill() -> false on KimiLinearModel's impl LanguageModel, next to the existing supports_padded_prefill / supports_batching overrides at kimi_linear.rs:1581-1587, mirroring KimiK3Model at src/models/kimi_k3.rs:2031. One method, no contract change, correct answers. It gives up the per-chunk prefill memory bound of #672 for this family: transients become one graph over the whole prompt instead of one per chunk.

Option 2, proper. Teach forward_for_sequence to tell "new prompt" from "continuation" by something other than the token count, so chunked prefill can carry state across chunks. This changes the shared ModelOwnedSequenceState contract (src/models/model_owned.rs:22), so it must be checked against every family that resets state inside a forward, not just this one. If option 2 lands, remove Kimi K3's override in the same change so the two families do not drift.

Decision criterion: ship option 1 unless option 2 is already in hand. Option 1 is strictly better than silent truncation and is a one-line revert once option 2 exists.

Scope

In scope: src/models/kimi_linear.rs and src/models/kimi_linear_tests.rs.

Out of scope: fixing the five sibling families in the audit below. They carry the same defect but each needs its own validation checkpoint; file follow-ups rather than widening this issue.

Audit: in-forward state resets on ModelOwnedSequenceState

The defect shape is a replace_internal reached from a forward path, not from make_caches / reset_runtime_state / set_kv_cache_layer_modes, which the generator calls once per prompt and are safe. Of the 42 replace_internal call sites under src, seven are in-forward, all with the identical if seq_id.is_none() && seq_len > 1 guard:

Checked and not affected, because their replace_internal calls are all in make_caches, reset_runtime_state or set_kv_cache_layer_modes: qwen3_5.rs (3329, 3344), mamba.rs, mamba2.rs, jamba.rs, nemotron_h.rs, falcon_h1.rs, granitemoehybrid.rs, plamo2.rs, lfm2.rs, deepseek_v4.rs, llama4.rs, gemma3.rs, gemma4.rs, muse_glimmer.rs, inkling/runtime.rs, vision/unlimited_ocr.rs (which opts out separately at unlimited_ocr.rs:275).

Implementation Notes

• Reuse: for option 1, copy the shape and the doc-comment reasoning from kimi_k3.rs:2011-2033. The two families share the failure mode; do not reword it into a second, divergent account.
• Reuse: src/models/kimi_k3_tests.rs:1640 (kimi_k3_opts_out_of_chunked_prefill) is the existing per-model test shape. generate.rs:3107 (supports_chunked_prefill_default_and_override) and generate.rs:3139 (chunked_prefill_matches_single_pass) are the trait-level and equivalence shapes.
• No call-site plumbing: effective_prefill_chunk (generate.rs:285) already ANDs in the model's answer, and LoadedModel forwards the method (src/loaded_model.rs:410), so option 1 reaches the CLI and bench paths on its own.
• Edge case: a prompt exactly MLXCEL_PREFILL_CHUNK long stays single-pass (the gate requires prompt_len > configured), so the bug starts at chunk + 1 tokens. MLXCEL_PREFILL_CHUNK=0 also forces single-pass and is the current user-side workaround.
• Edge case: option 2 must preserve the two-unrelated-prompts-in-a-row property the reset exists to provide. Whatever signal replaces the token count has to be cleared by make_caches / reset_runtime_state, which the generator already calls at the start of each prompt.
• Error handling: neither option adds a panic on the prefill path. The unwrap_or_else(|err| panic!("KimiLinear {err}")) at kimi_linear.rs:1567 stays as-is; it guards the server path, where the scheduler must have called prepare_sequence_state first.

Acceptance Criteria

• A prompt longer than one chunk is no longer answered from its last chunk alone, by option 1 or option 2.
• A test through the LanguageModel interface asserts it: for option 1, that supports_chunked_prefill() is false on a tiny loaded model; for option 2, that a prompt longer than one chunk yields the same last-position logits chunked and unchunked.
• The audit table above is confirmed or corrected in a comment on this issue. Option 1 needs the table re-checked at merge time; option 2 needs a fresh audit of every ModelOwnedSequenceState user.
• If option 2 lands, KimiK3Model::supports_chunked_prefill (kimi_k3.rs:2031) is removed in the same change and kimi_k3_opts_out_of_chunked_prefill is updated or deleted.
• Follow-ups are filed for the five other affected families, referencing this issue.
• Real-checkpoint validation is recorded in the PR body: prompt length, both answers, flag settings.

Verification

cargo test --workspace --profile test-fast --features metal,accelerate kimi_linear
cargo test --profile test-fast -p mlxcel-core generate::tests::supports_chunked_prefill
cargo clippy --workspace --all-targets -- -D warnings
cargo fmt --check

No Kimi Linear conversion is in the benchmark store today (models/ holds kimi-k3-2l-mla-synth, kimi-k3-8l-mxfp4, kimi-k3-tokenizer, kimi-vl-a3b-thinking-4bit), so fetch one first:

./target/release/mlxcel download mlx-community/Kimi-Linear-48B-A3B-Instruct-4bit
./target/release/mlxcel generate -m models/Kimi-Linear-48B-A3B-Instruct-4bit -p "$(cat long_prompt.txt)" -n 64
MLXCEL_PREFILL_CHUNK=0 ./target/release/mlxcel generate -m models/Kimi-Linear-48B-A3B-Instruct-4bit -p "$(cat long_prompt.txt)" -n 64

A pass is the two runs agreeing. Build long_prompt.txt at over 2048 tokens with the answer-bearing fact in its first 2048 tokens and the question at the end, so a truncated prefill cannot guess it; before the fix the two runs disagree.

Technical Considerations

Related: #1334 and PR #1741 (where the defect was found, and where Kimi K3 took option 1), #672 (chunked prefill's memory motivation), #674 (the supports_chunked_prefill policy and its default).