Skip to content

fix(server): pin attention-sink prefix on --max-kv-size dense trim#759

Merged
inureyes merged 2 commits into
mainfrom
fix/718-dense-batched-decode-max-kv-size
Jul 12, 2026
Merged

fix(server): pin attention-sink prefix on --max-kv-size dense trim#759
inureyes merged 2 commits into
mainfrom
fix/718-dense-batched-decode-max-kv-size

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

--max-kv-size N with a prompt longer than N produced degenerate, repeating output on the DENSE batched-decode backend (the model echoed the prompt's question in a loop). This was not Int8-specific: --kv-cache-mode fp16 --decode-storage-backend dense --max-kv-size 256 failed identically, while the default paged backend stayed coherent because pool-backed Fp16 sequences make trim_front a no-op. Int8 always runs on the dense backend, so it always exposed the path. Reproduced and fixed locally on Apple Silicon (metal).

Root cause

Attention-sink loss, not a position bug. Instrumenting the decode loop confirmed the runtime bookkeeping was already correct: offset stays monotonic (504, 505, 506, ...), live_start advances by the trim depth, buffer_idx() = offset - live_start tracks the physical length, and RoPE positions stay aligned across prefill and every decode step. The defect was that dense KVCache::trim_front dropped the oldest tokens starting at position 0, discarding the leading BOS/opening tokens a transformer dumps its excess attention onto. Without those sink tokens, decode attention collapses into repetition (the StreamingLLM effect). Upstream mlx-lm avoids this with RotatingKVCache(max_size=max_kv_size, keep=4), which the --max-kv-size path is documented to mirror. The cache-layer unit tests added in PR #717 could not catch this because the cache does not perform RoPE and the tests only checked dequantized values, not the model-level attention-sink behavior.

The fix

  • src/lib/mlxcel-core/src/cache.rs: add KVCache::trim_front_keep_sink(n, keep). With keep > 0 it rearranges the physical buffer into [first keep slots] ++ [most-recent window] and drops the n tokens after the sink instead of the n oldest overall (mirroring mlx-lm RotatingKVCache._trim). Exactly n tokens are still removed, so live_start still advances by n and no other bookkeeping changes; the INT8 scale sidecars are sliced in lockstep via a shared slice_keep_sink helper. trim_front(n) now forwards keep = 0 (byte-identical to the old front trim, so existing tests and the gemma/exaone callers are unaffected).
  • src/server/batch/scheduler.rs: enforce_max_kv_size_for calls trim_front_keep_sink(excess, MAX_KV_SIZE_SINK_KEEP) where MAX_KV_SIZE_SINK_KEEP = 4, clamped so the sink never leaves the recent window without room. Removes the Int8 + --max-kv-size startup warning added in PR fix(cuda/kv): correct Int8 KV per-token scale and document CUDA KV mode status #717 now that the path is correct.
  • docs/benchmark_results/cuda-kv-quant-modes-gb10-2026-07-10.md: updates the "Known limitation" paragraph to "Resolved", with the root cause and fix.
  • tests/max_kv_size_dense_e2e.rs: new gated real-server E2E test that boots mlxcel-server and asserts non-degenerate output for both fp16 + dense and int8 with a prompt longer than the cap.

Real-model validation (llama-3.1-8b-4bit, metal, 504-token prompt, --max-kv-size 256)

Before (both degenerate):

  • fp16 --decode-storage-backend dense: " Mercury is the closest to the Sun, and what is special about the planet Saturn is which planet is the closest to the Sun, and what is special about the planet Saturn is which planet is ..." (infinite question echo)
  • int8: " Mercury is the closest to the Sun, and which planet is the closest to the Sun. The planet is the closest to the Sun, and what is special about the planet. The planet is ..."

After (both coherent, matching the paged control):

  • fp16 --decode-storage-backend dense: " Mercury is the closest planet to the Sun. Saturn is known for its bright and extensive ring system. ..."
  • int8: " Mercury is the closest planet to the Sun. Saturn is known for its bright and extensive ring system. ..."
  • paged control (--kv-cache-mode fp16 --max-kv-size 256): unchanged from before the fix.

Test plan

  • cargo check --features metal,accelerate --lib --tests
  • cargo test --release -p mlxcel-core --features metal,accelerate --lib kv_cache_trim_front (12 passed, incl. 4 new sink-trim tests)
  • cargo test --release -p mlxcel-core --features metal,accelerate --lib int8_kv (3 PR fix(cuda/kv): correct Int8 KV per-token scale and document CUDA KV mode status #717 regression tests still pass)
  • cargo test --test max_kv_size_dense_e2e --release --features metal,accelerate -- --ignored --nocapture (passes on the real model; both cases coherent)
  • cargo fmt --all -- --check, scoped cargo clippy -- -D warnings on the changed crates, python3 scripts/ci/check_cross_repo_refs.py
  • Full cargo build --release --features metal,accelerate and the three real-model scenarios re-run

Closes #718

On the dense batched-decode backend, combining `--max-kv-size N` with a prompt longer than N produced degenerate repeating output (the model echoed the prompt's question in a loop). This was not Int8-specific: `--kv-cache-mode fp16 --decode-storage-backend dense --max-kv-size 256` on a 504-token prompt failed identically, while the default paged backend stayed coherent because pool-backed Fp16 sequences make `trim_front` a no-op. Int8 always runs on the dense backend, so it always exposed the path.

Root cause is attention-sink loss, not a position bug. The runtime bookkeeping was already correct: `offset` stays monotonic, `live_start` advances by the trim depth, `buffer_idx() = offset - live_start` tracks the physical length, and RoPE positions stay aligned (verified by instrumenting offset/live_len across prefill and every decode step). The defect was that dense `KVCache::trim_front` dropped the oldest tokens starting at position 0, discarding the leading BOS/opening tokens a transformer dumps its excess attention onto. Without those sink tokens, decode attention collapses into repetition (the StreamingLLM effect). Upstream mlx-lm avoids this by using `RotatingKVCache(max_size=max_kv_size, keep=4)`, which the `--max-kv-size` path is documented to mirror.

The fix adds `KVCache::trim_front_keep_sink(n, keep)`: with `keep > 0` it rearranges the physical buffer into `[first keep slots] ++ [most-recent window]` and drops the `n` tokens after the sink instead of the `n` oldest overall. Exactly `n` tokens are still removed, so `live_start` still advances by `n` and no other bookkeeping changes; the INT8 scale sidecars are sliced in lockstep. `trim_front(n)` forwards `keep = 0` (byte-identical to the old front trim, so existing tests and other callers are unaffected). `enforce_max_kv_size_for` now calls `trim_front_keep_sink(excess, 4)`. The Int8-plus-max-kv-size startup warning from PR #717 is removed now that the path is correct, and the known-limitation note in the GB10 KV-quant doc is updated.

Validated on llama-3.1-8b-4bit (metal) with a 504-token prompt and `--max-kv-size 256`: both `--kv-cache-mode fp16 --decode-storage-backend dense` and `--kv-cache-mode int8` now answer coherently ("Mercury is the closest planet to the Sun. Saturn is known for its bright and extensive ring system."), and the paged control output is unchanged. Added cache-layer unit tests for the sink-preserving trim (Fp16 and Int8, keep=0 equivalence, prefix stability across repeated trims) and a gated real-server E2E test (`tests/max_kv_size_dense_e2e.rs`) that drives both failing configurations and asserts non-degenerate output.

Closes #718
@inureyes inureyes added type:bug Bug fixes, error corrections, or issue resolutions priority:medium Medium priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers labels Jul 12, 2026
`run_completion` in tests/max_kv_size_dense_e2e.rs killed the spawned mlxcel-server before panicking on a non-200 response or a request error, but not on a malformed-JSON success response: `resp.json().await.expect(...)` could panic while the child process was still running, orphaning it. Fold the JSON parse into the same cleanup-aware match arm so `stop_server` always runs first on every panic path.

Also correct two rustdoc comments left over from the plain trim_front era: `BatchScheduler::enforce_max_kv_size_for` and the `max_kv_size` field docs in scheduler.rs and config.rs said the cap drops the oldest tokens, which no longer matches trim_front_keep_sink pinning an attention-sink prefix and dropping the excess after it.

Validation:
- cargo test -p mlxcel-core --lib kv_cache_trim_front (12 passed)
- cargo test --test max_kv_size_dense_e2e --no-run (compiles; model-gated)
- cargo fmt --all -- --check
- cargo clippy -p mlxcel-core --lib --tests -- -D warnings
- cargo clippy --lib --tests -- -D warnings
- python3 scripts/ci/check_cross_repo_refs.py

Refs #718
@inureyes inureyes added the status:done Completed label Jul 12, 2026
@inureyes
inureyes merged commit abf9d3e into main Jul 12, 2026
5 checks passed
@inureyes
inureyes deleted the fix/718-dense-batched-decode-max-kv-size branch July 12, 2026 19:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers priority:medium Medium priority status:done Completed type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix(server): dense batched-decode + --max-kv-size mis-decodes prompts longer than the cap

1 participant