fix(server): pin attention-sink prefix on --max-kv-size dense trim#759
Merged
Conversation
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
`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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
--max-kv-size Nwith 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 256failed identically, while the default paged backend stayed coherent because pool-backed Fp16 sequences maketrim_fronta 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:
offsetstays monotonic (504, 505, 506, ...),live_startadvances by the trim depth,buffer_idx() = offset - live_starttracks the physical length, and RoPE positions stay aligned across prefill and every decode step. The defect was that denseKVCache::trim_frontdropped 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 withRotatingKVCache(max_size=max_kv_size, keep=4), which the--max-kv-sizepath 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: addKVCache::trim_front_keep_sink(n, keep). Withkeep > 0it rearranges the physical buffer into[first keep slots] ++ [most-recent window]and drops thentokens after the sink instead of thenoldest overall (mirroring mlx-lmRotatingKVCache._trim). Exactlyntokens are still removed, solive_startstill advances bynand no other bookkeeping changes; the INT8 scale sidecars are sliced in lockstep via a sharedslice_keep_sinkhelper.trim_front(n)now forwardskeep = 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_forcallstrim_front_keep_sink(excess, MAX_KV_SIZE_SINK_KEEP)whereMAX_KV_SIZE_SINK_KEEP = 4, clamped so the sink never leaves the recent window without room. Removes the Int8 +--max-kv-sizestartup 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 bootsmlxcel-serverand asserts non-degenerate output for bothfp16 + denseandint8with 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. ..."--kv-cache-mode fp16 --max-kv-size 256): unchanged from before the fix.Test plan
cargo check --features metal,accelerate --lib --testscargo 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, scopedcargo clippy -- -D warningson the changed crates,python3 scripts/ci/check_cross_repo_refs.pycargo build --release --features metal,accelerateand the three real-model scenarios re-runCloses #718