add TurboQuant KV cache compression

Libraries/MLXLLM/LLMModel.swift

@@ -1,5 +1,6 @@
 // Copyright © 2024 Apple Inc.
 
+import Foundation
 import MLX
 import MLXLMCommon
 
@@ -40,11 +41,15 @@ extension LLMModel {
                 // prefill cannot be interrupted, so apps cannot stop GPU submissions
                 // in time when entering the background. See ml-explore/mlx-swift-examples#230.
                 try Task.checkCancellation()
-                let input = y[.newAxis, ..<prefillStepSize]
-                let output = self(input, cache: cache.isEmpty ? nil : cache, state: state)
-                state = output.state
-                asyncEval(cache)
-                y = y[prefillStepSize...]
+                // Pool per chunk: long prompts run hundreds of chunk forwards
+                // before returning to any autorelease boundary.
+                autoreleasepool {
+                    let input = y[.newAxis, ..<prefillStepSize]
+                    let output = self(input, cache: cache.isEmpty ? nil : cache, state: state)
+                    state = output.state
+                    asyncEval(cache)
+                    y = y[prefillStepSize...]
+                }
             }
 
             // Single sync after the loop to flush any remaining async work.

Libraries/MLXLMCommon/AttentionUtils.swift

@@ -51,7 +51,25 @@ public func attentionWithCacheUpdate(
             mask: mask
         )
     }
-    if let quantizedKVCache = cache as? QuantizedKVCacheProtocol {
+    if let turboCache = cache as? TurboQuantKVCache {
+        let L = queries.dim(2)
+        if L > 1 && !turboCache.isCompressed {
+            // Prefill (L>1) on a raw cache: plain update + standard SDPA, // zero overhead; compression is deferred to the first decode step.
+            let (cachedKeys, cachedValues) = turboCache.update(keys: keys, values: values)
+            return MLXFast.scaledDotProductAttention(
+                queries: queries, keys: cachedKeys, values: cachedValues,
+                scale: scale, mask: mask
+            )
+        }
+        // Decode (L=1) or any call once the cache is compressed (speculative
+        // verify chunks, multi-turn re-prefill): the compressed path. The raw
+        // update() path is invalid after compression, its raw buffers are
+        // gone. First decode call triggers compressRawCache().
+        return turboCache.compressedAttention(
+            queries: queries, keys: keys, values: values,
+            scale: scale, mask: mask
+        )
+    } else if let quantizedKVCache = cache as? QuantizedKVCacheProtocol {
         let (quantizedKeys, quantizedValues) = quantizedKVCache.updateQuantized(
             keys: keys, values: values)
         return quantizedScaledDotProductAttention(

Libraries/MLXLMCommon/Documentation.docc/Documentation.md

@@ -7,6 +7,7 @@ Common language model code.
 - <doc:model-compatibility>
 - <doc:upgrade>
 - <doc:wired-memory>
+- <doc:kv-cache-quantization>
 
 ## Other MLX Libraries Packages
 

Libraries/MLXLMCommon/Documentation.docc/kv-cache-quantization.md

@@ -0,0 +1,109 @@
+# KV Cache Quantization
+
+Reduce the memory footprint of long-context generation by compressing the
+key/value cache.
+
+## Overview
+
+At long context lengths the KV cache, not the weights, dominates memory. Two
+mechanisms are available through ``GenerateParameters``:
+
+- **Affine quantization**: `kvBits`, `kvGroupSize`, and `quantizedKVStart`
+  quantize both K and V with MLX's affine scheme. The named schemes
+  `"affine4"` and `"affine8"` are shorthands for the same path.
+- **TurboQuant schemes via `kvScheme`**: vectors are rotated with a
+  Walsh-Hadamard transform and quantized against a Lloyd-Max codebook.
+  TurboQuant schemes are asymmetric. Keys and values can use different
+  precision because attention quality is far more sensitive to key error
+  (softmax amplifies it) than to value error (linear averaging smooths it).
+
+```swift
+var parameters = GenerateParameters()
+parameters.kvScheme = "turbo8v3"   // recommended default
+```
+
+Prefill is unaffected: the cache stores raw fp16 during prompt processing,
+compresses on the first decode step, and encodes each new token incrementally
+afterwards. Compressed caches round-trip through prompt-cache save/restore.
+
+## Scheme reference
+
+Scheme names read `turbo<K-bits>v<V-bits>`; `0` means keys stay fp16.
+
+| scheme | keys | values | KV compression | character |
+|---|---|---|---|---|
+| `affine8` | 8-bit affine | 8-bit affine | 1.88x | near-lossless on most models, full decode speed |
+| `affine4` | 4-bit affine | 4-bit affine | 3.56x | collapses on some families; validate first |
+| `turbo0v4` | fp16 | 4-bit turbo | 1.58x | safest start; beats affine8 quality on most models tested |
+| `turbo0v3` | fp16 | 3-bit turbo | 1.66x | light value compression |
+| `turbo0v2` | fp16 | 2-bit turbo | 1.58x† | aggressive value compression |
+| `turbo8v4` | 8-bit affine | 4-bit turbo | 2.51x | conservative asymmetric |
+| `turbo8v3` | 8-bit affine | 3-bit turbo | 2.75x | recommended default |
+| `turbo8v2` | 8-bit affine | 2-bit turbo | 2.32x† | memory-bound long context |
+| `turbo4`, `turbo3`, `turbo2` | turbo | turbo | up to 3.4x† | maximum compression; key sensitivity varies strongly by family |
+
+† boundary-layer protection auto-engages for fragile schemes (turbo-quantized
+keys or 2-bit values): the first and last two attention layers use 8-bit
+affine instead of the turbo scheme.
+
+## Choosing a scheme
+
+Start asymmetric and light, verify output quality on your model, then
+compress further:
+
+1. `turbo0v4`: keys untouched. If this is not faithful, the model is
+   unusually quantization-sensitive; stop here.
+2. `turbo8v4`: 8-bit keys are near-lossless at half the key bytes.
+3. `turbo8v3`: the sweet spot for most dense models.
+4. `turbo8v2`: memory-bound long context, after validating step 3.
+5. Symmetric `turbo4`, `turbo3`, `turbo2`: maximum compression. Key
+   compression is where models break. Per-dimension key calibration
+   (computed automatically from the prefill cache at compression time)
+   equalizes post-rotation key variance and recovers most of the gap on
+   sensitive families; sensitivity still varies, so validate on your model
+   before deploying.
+
+## Family sensitivity (measured)
+
+WikiText-2 decode-time KL divergence against an fp16 cache with identical
+weights, so the numbers isolate cache compression from weight quantization:
+
+- Mistral-class models tolerate even symmetric `turbo4` (KLD 0.040 at 2.8x
+  on Mistral-7B).
+- qk-norm families (Qwen3) and small models are the most sensitive to key
+  compression. Per-dimension key calibration recovers most of it: turbo4
+  KLD 2.65 to 0.15 on Qwen3-1.7B, 2.76 to 0.036 on Phi-4-mini, 0.62 to
+  0.060 on Qwen2.5-7B, all at about 3.3x. The asymmetric family stays
+  healthy with or without calibration.
+- Qwen2.5-class models carry large key outliers in their first and last
+  layers; attention scores are computed in f32 for this reason, and affine8
+  is notably weaker on this family (KLD 0.041) than the turbo value schemes
+  (turbo0v4 KLD 0.005).
+- Phi-family models are unusually affine-friendly (affine8 KLD 0.0004).
+- Symmetric key compression improves with model size (KLD 2.7 at 1.7B, 0.6
+  at 7B on the same family).
+
+## Performance
+
+TurboQuant decode runs through JIT-compiled Metal kernels; kernels are
+compiled once per model shape, not per call. Single-token decode is a
+SIMD-parallel flash kernel (packed, raw fp16, or 8-bit affine key scoring
+with inline value dequantization). Measured on Qwen3-1.7B (M5 Max, fp16
+150 tok/s): turbo8v3 114, turbo4 122, turbo0v4 102. Prefill stays raw
+fp16, so prefill throughput is unaffected. Use TurboQuant when memory is
+the constraint (long contexts, larger models per machine); use affine8
+when exact fp16-parity decode matters more than footprint.
+
+## Limitations
+
+- Rotating and sliding-window cache layers (most Gemma layers) are not
+  converted; their storage is already bounded by the window size and their
+  eviction layout does not fit the sequential-append compression path. On
+  mixed models the global (non-rotating) layers, where cache memory
+  actually grows, do compress, and a one-time notice lists the layers that
+  kept fp16 rotating caches. Hybrid recurrent layers are likewise left
+  untouched.
+- Unrecognized scheme strings are ignored.
+- For memory estimation with wired limits see <doc:wired-memory>; effective
+  bytes per element follow from the table above (for example `turbo8v3` is
+  about 0.73 bytes per K/V element pair average against 4 for fp16).

Libraries/MLXLMCommon/Evaluate.swift

@@ -74,8 +74,21 @@ public struct GenerateParameters: Sendable {
     public var quantizedKVStart: Int
 
     /// KV cache compression scheme. Overrides kvBits when set.
-    /// Built-in: "affine4", "affine8" (equivalent to kvBits 4/8).
-    /// Extensible for custom schemes (e.g. WHT-based compression).
+    ///
+    /// Built-in affine: "affine4", "affine8" (equivalent to kvBits 4/8).
+    ///
+    /// TurboQuant schemes (`turbo<K>v<V>`, keys × values). Start light,
+    /// verify quality on your model, then ratchet V:
+    /// - "turbo0v4"  FP16 K + 4-bit V, the safest start
+    /// - "turbo0v3"/"turbo0v2"  FP16 K, more aggressive V
+    /// - "turbo8v4"  8-bit affine K + 4-bit V, conservative asymmetric
+    /// - "turbo8v3"  recommended default (near-lossless K, compressed V)
+    /// - "turbo8v2"  aggressive V (boundary-layer protection auto-engages)
+    /// - "turbo4"/"turbo4v2"/"turbo3"/"turbo2"  turbo-quantized keys:
+    ///   maximum compression; K sensitivity varies by model family, so
+    ///   validate on your model (asym is the recommended starting point).
+    ///
+    /// Unrecognized schemes are ignored. See `resolveTurboScheme`.
     public var kvScheme: String?
 
     /// Sampling temperature
@@ -746,17 +759,35 @@ public struct TokenIterator: TokenIteratorProtocol {
             return nil
         }
 
-        // save current value -- this will be returned
-        let previousY = y
-
-        // compute the next state and async eval the next token
-        let token = step(previous: previousY)
-        y = .init(tokens: token)
-        asyncEval(token)
+        // Drain autoreleased MLX temporaries every step: a full model forward
+        // produces hundreds of autoreleased wrapper objects, and the token
+        // loop never returns to a pool boundary on its own; without this,
+        // long generations grow host memory without bound.
+        return autoreleasepool {
+            // save current value -- this will be returned
+            let previousY = y
 
-        tokenCount += 1
+            // compute the next state and async eval the next token
+            let token = step(previous: previousY)
+            y = .init(tokens: token)
+            // Evaluate the cache state together with the token: caches update
+            // through functional ops (concatenation, slice assignment), and an
+            // unevaluated chain of those updates keeps every prior step's
+            // intermediates alive. Python mlx-lm settles cache state the same
+            // way in its generation loop.
+            asyncEval([token] + cache.flatMap { $0.state })
+
+            tokenCount += 1
+
+            // Periodically return freed buffers that cannot be reused (odd or
+            // monotonically growing sizes accumulate in the pool otherwise).
+            // Matches mlx-lm's clear cadence.
+            if tokenCount % 256 == 0 {
+                MLX.Memory.clearCache()
+            }
 
-        return previousY.tokens.item(Int.self)
+            return previousY.tokens.item(Int.self)
+        }
     }
 }
 
@@ -861,7 +892,8 @@ public struct SpeculativeTokenIterator: TokenIteratorProtocol {
                 cache: &cache,
                 kvBits: parameters.kvBits,
                 kvGroupSize: parameters.kvGroupSize,
-                quantizedKVStart: parameters.quantizedKVStart
+                quantizedKVStart: parameters.quantizedKVStart,
+                kvScheme: parameters.kvScheme
             )
         }
 
@@ -1018,10 +1050,12 @@ public struct SpeculativeTokenIterator: TokenIteratorProtocol {
             return token
         }
 
-        // Run a new speculation round
+        // Run a new speculation round. Pooled: a round runs draft + verify
+        // forwards whose autoreleased MLX temporaries otherwise accumulate
+        // across the whole generation.
         pendingTokens.removeAll(keepingCapacity: true)
         pendingIndex = 0
-        speculateRound()
+        autoreleasepool { speculateRound() }
 
         if pendingTokens.isEmpty {
             return nil
@@ -1168,7 +1202,7 @@ private func runSynchronousGenerationLoop(
     var iterator = iterator
     var stopReason: GenerateStopReason?
 
-    while let token = iterator.next() {
+    while let token = autoreleasepool(invoking: { iterator.next() }) {
         // Compute the timing for the prompt.
         if promptTime == 0 {
             let now = Date.timeIntervalSinceReferenceDate
@@ -1864,12 +1898,12 @@ private func generateLoopTask<Handler: TokenLoopHandler>(
             // Check cancellation BEFORE iterator.next(): next() calls asyncEval() to
             // pipeline the next GPU evaluation, so checking after it (the previous
             // `while let token = iterator.next()` form) allowed one extra asyncEval to be
-            // submitted post-cancellation — which faults if the app has backgrounded
+            // submitted post-cancellation, which faults if the app has backgrounded
             // (kIOGPUCommandBufferCallbackErrorBackgroundExecutionNotPermitted). The
             // post-loop block below assigns `.cancelled`; Stream().synchronize() still
             // settles any in-flight evaluation at the end of the task body.
             tokenLoop: while !Task.isCancelled {
-                guard let token = iterator.next() else { break }
+                guard let token = autoreleasepool(invoking: { iterator.next() }) else { break }
 
                 if promptTime == 0 {
                     let now = Date.timeIntervalSinceReferenceDate

Libraries/MLXLMCommon/KVCache.swift

@@ -797,7 +797,7 @@ public class RotatingKVCache: BaseKVCache, CustomDebugStringConvertible {
     }
 }
 
-private func resolvedKVQuantizationGroupSize(
+func resolvedKVQuantizationGroupSize(
     requested: Int,
     keyHeadDim: Int,
     valueHeadDim: Int
@@ -1577,6 +1577,7 @@ private func cacheClassName(_ cache: KVCache) -> String {
     case is ArraysCache: return "ArraysCache"
     case is RotatingKVCache: return "RotatingKVCache"
     case is QuantizedKVCache: return "QuantizedKVCache"
+    case is TurboQuantKVCache: return "TurboQuantKVCache"
     case is KVCacheSimple: return "KVCache"
     case is CacheList: return "CacheList"
     default: return "KVCache"
@@ -1730,6 +1731,22 @@ private func restoreCacheFromMetaState(
         cache.restoreFromMetaState(state: state, savedMetaState: metaState)
         return cache
 
+    case "TurboQuantKVCache":
+        guard metaState.count >= 5,
+            let bits = Int(metaState[1]),
+            let keyBits = Int(metaState[2]),
+            let valueBits = Int(metaState[3]),
+            let seed = UInt64(metaState[4])
+        else {
+            throw KVCacheError(
+                message: "Invalid TurboQuantKVCache metaState")
+        }
+        let cache = TurboQuantKVCache(
+            bits: bits, keyBits: keyBits, valueBits: valueBits, seed: seed)
+        cache.state = state
+        cache.metaState = metaState
+        return cache
+
     case "CacheList":
         return try CacheList.fromState(state: state, metaState: metaState)
 
@@ -2004,15 +2021,28 @@ public func resolveAffineScheme(_ scheme: String?) -> (bits: Int, groupSize: Int
 ///   - kvGroupSize: Group size for quantization
 ///   - quantizedKVStart: Token count threshold to begin quantizing
 ///   - kvScheme: Scheme selector; overrides kvBits when it names a built-in
-///     affine scheme ("affine4", "affine8"). Unrecognized schemes are left to
-///     custom cache implementations and do not quantize here.
+///     affine scheme ("affine4", "affine8") or a TurboQuant scheme
+///     ("turbo4", "turbo4v2", ...). Unrecognized schemes are left to custom
+///     cache implementations and do not quantize here.
 public func maybeQuantizeKVCache(
     cache: inout [KVCache],
     kvBits: Int?,
     kvGroupSize: Int = 64,
     quantizedKVStart: Int = 0,
     kvScheme: String? = nil
 ) {
+    // TurboQuant schemes convert eligible layers to TurboQuantKVCache
+    // (handled in TurboQuantKVCache.swift to keep this file scheme-agnostic).
+    if let scheme = kvScheme, let turbo = resolveTurboScheme(scheme) {
+        maybeTurboQuantizeKVCache(
+            cache: &cache,
+            keyBits: turbo.keyBits,
+            valueBits: turbo.valueBits,
+            quantizedKVStart: quantizedKVStart
+        )
+        return
+    }
+
     // Resolve effective bits: kvScheme overrides kvBits.
     let effectiveBits: Int
     let effectiveGroupSize: Int

Libraries/MLXLMCommon/MTPSpeculativeTokenIterator.swift

@@ -423,7 +423,7 @@ public struct MTPSpeculativeTokenIterator: TokenIteratorProtocol {
         // Run a new speculation round (may transition to passthrough).
         pendingTokens.removeAll(keepingCapacity: true)
         pendingIndex = 0
-        speculateRound()
+        autoreleasepool { speculateRound() }
 
         if pendingTokens.isEmpty {
             // speculateRound chose passthrough -- fall through.