add TurboQuant KV cache compression

WHT rotation + Lloyd-Max codebook quantization for the KV cache, with
asymmetric K/V bit-widths, routed through the kvScheme parameter. Two-phase
design: raw FP16 during prefill (zero overhead), batch-compress on the
first decode step, fused Metal kernels for encode and compressed-domain
attention afterwards.

Stacked on the kvScheme plumbing branch; resolveTurboScheme +
maybeTurboQuantizeKVCache convert eligible KVCacheSimple layers when a
turbo scheme is selected.

Changes from the earlier draft of this branch:

- f32 rotation matrices (bf16 rounding in the rotation compounds across
  layers on larger models)
- norm correction applied on the WHT path too, matching the dense path
  and the fused encode kernel (quantized unit vectors land ~5-10% short
  without it)
- the codec setup crash this branch used to hit is fixed by computing
  tokenCount from the leading dims instead of an out-of-range dim index,
  so the external array-metadata guard is no longer needed
- raw-key schemes (turbo0v2/turbo0v4) are not exposed via kvScheme yet:
  the rawKeyMode decode path produces NaN from the second decode step on
  real checkpoints while staying unit-clean at every probed shape, dtype,
  and magnitude — tracked for a dedicated debugging pass; the mode stays
  reachable through the TurboQuantKVCache initializer
- CI-runnable test suite: rotation orthogonality, codec round trip,
  cache offsets/trim, compressed-vs-FP16 attention parity (incl. GQA
  head-mapping with head-distinguishable values and the kvScheme
  conversion route), and an end-to-end synthetic-model generation test

Author: TheTom

Libraries/MLXLMCommon/AttentionUtils.swift

@@ -51,7 +51,23 @@ 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 {
+            // Prefill (L>1): raw update + standard SDPA. Zero overhead.
+            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): compressed cache path. First call triggers
+        // compressRawCache() inside compressedAttention.
+        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/Evaluate.swift

@@ -73,8 +73,8 @@ 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: "affine4", "affine8" (equivalent to kvBits 4/8), plus the
+    /// TurboQuant family ("turbo4", "turbo3", "turbo2", "turbo4v2", "turbo4v3").
     public var kvScheme: String?
 
     /// Sampling temperature
@@ -810,7 +810,8 @@ public struct SpeculativeTokenIterator: TokenIteratorProtocol {
                 cache: &cache,
                 kvBits: parameters.kvBits,
                 kvGroupSize: parameters.kvGroupSize,
-                quantizedKVStart: parameters.quantizedKVStart
+                quantizedKVStart: parameters.quantizedKVStart,
+                kvScheme: parameters.kvScheme
             )
         }
 

Libraries/MLXLMCommon/KVCache.swift

@@ -1529,6 +1529,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"
@@ -1682,6 +1683,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)
 
@@ -1948,15 +1965,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