Fix compiled decoding with fixed-capacity KV cache

[aleroot]

Proposed changes

Fixes #406.

This PR fixes incorrect token generation when autoregressive decoding is wrapped in MLX.compile. The failure could cause generated text to enter repeating-token loops even though the same model decoded correctly without compilation.

Root cause

During generation, a model keeps a KV cache containing information from tokens it has already processed. Each new token must be written into the next position in that cache.

The existing KVCacheSimple stores this write position as a Swift Int. When MLX compiles the decode function, it traces the function once and reuses the resulting computation graph. The Swift integer and the array slices derived from it are captured as constants during that trace.

As a result, later decode calls can continue using the position observed during compilation instead of the current position. In simpler terms, the compiled decoder keeps writing to and reading from the wrong place in its memory. This is the root cause of the repeating generation reported in #406.

How this PR fixes it

This PR introduces FixedCapacityKVCache, a cache designed for reusable compiled graphs:

• It allocates the complete cache shape before compilation, keeping tensor shapes stable between tokens.
• It stores the current write position as an MLXArray instead of a Swift integer.
• The position is passed through the compiled graph as both input and output, so it advances on every call.
• New keys and values are written using tensor-based indices rather than Swift array slices.
• Unused cache positions are masked so the model cannot attend to memory that has not been written yet.

A CompiledDecodeSession owns the cache and provides the safe high-level interface. It:

• Performs prompt processing eagerly before compiling the single-token decode step.
• Verifies prompt and token shapes.
• Tracks capacity on the Swift side.
• Prevents an out-of-range cache write before it reaches the compiled graph.

Why this is a separate cache

KVCacheSimple remains the better default for normal, non-compiled generation. It grows as needed and only performs attention over tokens that have actually been written.

A compiled graph benefits from the opposite behavior: fixed shapes and tensor-managed state. FixedCapacityKVCache therefore allocates the requested capacity up front and performs attention over that capacity, masking unused positions. This uses more attention work when the cache is only partially filled, but it makes compiled decoding correct and allows the graph to be reused.

Keeping these implementations separate avoids making ordinary eager generation slower to solve a compilation-specific problem.

Supported models

Support is capability-based through FixedCapacityKVCacheProviding. A model should conform only after its position handling and cache topology have been audited for compiled execution.

This PR initially enables:

• Qwen 2 and Qwen 2.5 text models through Qwen2Model.
• Llama and the Mistral checkpoints implemented by LlamaModel.

Models with sliding-window, recurrent, hybrid, vision-language, or custom cache layouts are not assumed to be compatible.

Testing

The new tests cover:

• Compiled versus eager logit and token parity.
• Qwen2 and Llama model paths.
• Cache position advancement through repeated compiled calls.
• Prompt prefill and multi-token continuation.
• Causal and unused-capacity masking.
• Cache copying, trimming, and serialization.
• Invalid shapes and capacity exhaustion.
• An opt-in regression test using the Qwen2.5-7B checkpoint from KVCacheSimple.offset (Swift Int) breaks MLX.compile() decode — cache write position frozen across steps #406.

The Qwen2.5-7B regression was also run locally for 16 decode steps. Compiled and eager decoding produced matching logits and selected tokens at every step.

Longer-term direction

This issue exposes a broader distinction between ordinary Swift state and state that must change inside a reusable MLX computation graph.

Long term, compiled generation should probably move toward:

• Representing all decode-time state that affects computation as tensors carried through the graph.
• Making compiled-decoding support an explicit model capability rather than assuming every cache and model architecture is compatible.
• Keeping compilation lifecycle, prompt prefill, cache ownership, and capacity validation behind a session-style API.
• Adding compiled-versus-eager parity tests whenever a new model family adopts this capability.
• Exploring more efficient fixed-shape or paged cache strategies later, without weakening the correctness guarantees introduced here.

Checklist

• I have read the [CONTRIBUTING](https://github.com/ml-explore/mlx-swift-lm/blob/main/CONTRIBUTING.md) document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[Augustas11]

Thanks for putting this together. I tested PR head edcda687fb4a2a61ae6028e2ff81a077c14a03e6 in our downstream macprovider spike by pinning mlx-swift-lm to that revision and replacing our bench-only manual compiled decode path with CompiledDecodeSession.

Environment/config:

• macOS arm64, Apple Swift 6.3.3 toolchain
• mlx-swift-lm: PR head edcda687...
• mlx-swift: resolved to 0.31.6
• model: mlx-community/Qwen2.5-7B-Instruct-4bit from local HF cache
• bench: greedy decode, prefill target 128, decode target 32, one warmup plus one measured run

Validation results:

• swift build --product macprovider-cli: pass
• focused downstream guards passed: DecodeBenchFlagTests, KVConversationColdTierTests/testEligibleServeNewCacheProducesKVCacheSimple, ModelRuntimeSwapTests/testServingDecodeRunsThroughBlockingInferenceExecutor
• uncompiled generate(...) vs CompiledDecodeSession parity: pass for both warmup and measured sample
• measured sample emitted 31 token ids on both paths, exact token-id match. The session path produced the same content sequence; the downstream harness needed to apply the same stop-token exclusion convention as generate(...) for EOS 151645.

One measured-run smoke result, not a perf claim: uncompiled prefill/decode 543.0 / 27.9 tok/s; compiled-session prefill/decode 572.6 / 26.0 tok/s.

Only downstream compatibility change needed outside the compiled-path replacement was handling the new throwing LanguageModel.newCache(...) API in our runtime/tests. This looks good for our T2-01 parity spike and addresses the frozen cache-position failure mode we were tracking. Thanks again for the fix.