Skip to content

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

Description

@inureyes

Follow-up to #635 / PR #717.

Problem / Background

On the DENSE batched-decode backend, combining --max-kv-size N with a prompt longer than N produces degenerate output: the model emits repeating tokens such as {"{"{" or the function of the function. This was observed on GB10 (CUDA) with llama-3.1-8b-4bit, but the reasoning is not CUDA-specific. It is a decode-orchestration bug on the dense path, not a kernel or platform issue.

Evidence that it is pre-existing and mode-independent (NOT an Int8 KV bug)

On a ~520-token prompt with --max-kv-size 256:

  • --kv-cache-mode int8 --max-kv-size 256: garbage output.
  • --kv-cache-mode fp16 --decode-storage-backend dense --max-kv-size 256: garbage output (identical failure).
  • --kv-cache-mode fp16 on the DEFAULT paged backend + --max-kv-size 256: coherent output.

The paged path is coherent because genuine Fp16 sequences are pool-backed: KVCache::trim_front is a no-op there and capacity is a pool-side concern. Int8 always uses the dense backend (only genuine Fp16 sequences are pool-backed), so Int8 + --max-kv-size always exposes this path. The common factor across the two failing runs is the dense backend plus a front-trim, not the KV quantization mode.

Where the bug is NOT

The KV-cache-layer front-trim itself is proven correct by unit tests added in PR #717:

  • int8_kv_trim_front_then_fetch_matches_untrimmed_tail
  • int8_kv_prefill_grow_trim_then_decode_tracks_fp16

These drive prefill across the buffer-grow boundary, front-trim, then decode appends, and confirm the dequantized visible window tracks the Fp16 reference. So KVCache::trim_front / update_and_fetch are correct in isolation.

Where the bug likely IS

The failure is in the server's dense batched-decode orchestration after a trim_front advances live_start (offset stays monotonic, live_start > 0). Candidate areas to investigate:

  • src/server/batch/scheduler.rs: enforce_max_kv_size_for (~line 2082) and execute_batched_decode (~line 4317).
  • src/lib/mlxcel-core/src/cache.rs: BatchedAttentionMetadata::from_kv_caches rope_offsets / kv_lens usage (~line 5198).
  • src/models/llama3.rs: the dense path in Llama3Attention::forward_split_attention (~lines 272-438) when live_start > 0.

The metadata builder itself looks correct (rope_offsets = monotonic offset, kv_lens = live_len + query_len), so the divergence is likely in how the dense decode consumes that metadata, or in a prefill/decode RoPE-position inconsistency after the trim.

Current mitigation (already shipped)

PR #717 emits a startup warning when --max-kv-size is combined with Int8 (which forces the dense backend), telling operators the cap is unreliable there. This issue tracks the actual fix so --max-kv-size produces correct output on the dense backend (and therefore under Int8).

Proposed Solution

Fix the dense-decode trim path so that after trim_front advances live_start, the batched decode consumes rope_offsets / kv_lens and reads the visible KV window consistently, keeping prefill and decode RoPE positions aligned. Start by reproducing on the dense backend, then narrow to the consumption site (scheduler decode vs. attention forward) versus a RoPE-position inconsistency after trim.

Acceptance Criteria

  • A server integration or scripted E2E test sends a prompt longer than --max-kv-size on the dense backend (both --kv-cache-mode fp16 --decode-storage-backend dense and --kv-cache-mode int8) and asserts coherent, non-degenerate output.
  • The fix makes --max-kv-size correct on the dense decode backend.
  • Update or remove 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 once the path is correct.
  • Update the "Known limitation" paragraph in docs/benchmark_results/cuda-kv-quant-modes-gb10-2026-07-10.md to reflect the fix.

Technical Considerations

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:coremlxcel-core: MLX FFI, primitives, KV cache, layerspriority:mediumMedium prioritystatus:readyReady to be worked ontype:bugBug fixes, error corrections, or issue resolutions

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions