fix(cuda/kv): correct Int8 KV per-token scale and document CUDA KV mode status#717
Merged
Conversation
…de status
Validate every KVCacheMode on CUDA (GB10) and fix the one real defect that made Int8 KV unusable there. The single-stream Int8 quantizer floored every per-token scale to a minimum of 1.0 (`maximum(scale, 1.0)`), but per-token KV absmax is almost always well below 127 so the true scale (`absmax / 127`) is far below 1.0. The floor collapsed quantization into round-to-nearest-integer and produced degenerate greedy decodes ("the function of the function of ..."). The guard now substitutes 1.0 only where the scale is exactly zero (all-zero token) and keeps the true scale otherwise, so the max-magnitude coordinate maps to +-127 and uses the full INT8 range. After the fix, Int8 greedy output on llama-3.1-8b-4bit is byte-identical to Fp16 on the probe prompt. The batched server path (`cache/batch_quant.rs`) uses MLX-native affine quantize and was never affected.
Audit result: all six modes (Fp16, Int8, Turbo4Asym, Turbo3Asym, Turbo4, Turbo4Delegated) run on CUDA with no Metal-only abort. Every custom Turbo/Sparse-V kernel is gated to macOS via `kernel_enabled()` and falls through to a plain-MLX graph path on CUDA, so there is nothing to port or hard-gate; the forced-fallback env-var paths were also exercised and are crash-free.
Benchmark (warmed, `mlxcel-bench-decode`, llama-3.1-8b-4bit): Int8 decode vs Fp16 is 0.74x at 2K, 0.53x at 8K, 0.31x at 32K. Int8 is slower at every length and the gap widens with context because the path dequantizes the whole INT8 window to Fp16 each step and then runs standard Fp16 SDPA (no fused INT8-KV kernel cuts the read bandwidth). What Int8 buys is memory: at 32K the MLX peak drops from 10.37 GB to 8.18 GB, matching a halved Fp16 KV cache. The dated status matrix, the 2K/8K/32K table, and the raw CSV are committed; the main KV-cache doc gains a CUDA subsection recommending Fp16 for speed and Int8 only for long-context memory relief. A fused INT8-KV / paged-attention kernel (#634) is the prerequisite for turning the smaller footprint into a decode win.
Server integration: Int8 chunked prefill + decode and prompt-cache donate/adopt were verified end to end (a second prefix-sharing request reported cached_tokens=256 with coherent output). The KV-cache-layer Int8 front-trim is proven correct by new unit tests (prefill across the buffer-grow boundary, trim, then decode appends track the Fp16 reference within one quant step). Separately, `--max-kv-size` combined with the dense batched-decode backend mis-decodes prompts longer than the cap; this is pre-existing and mode-independent (`--kv-cache-mode fp16 --decode-storage-backend dense --max-kv-size N` fails identically), and Int8 always uses the dense backend so it always exposes it. A startup warning now flags the Int8 + `--max-kv-size` combination so operators do not silently get garbage, and the docs record it as a follow-up to fix in the dense-decode orchestration.
Tests: added `int8_kv_fetch_recovers_values_within_one_quant_step`, `int8_kv_trim_front_then_fetch_matches_untrimmed_tail`, and `int8_kv_prefill_grow_trim_then_decode_tracks_fp16`. Verified with `cargo clippy --features cuda --lib --tests -- -D warnings` clean and the cache::tests::int8 / kv_cache_trim_front / int8 detach-adopt suites green.
Refs #635
inureyes
added
type:performance
This was referenced Jul 9, 2026
…w-up #718 Review follow-ups on PR #717: the status note undersold the scale fix (8-bit uniform batched serving maps to KVCacheMode::Int8 and reuses quantize_per_token, so it is also fixed), and the deferred dense-decode trim defect now cites its tracking issue #718 in the doc and the scheduler comment.
cargo fmt --check flagged two lines inside the three new int8 KV-cache regression tests added for issue #635: a for-loop header in int8_kv_fetch_recovers_values_within_one_quant_step and a chained iterator in int8_kv_prefill_grow_trim_then_decode_tracks_fp16. Reformat both with rustfmt scoped to cache.rs only; no other lines in the file change. Validation: - cargo fmt --check -p mlxcel-core and -p mlxcel now both exit 0 - cargo check --features cuda --lib --tests -p mlxcel-core compiles clean - cargo test -p mlxcel-core --lib --features cuda for the three int8 KV regression tests: 3 passed Refs #635
6 tasks
inureyes
added a commit
that referenced
this pull request
Jul 12, 2026
) ## 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 #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 - [x] `cargo check --features metal,accelerate --lib --tests` - [x] `cargo test --release -p mlxcel-core --features metal,accelerate --lib kv_cache_trim_front` (12 passed, incl. 4 new sink-trim tests) - [x] `cargo test --release -p mlxcel-core --features metal,accelerate --lib int8_kv` (3 PR #717 regression tests still pass) - [x] `cargo test --test max_kv_size_dense_e2e --release --features metal,accelerate -- --ignored --nocapture` (passes on the real model; both cases coherent) - [x] `cargo fmt --all -- --check`, scoped `cargo clippy -- -D warnings` on the changed crates, `python3 scripts/ci/check_cross_repo_refs.py` - [x] Full `cargo build --release --features metal,accelerate` and the three real-model scenarios re-run Closes #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
1 participant
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
Validates every
KVCacheModeon CUDA (GB10) and fixes the one real defect that made Int8 KV unusable there: the single-stream Int8 quantizer floored every per-token scale to a minimum of 1.0, collapsing quantization into round-to-nearest-integer and producing degenerate greedy decodes. All six modes now run on CUDA with no Metal-only abort, Int8 greedy output is byte-identical to Fp16 on the probe prompt, and the CUDA behavior/perf is benchmarked and documented.The Int8 bug and fix
quantize_per_tokencomputedscale = absmax / 127(almost always well below 1.0 for KV activations) and then appliedmaximum(scale, 1.0), forcing the scale up to a floor of 1.0. That mapped the max-magnitude coordinate to onlyround(absmax)instead of +-127, wasting the INT8 range and destroying precision. On llama-3.1-8b-4bit this turned a coherent tool-call into"the function of the function of ...". The guard now substitutes 1.0 only where the scale is exactly zero (all-zero token) and keeps the trueabsmax / 127scale otherwise. After the fix Int8 greedy output is byte-identical to Fp16 on the validation prompt. The batched server path (cache/batch_quant.rs) uses MLX-native affinequantizeand was never affected.Audit status matrix (GB10, llama-3.1-8b-4bit)
No mode aborts on CUDA: every custom Turbo/Sparse-V Metal kernel is gated to macOS via
kernel_enabled()and falls through to a plain-MLX graph path, so there is nothing to port or hard-gate. The forced-fallback env-var paths (MLXCEL_TURBO4_ASYM_DEQUANT_SDPA=0,MLXCEL_TURBO4_DELEGATED_DEQUANT_SDPA=0) were also exercised and are crash-free.Int8 vs Fp16 decode benchmark (warmed)
Int8 decode is slower at every length and the gap widens with context, because the path dequantizes the whole INT8 window to Fp16 each step and then runs standard Fp16 SDPA, so the attention kernel still reads a full Fp16 KV stream while the per-step dequant adds work. The 32K peak drops by 2.19 GB, matching a halved Fp16 KV cache. Net: on CUDA, Int8 KV is a memory-capacity lever, not a decode-speed one. A fused INT8-KV / paged-attention kernel (#634) is the prerequisite for a decode win, and is out of scope here.
What changed
src/lib/mlxcel-core/src/cache.rs: fix the Int8 per-token scale zero-guard; add three regression tests (int8_kv_fetch_recovers_values_within_one_quant_step,int8_kv_trim_front_then_fetch_matches_untrimmed_tail,int8_kv_prefill_grow_trim_then_decode_tracks_fp16).src/server/batch/scheduler.rs: startup warning when--max-kv-sizeis combined with Int8 (which forces the dense decode backend where the cap currently mis-decodes; see below).docs/benchmark_results/cuda-kv-quant-modes-gb10-2026-07-10.md: dated status matrix, benchmark table, analysis, CUDA recommendation, server-integration notes.docs/turbo-kv-cache.md: CUDA subsection with the measured Int8 numbers and recommendation.benchmarks/cuda_gb10_issue635_kvquant_2026-07-10.csv: raw benchmark rows.Server integration and a pre-existing limitation
Int8 chunked prefill + decode and prompt-cache donate/adopt were verified end to end (a second prefix-sharing request reported
cached_tokens: 256with coherent output). The KV-cache-layer Int8 front-trim is proven correct by the new unit tests. Separately,--max-kv-sizeon the dense batched-decode backend mis-decodes prompts longer than the cap. This is pre-existing and mode-independent:--kv-cache-mode fp16 --decode-storage-backend dense --max-kv-size Nfails identically, while Fp16 on the default paged backend stays coherent. Int8 always uses the dense backend, so it always exposes this path. Because the KV-cache-layer Int8 trim is correct and Fp16-dense fails the same way, the defect is in the dense batched-decode orchestration after a trim, not in the Int8 cache. The new startup warning prevents silent garbage; a full fix of the dense-decode trim path is left as a follow-up.Test plan
cargo clippy --features cuda --lib --tests -- -D warningscleancargo test --release --features cuda -p mlxcel-core --lib cache::tests::int8(3 passed)cargo test --release --features cuda -p mlxcel-core --lib kv_cache_trim_front(8 passed)cargo test --release --features cuda -p mlxcel-core --lib int8(20 passed, incl. detach/adopt round-trip)Closes #635