lablup / lablup/mlxcel

fix(kimi_linear): chunked prefill discards the recurrent state

Open
#1,749 0 comments 0 reactions 0 assignees View on GitHub
area:inference area:models priority:medium status:ready type:bug
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## 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:

```rust
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:

| Model | Reset site | `supports_chunked_prefill` | Verdict |
|---|---|---|---|
| `kimi_linear.rs` | 1559 | default `true` | Affected. This issue. |
| `kimi_k3.rs` | 1982 | `false` at 2031 | Already mitigated (PR #1741). |
| `qwen3_next.rs` | 1669 | default `true` | Affected. Follow-up needed. |
| `rwkv7.rs` | 1016 | default `true` | Affected. Follow-up needed. |
| `bailing_moe_linear.rs` | 2509 | default `true` | Affected. Follow-up needed. |
| `recurrent_gemma.rs` | 1039 | default `true` | Affected. Follow-up needed. |
| `afmoe.rs` | 1226 | default `true` | Affected. Follow-up needed. |

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

```bash
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:

```bash
./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).

Contributor guide

Open the contributing guide

Research direction

Start in src/models/kimi_linear.rs at the LanguageModel implementation and compare its prefill capability methods with src/models/kimi_k3.rs:2011-2033. Read kimi_linear_tests.rs and the existing kimi_k3_opts_out_of_chunked_prefill test, then run the listed cargo tests; done means the LanguageModel path no longer uses chunked prefill for Kimi Linear and the regression test and verification commands pass.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
machine-learning, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
74/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.