From df227d703b02ca1eb41ae1f2e57e89c7a8eae434 Mon Sep 17 00:00:00 2001 From: TheTom Date: Thu, 16 Jul 2026 15:07:20 -0500 Subject: [PATCH 1/3] Add TurboQuant KV cache compression Walsh-Hadamard rotation + Lloyd-Max codebook KV compression behind the kvScheme parameter (#230), with JIT Metal kernels: fused encode, a SIMD-parallel two-pass online-softmax flash decode (packed-K, raw-fp16-K, and 8-bit-affine-K score variants), and an inline sparse value kernel. Schemes are asymmetric because keys are far more quality-sensitive than values (softmax amplifies key error, value error averages out): the recommended tiers keep K at fp16 (turbo0v4/v3/v2) or 8-bit affine (turbo8v4/v3/v2) while V takes the compression; symmetric turbo4/3/2 compress both sides. Symmetric key quality rides on per-dimension key calibration: at compression time the cache derives a per-dimension scale from the post-rotation key variance of the prefill cache and folds it into the query side, leaving the score kernels untouched; the calibrated encode rotates through the codec's MLX ops so kernel and reference quantize identical values. Boundary-layer protection converts the first and last attention layers to 8-bit affine for fragile schemes. Rotating (sliding-window) cache layers stay fp16 with a one-time notice; global layers on mixed models compress normally. Per-call kernel values travel in a parameter buffer rather than template arguments, so each shape compiles once. Caches stay raw fp16 through prefill, compress on the first decode step, and encode incrementally; state round-trips through prompt cache save/restore, and speculative decoding routes through the compressed path. Generation settles cache state with the token eval and drains autoreleased temporaries per step so long generations hold steady memory. --- Libraries/MLXLLM/LLMModel.swift | 15 +- Libraries/MLXLMCommon/AttentionUtils.swift | 20 +- Libraries/MLXLMCommon/Evaluate.swift | 68 +- Libraries/MLXLMCommon/KVCache.swift | 36 +- .../MTPSpeculativeTokenIterator.swift | 2 +- Libraries/MLXLMCommon/TurboQuantKVCache.swift | 1765 ++++++++++++ Libraries/MLXLMCommon/TurboQuantKernels.swift | 2367 +++++++++++++++++ Libraries/MLXVLM/Models/Gemma4.swift | 383 +-- Tests/MLXLMTests/Gemma4PoolingTests.swift | 84 - Tests/MLXLMTests/Gemma4ResizeTests.swift | 113 + 10 files changed, 4593 insertions(+), 260 deletions(-) create mode 100644 Libraries/MLXLMCommon/TurboQuantKVCache.swift create mode 100644 Libraries/MLXLMCommon/TurboQuantKernels.swift delete mode 100644 Tests/MLXLMTests/Gemma4PoolingTests.swift create mode 100644 Tests/MLXLMTests/Gemma4ResizeTests.swift diff --git a/Libraries/MLXLLM/LLMModel.swift b/Libraries/MLXLLM/LLMModel.swift index 87aeea0a5..2b608818a 100644 --- a/Libraries/MLXLLM/LLMModel.swift +++ b/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, .. 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( diff --git a/Libraries/MLXLMCommon/Evaluate.swift b/Libraries/MLXLMCommon/Evaluate.swift index a0f920082..df34922e1 100644 --- a/Libraries/MLXLMCommon/Evaluate.swift +++ b/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 (`turbov`, 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( // 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 diff --git a/Libraries/MLXLMCommon/KVCache.swift b/Libraries/MLXLMCommon/KVCache.swift index c4fa11fa2..d4ba7db42 100644 --- a/Libraries/MLXLMCommon/KVCache.swift +++ b/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,8 +2021,9 @@ 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?, @@ -2013,6 +2031,18 @@ public func maybeQuantizeKVCache( 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 diff --git a/Libraries/MLXLMCommon/MTPSpeculativeTokenIterator.swift b/Libraries/MLXLMCommon/MTPSpeculativeTokenIterator.swift index b990a2cdb..b0a75fdae 100644 --- a/Libraries/MLXLMCommon/MTPSpeculativeTokenIterator.swift +++ b/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. diff --git a/Libraries/MLXLMCommon/TurboQuantKVCache.swift b/Libraries/MLXLMCommon/TurboQuantKVCache.swift new file mode 100644 index 000000000..9f8ff9af7 --- /dev/null +++ b/Libraries/MLXLMCommon/TurboQuantKVCache.swift @@ -0,0 +1,1765 @@ +// Copyright © 2026 Apple Inc. +// +// Implements TurboQuant Algorithm 1 (MSE-optimal, arXiv:2504.19874) for KV cache: +// rotation Π + optimal Lloyd-Max scalar codebook quantization on Beta distribution. +// +// Both keys and values use Algorithm 1 (MSE-only at b bits). QJL (Algorithm 2) +// is omitted, Tom Turney's research shows no quality benefit on Apple Silicon, +// and at 4-bit the MSE bias is negligible (paper Section 3.2: bias = 2/π, +// diminishing with bit-width). +// +// Enhancements beyond paper: +// - Norm extraction/restoration: paper assumes ||x||=1; we store norms for arbitrary vectors +// - Norm correction: store ||x|| / ||ỹ|| for dense rotation path (WHT skips, orthogonal preserves norms) +// - WHT rotation option: O(d log d) butterfly in Metal kernel for power-of-2 dims +// - Two-phase architecture: raw prefill → batch compress → compressed decode +// - Pre-rotated queries: q' = Π·q computed once, reused for all cached keys +// +// References: +// - TurboQuant: https://arxiv.org/abs/2504.19874 +// - QJL: https://arxiv.org/abs/2406.03482 +// - PolarQuant: https://arxiv.org/abs/2502.02617 + +import Foundation +import MLX +import MLXNN + +extension DType { + fileprivate var bytesPerElement: Int { + switch self { + case .bfloat16, .float16: return 2 + case .float32: return 4 + case .int32, .uint32: return 4 + case .int16, .uint16: return 2 + case .int8, .uint8: return 1 + default: return 4 + } + } +} + +// MARK: - Codebook Generation + +/// Optimal Lloyd-Max codebook centroids for Beta-distributed coordinates. +/// +/// After random orthogonal rotation, each coordinate of a unit-sphere vector +/// follows Beta distribution f_X(x) ∝ (1-x²)^((d-3)/2) on [-1,1]. +/// For large d, this converges to N(0, 1/d). +enum TurboQuantCodebook { + + // MARK: - Pre-computed Centroids + + /// Pre-computed Lloyd-Max centroids for common (dim, bits) pairs. + /// Generated offline via 100-iteration weighted k-means on 32K-point Beta PDF grid. + /// Avoids ~50ms runtime codebook generation per codec. + /// + /// Additional dims (80, 96) are lazily generated on first access to support + /// Qwen3-4B (head_dim=80) and similar models without startup cost for unused dims. + nonisolated(unsafe) private static var precomputed: [Int: [Int: [Float]]] = [ + 64: [ + 2: [-0.18745463, -0.05649366, 0.05649367, 0.18745449], + 3: [ + -0.26375133, -0.16599470, -0.09368263, -0.03040462, 0.03040464, 0.09368261, + 0.16599482, 0.26375186, + ], + 4: [ + -0.32913971, -0.25096416, -0.19681059, -0.15295772, -0.11478586, -0.08000945, + -0.04726735, -0.01563822, 0.01563822, 0.04723797, 0.07994876, 0.11472529, + 0.15289739, 0.19675052, 0.25090477, 0.32908401, + ], + ], + 128: [ + 2: [-0.13302007, -0.03998107, 0.03998102, 0.13302033], + 3: [ + -0.18828832, -0.11801215, -0.06648001, -0.02156330, 0.02156329, 0.06648005, + 0.11801218, 0.18828897, + ], + 4: [ + -0.23639172, -0.17934021, -0.14023653, -0.10881814, -0.08157559, -0.05678632, + -0.03350975, -0.01108178, 0.01108178, 0.03350975, 0.05678631, 0.08157560, + 0.10881804, 0.14023650, 0.17934017, 0.23639278, + ], + ], + 256: [ + 2: [-0.09420358, -0.02827190, 0.02827190, 0.09420330], + 3: [ + -0.13371243, -0.08361249, -0.04704370, -0.01524900, 0.01524901, 0.04704368, + 0.08361248, 0.13371260, + ], + 4: [ + -0.16852295, -0.12754069, -0.09961203, -0.07719406, -0.05781249, -0.04021866, + -0.02370371, -0.00783269, 0.00783269, 0.02370371, 0.04021868, 0.05781246, + 0.07719407, 0.09961203, 0.12754090, 0.16852276, + ], + ], + ] + + /// Lock for thread-safe lazy population of precomputed centroids. + private static let centroidLock = NSLock() + + /// Dims that should be lazily pre-populated (non-power-of-2 dims used by real models). + /// These fall back to dense rotation path since WHT requires power-of-2, but still + /// benefit from cached centroids to avoid ~50ms runtime k-means per codec init. + /// + /// - 80: Qwen3-4B (head_dim=80) + /// - 96: Various smaller models + private static let lazyDims: [Int] = [80, 96] + private static let lazyBits: [Int] = [2, 3, 4] + + /// Ensure centroids for a given dim are populated. Thread-safe, generates once. + private static func ensureCentroidsPopulated(dim: Int) { + centroidLock.lock() + let exists = precomputed[dim] != nil + centroidLock.unlock() + guard !exists else { return } + + // Generate all bit-widths for this dim + var dimTable: [Int: [Float]] = [:] + for bits in lazyBits { + dimTable[bits] = generateCentroids(dim: dim, bits: bits) + } + + centroidLock.lock() + // Double-check after lock (another thread may have populated) + if precomputed[dim] == nil { + precomputed[dim] = dimTable + } + centroidLock.unlock() + } + + // MARK: - Public API + + /// Codebook centroids for (dim, bits). Uses pre-computed table for common configs, + /// lazily generates and caches for known model dims (80, 96), falls back to + /// runtime generation for truly uncommon ones. + /// + static func codebook(dim: Int, bits: Int) -> MLXArray { + if let dimTable = precomputed[dim], let centroids = dimTable[bits] { + return MLXArray(centroids) + } + // Lazy populate for known model dims (Qwen3-4B dim=80, etc.) + if lazyDims.contains(dim) { + ensureCentroidsPopulated(dim: dim) + if let dimTable = precomputed[dim], let centroids = dimTable[bits] { + return MLXArray(centroids) + } + } + let centroids = generateCentroids(dim: dim, bits: bits) + return MLXArray(centroids) + } + + /// Codebook boundaries (midpoints between adjacent centroids). + static func boundaries(dim: Int, bits: Int) -> MLXArray { + let centroids: [Float] + if let dimTable = precomputed[dim], let cached = dimTable[bits] { + centroids = cached + } else if lazyDims.contains(dim) { + ensureCentroidsPopulated(dim: dim) + if let dimTable = precomputed[dim], let cached = dimTable[bits] { + centroids = cached + } else { + centroids = generateCentroids(dim: dim, bits: bits) + } + } else { + centroids = generateCentroids(dim: dim, bits: bits) + } + var bounds = [Float]() + for i in 0 ..< centroids.count - 1 { + bounds.append((centroids[i] + centroids[i + 1]) / 2.0) + } + return MLXArray(bounds) + } + + /// Generate codebook centroids via weighted k-means on Beta distribution. + /// Used as fallback for uncommon (dim, bits) pairs not in the pre-computed table. + static func generateCentroids(dim: Int, bits: Int) -> [Float] { + let levels = 1 << bits + let gridSize = 32768 + let sigma = 1.0 / sqrt(Float(dim)) + + // Generate grid points and PDF weights + var grid = [Float](repeating: 0, count: gridSize) + var weights = [Float](repeating: 0, count: gridSize) + for i in 0 ..< gridSize { + let x = -1.0 + 2.0 * Float(i) / Float(gridSize - 1) + grid[i] = x + // Beta PDF ∝ (1 - x²)^((d-3)/2), approximated by Gaussian for large d + let exponent = Float(dim - 3) / 2.0 + let w = pow(max(1.0 - x * x, 1e-30), exponent) + weights[i] = w + } + + // Initialize centroids via quantiles + let totalW = weights.reduce(0, +) + var centroids = [Float](repeating: 0, count: levels) + var cumW: Float = 0 + var ci = 0 + for i in 0 ..< gridSize { + cumW += weights[i] + let target = (Float(ci) + 0.5) / Float(levels) * totalW + if cumW >= target && ci < levels { + centroids[ci] = grid[i] + ci += 1 + } + } + // Fill remaining + while ci < levels { + centroids[ci] = centroids[ci - 1] + sigma + ci += 1 + } + + // K-means iterations + for _ in 0 ..< 100 { + var sums = [Float](repeating: 0, count: levels) + var counts = [Float](repeating: 0, count: levels) + for i in 0 ..< gridSize { + var bestJ = 0 + var bestDist = Float.infinity + for j in 0 ..< levels { + let d = abs(grid[i] - centroids[j]) + if d < bestDist { + bestDist = d + bestJ = j + } + } + sums[bestJ] += grid[i] * weights[i] + counts[bestJ] += weights[i] + } + for j in 0 ..< levels { + if counts[j] > 0 { centroids[j] = sums[j] / counts[j] } + } + } + + return centroids.sorted() + } +} + +// MARK: - Rotation Matrix + +/// Random orthogonal rotation matrix generation. +/// +/// TurboQuant Algorithm 1 line 2: Π ∈ ℝ^(d×d) via QR decomposition +/// on random Gaussian matrix. Sign-corrected for determinism. +enum TurboQuantRotation { + + /// Generate a deterministic random orthogonal rotation matrix (dense, d×d). + /// Uses QR decomposition on CPU (not yet GPU-supported in MLX). + static func rotationMatrix(dim: Int, seed: UInt64) -> MLXArray { + let key = MLXRandom.key(seed) + let gaussian = MLXRandom.normal([dim, dim], key: key) + + // QR on CPU (MLX GPU QR not supported yet). MLX's QR runs in f32 + // regardless of input dtype, so orthogonality is good to ~1e-3, // fine for the dense fallback path, whose matched-norm scales absorb + // the residual. Real models (pow2 head dims) take the exact WHT path. + let (q, r) = MLXLinalg.qr(gaussian, stream: .cpu) + let diagR = r.diagonal(stream: .cpu) + let signs = sign(diagR, stream: .cpu) + let result = q * expandedDimensions(signs, axis: 0) + eval(result) + return result + } + + /// Generate a Hadamard matrix of size dim × dim via recursive Kronecker product. + /// Requires dim to be a power of 2. The resulting matrix H satisfies H·H = dim·I. + static func hadamardMatrix(dim: Int) -> MLXArray { + precondition(dim > 0 && (dim & (dim - 1)) == 0, "dim must be power of 2") + // Build recursively: H_1 = [[1]], H_2n = [[H_n, H_n], [H_n, -H_n]] + var h: [[Float]] = [[1.0]] + var size = 1 + while size < dim { + var newH = [[Float]](repeating: [Float](repeating: 0, count: size * 2), count: size * 2) + for i in 0 ..< size { + for j in 0 ..< size { + newH[i][j] = h[i][j] + newH[i][j + size] = h[i][j] + newH[i + size][j] = h[i][j] + newH[i + size][j + size] = -h[i][j] + } + } + h = newH + size *= 2 + } + let flat = h.flatMap { $0 } + let result = MLXArray(flat, [dim, dim]) + eval(result) + return result + } + + /// Generate WHT sign vector: random ±1 per dimension, length d. + /// Used with Walsh-Hadamard Transform for O(d log d) rotation. + static func whtSigns(dim: Int, seed: UInt64) -> MLXArray { + let key = MLXRandom.key(seed) + // Random bits → ±1 + // Generate random ±1 signs using uniform random + let uniform = MLXRandom.uniform(low: 0, high: 1, [dim], key: key) + let signs = MLX.where( + uniform .> MLXArray(Float(0.5)), MLXArray(Float(1.0)), MLXArray(Float(-1.0))) + eval(signs) + return signs + } + + /// Apply WHT butterfly on the last dimension of x. Shape-preserving. + /// Computes unnormalized Walsh-Hadamard transform: H * x along last dim. + private static func whtButterfly(_ x: MLXArray) -> MLXArray { + let dim = x.dim(-1) + let logDim = Int(log2(Double(dim))) + let origShape = x.shape + // Flatten leading dims: [N, dim] + let N = origShape.dropLast().reduce(1, *) + var y = x.reshaped([N, dim]) + + for s in 0 ..< logDim { + let halfBlock = 1 << s + let blockSize = halfBlock << 1 + let numBlocks = dim / blockSize + // Reshape to [N, numBlocks, 2, halfBlock] + y = y.reshaped([N, numBlocks, blockSize]) + let a = y[0..., 0..., .. MLXArray { + let dim = x.dim(-1) + precondition(dim > 0 && (dim & (dim - 1)) == 0, "dim must be power of 2") + let signed = x * signs + let transformed = whtButterfly(signed) + return transformed * MLXArray(Float(1.0 / sqrt(Float(dim)))) + } + + /// Apply SRHT inverse rotation: x = diag(signs) * H * y / sqrt(dim) + /// WHT is self-inverse up to scale. Inverse of (H·D/√d) is (D·H/√d). + static func fwhtInverse(_ y: MLXArray, signs: MLXArray) -> MLXArray { + let dim = y.dim(-1) + precondition(dim > 0 && (dim & (dim - 1)) == 0, "dim must be power of 2") + let transformed = whtButterfly(y) + return transformed * MLXArray(Float(1.0 / sqrt(Float(dim)))) * signs + } +} + +// MARK: - Bit Packing + +/// Efficient bit packing/unpacking for codebook indices. +enum TurboQuantPacking { + + /// Number of uint32 words needed to pack `count` values at `bits` each. + static func packedWidth(count: Int, bits: Int) -> Int { + (count * bits + 31) / 32 + } + + /// Pack b-bit indices into uint32 words. + /// Input: [rows, count] as uint32 (values 0..2^bits-1) + /// Output: [rows, packedWidth] as uint32 + static func packLowBit(_ indices: MLXArray, bits: Int) -> MLXArray { + let count = indices.dim(-1) + let batchShape = Array(indices.shape.dropLast()) + let rows = batchShape.reduce(1, *) + let flat = indices.reshaped([rows, count]) + let pw = packedWidth(count: count, bits: bits) + let mask = UInt32((1 << bits) - 1) + + var wordArrays = [MLXArray]() + for w in 0 ..< pw { + var word = MLXArray.zeros([rows], dtype: .uint32) + for d in 0 ..< count { + let bitOffset = d * bits + let wordIdx = bitOffset / 32 + let offset = bitOffset % 32 + let spill = offset + bits - 32 + + if wordIdx == w { + let shifted = + (flat[0..., d].asType(.uint32) & MLXArray(mask)) << MLXArray(UInt32(offset)) + word = word | shifted + } + if spill > 0 && wordIdx + 1 == w { + let shifted = + (flat[0..., d].asType(.uint32) & MLXArray(mask)) + >> MLXArray(UInt32(bits - spill)) + word = word | shifted + } + } + wordArrays.append(expandedDimensions(word, axis: -1)) + } + let packed = concatenated(wordArrays, axis: -1) // [rows, pw] + return packed.reshaped(batchShape + [pw]) + } + + /// Unpack b-bit indices from uint32 words. + /// Input: [rows, packedWidth] as uint32 + /// Output: [rows, count] as uint32 + static func unpackLowBit(_ packed: MLXArray, bits: Int, count: Int) -> MLXArray { + let shape = packed.shape + let batchShape = Array(shape.dropLast()) + let rows = batchShape.reduce(1, *) + let flat = packed.reshaped([rows, -1]) + let mask = UInt32((1 << bits) - 1) + + var dimArrays = [MLXArray]() + for d in 0 ..< count { + let bitOffset = d * bits + let wordIdx = bitOffset / 32 + let offset = bitOffset % 32 + let spill = offset + bits - 32 + + var value = (flat[0..., wordIdx] >> MLXArray(UInt32(offset))) & MLXArray(mask) + if spill > 0 { + let high = + (flat[0..., wordIdx + 1] << MLXArray(UInt32(bits - spill))) & MLXArray(mask) + value = value | high + } + dimArrays.append(expandedDimensions(value, axis: -1)) + } + let unpacked = concatenated(dimArrays, axis: -1) // [rows, count] + return unpacked.reshaped(batchShape + [count]) + } +} + +// MARK: - MSE Codec (TurboQuant Algorithm 1) + +/// State for MSE-quantized vectors. +struct MSECodecState { + var norms: MLXArray // [B, H, T], original vector L2 norms + var packedIndices: MLXArray // [B, H, T, PackedWidth], packed codebook indices + var tokenCount: Int + let dim: Int + let bits: Int +} + +/// MSE-optimal codec per TurboQuant Algorithm 1. +/// +/// QUANT: y ← Π·x, idx_j ← argmin|y_j - c_k| +/// DEQUANT: ỹ_j ← c_{idx_j}, x̃ ← Π^T · ỹ +class MSECodec { + let dim: Int + let bits: Int + let seed: UInt64 + + /// Codebook centroids [2^bits] + let codebook: MLXArray + /// Codebook boundaries for fast quantization [2^bits - 1] + let boundaries: MLXArray + + /// Whether to use WHT (power-of-2 dim) or dense rotation + let useWHT: Bool + /// WHT sign vector [dim], for O(d log d) Metal encode kernel (power-of-2 dims only) + let whtSigns: MLXArray? + /// Dense rotation matrix Π [dim, dim], used for decode/query rotation (single matmul, fast) + let rotation: MLXArray + /// Π^T, for forward rotation + let rotationT: MLXArray + + init(dim: Int, bits: Int, seed: UInt64 = 42) { + self.dim = dim + self.bits = bits + self.seed = seed + self.codebook = TurboQuantCodebook.codebook(dim: dim, bits: bits) + self.boundaries = TurboQuantCodebook.boundaries(dim: dim, bits: bits) + + // Use WHT for power-of-2 dims (O(d log d) Metal encode kernel) + let isPowerOf2 = dim > 0 && (dim & (dim - 1)) == 0 + self.useWHT = isPowerOf2 && dim <= 1024 + if useWHT { + let signs = TurboQuantRotation.whtSigns(dim: dim, seed: seed) + self.whtSigns = signs + // Build dense WHT rotation matrix for decode/query path (single matmul is faster + // than FWHT butterfly via MLX ops due to graph overhead) + let hadamard = TurboQuantRotation.hadamardMatrix(dim: dim) + let signsDiag = expandedDimensions(signs, axis: 0) + let whtRot = hadamard * signsDiag / MLXArray(Float(sqrt(Float(dim)))) + eval(whtRot) + self.rotation = whtRot + self.rotationT = whtRot.transposed() + } else { + self.whtSigns = nil + self.rotation = TurboQuantRotation.rotationMatrix(dim: dim, seed: seed) + self.rotationT = self.rotation.transposed() + } + } + + /// Encode vectors (Algorithm 1 QUANT) with optional norm correction and optional + /// per-dimension key calibration. + /// Input: [B, H, T, D] + /// Returns MSECodecState with norms and packed indices. + /// + /// WHT path: stores raw norms directly. WHT is an orthogonal transform that preserves + /// norms, so reconstruction_norm ≈ original_norm (within floating point error). + /// Skipping norm correction saves one codebook lookup, one norm computation, and one + /// division per encoded vector. + /// + /// Dense rotation path: stores `original_norm / reconstruction_norm` (norm correction). + /// This compensates for quantization error in the non-orthogonal rotation case. + /// + /// `scale`, when non-nil, is `[dim]` and divides the rotated vector elementwise + /// before quantization: quantize(rotate(k) / scale). This is TurboQuant key + /// calibration (equalize post-rotation per-dimension variance on + /// quantization-sensitive families, see TurboQuantKVCache.computeKeyCalibrationScale). + /// Dividing by scale breaks the orthogonal norm-preservation the WHT raw-norm + /// shortcut relies on, so norm correction always runs when scale is supplied, + /// even on the WHT path. + /// Rotation front half of `encode`: unit-normalize and rotate. + /// + /// Exposed separately so the calibrated fused-kernel dispatch quantizes + /// exactly these values. Boundary quantization is sensitive to the fp + /// reduce order of the rotation; real caches contain degenerate rows + /// (attention sinks) whose rotated coordinates sit near quantization + /// boundaries, so an in-kernel rotation with a different reduce order + /// flips many dimensions of exactly the rows softmax weights most. + /// Sharing these ops makes kernel and MLX encode indices identical. + func rotatedUnit(_ vectors: MLXArray) -> (rotated: MLXArray, norms: MLXArray) { + // Extract norms and normalize (paper assumes unit sphere; we store norms separately) + let norms = sqrt((vectors * vectors).sum(axis: -1)) + let safeNorms = maximum(norms, MLXArray(Float(1e-8))) + let unit = vectors / expandedDimensions(safeNorms, axis: -1) + + // Rotate: y ← Π · x (Algorithm 1 line 5) + return (matmul(unit, rotationT), norms) + } + + func encode(_ vectors: MLXArray, scale: MLXArray? = nil) -> MSECodecState { + let (rotated, norms) = rotatedUnit(vectors) + let calibrated = scale.map { rotated / $0 } ?? rotated + + // Quantize via boundary comparison (fast, no broadcast) + let indices = boundaryQuantize(calibrated) + + let storedNorms: MLXArray + if let scale { + // Calibrated path: undo the /scale divide before measuring the + // reconstruction, per-vector norm is computed AFTER scaling, which + // matches what the flash kernels actually reconstruct (query is + // pre-multiplied by scale, so the dot product resolves in the true + // rotated space, see TurboQuantKVCache.compressedAttention). + let reconstructed = codebook[indices] * scale + let reconNormSq = (reconstructed * reconstructed).sum(axis: -1) + let reconNorms = sqrt(maximum(reconNormSq, MLXArray(Float(1e-16)))) + storedNorms = norms / reconNorms + } else if useWHT { + // WHT fast path: orthogonal transform preserves norms, skip correction + storedNorms = norms + } else { + // Dense rotation path: norm correction compensates for quantization error + let reconstructed = codebook[indices] // [B,H,T,D], quantized approximation in rotated space + let reconNormSq = (reconstructed * reconstructed).sum(axis: -1) + let reconNorms = sqrt(maximum(reconNormSq, MLXArray(Float(1e-16)))) + storedNorms = norms / reconNorms // original_norm / reconstruction_norm + } + + // Pack indices + let packed = TurboQuantPacking.packLowBit(indices, bits: bits) + + return MSECodecState( + norms: storedNorms, + packedIndices: packed, + tokenCount: vectors.dim(2), + dim: dim, + bits: bits + ) + } + + /// Decode from state (Algorithm 1 DEQUANT). + /// `scale`, when non-nil, undoes the calibration divide applied at + /// `encode(_:scale:)`: multiply the codebook lookup elementwise by scale + /// before inverse rotation. Must match the scale passed to encode. + /// Returns: [B, H, T, D] + func decode(_ state: MSECodecState, scale: MLXArray? = nil) -> MLXArray { + // Unpack indices + let indices = TurboQuantPacking.unpackLowBit(state.packedIndices, bits: bits, count: dim) + + // Codebook lookup: ỹ_j ← c_{idx_j} (Algorithm 1 line 9) + var approx = codebook[indices] + if let scale { + approx = approx * scale + } + + // Inverse rotate: x̃ ← Π^T · ỹ (Algorithm 1 line 10) + let unrotated = matmul(approx, rotation) + + // Rescale by stored norms + return expandedDimensions(state.norms, axis: -1) * unrotated + } + + /// Decode in rotated space (skip inverse rotation). + /// Returns centroid values scaled by norm, still in Π-rotated coordinate space. + /// Used with pre-rotated queries for dequant-first SDPA. + func decodeRotated(_ state: MSECodecState) -> MLXArray { + let indices = TurboQuantPacking.unpackLowBit(state.packedIndices, bits: bits, count: dim) + let approx = codebook[indices] + return expandedDimensions(state.norms, axis: -1) * approx + } + + /// Pre-rotate queries for compressed-domain scoring. + /// q' ← Π · q (once per query, reused for all cached keys) + func prepareQueries(_ queries: MLXArray) -> MLXArray { + return matmul(queries, rotationT) + } + + /// Fast quantization via boundary comparison instead of argmin broadcast. + /// boundaries = sorted midpoints between adjacent centroids. + /// Returns uint32 indices in [0, 2^bits - 1]. + func boundaryQuantize(_ rotated: MLXArray) -> MLXArray { + // For each coordinate, count how many boundaries it exceeds + // This gives the codebook index directly + let ndim = rotated.ndim + let expanded = expandedDimensions(rotated, axis: -1) // [..., D, 1] + // Reshape boundaries to broadcast: [1, 1, ..., 1, numBoundaries] + var bShape = [Int](repeating: 1, count: ndim + 1) + bShape[ndim] = boundaries.count + let b = boundaries.reshaped(bShape) + let greater = (expanded .> b).asType(.uint32) // compare against all boundaries + let indices = greater.sum(axis: -1) // count exceeded = index + return indices.asType(.uint32) + } +} + +// MARK: - TurboQuantKVCache + +/// KV cache using TurboQuant compression with two-phase architecture: +/// +/// **Phase 1, Prefill** (L>1): Store raw K/V like KVCacheSimple. Zero overhead. +/// **Transition**: On first decode call, compress entire raw cache in one batch. +/// **Phase 2, Decode** (L=1): Encode 1 new token. Metal kernel scores against +/// all compressed tokens. Zero dequantization. +/// +/// Both keys and values: Algorithm 1 (MSE at b bits, no QJL) +public class TurboQuantKVCache: BaseKVCache { + + public let bits: Int // Legacy: used when keyBits == valueBits + public let keyBits: Int // Bit-width for key compression (0 = raw FP16, no compression) + public let valueBits: Int // Bit-width for value compression (can be lower, V compression is nearly free) + private let seed: UInt64 + + /// Affine-K mode: keys stored as 8-bit affine-quantized (upstream QuantizedKVCache + /// layout, groupSize 64) while values are TurboQuant compressed. The efficient K per + /// the asymmetric-kv paper: near-lossless like FP16 keys at half the bytes. + /// Enabled when keyBits == 8. + public let affineKeyMode: Bool + + /// Affine-K quantization group size. Defaults to 64, but callers that + /// already know the head dimension (e.g. `maybeTurboQuantizeKVCache`) + /// can pass a value pre-resolved via `resolvedKVQuantizationGroupSize` + /// so kernels and quantization agree. Resolved (and possibly adjusted) + /// against the actual head dimension the first time an affine-K encode + /// runs, see `resolveAffineKeyGroupSize(headDim:)`. + public private(set) var keyGroupSize: Int + + // Affine-K storage (wq/scales/biases triplet, grown in steps) + private var affKeyW: MLXArray? + private var affKeyScales: MLXArray? + private var affKeyBiases: MLXArray? + + /// Raw-K mode: keys stay at FP16 (uncompressed) while only values are TurboQuant compressed. + /// This is the single biggest quality finding from TurboQuant+, K precision dominates + /// quality via softmax amplification, V compression is nearly free (linear averaging). + /// Enabled when keyBits == 0. + public let rawKeyMode: Bool + + // Codecs (lazy init) + private var keyMSECodec: MSECodec? // keyBits for keys + private var valueMSECodec: MSECodec? // valueBits for values + + // Phase 1: Raw K/V storage (like KVCacheSimple), used during prefill + private var rawKeys: MLXArray? // [B, H, allocSteps, D] + private var rawValues: MLXArray? // [B, H, allocSteps, D] + private var rawAllocSteps = 0 + + // Phase 2: Compressed storage, used during decode + // MSE-only: packed indices + norms (no QJL, simpler, same quality) + private var keyPackedMSE: MLXArray? + private var keyNorms: MLXArray? + private var valPackedMSE: MLXArray? + private var valNorms: MLXArray? + private var compressedAllocSteps = 0 + + /// Per-dimension key calibration scale s ∈ ℝ^dim, computed once at + /// compressRawCache time from the full prefill K cache and reused for + /// every subsequent encodeNewToken call. `nil` means identity (no + /// calibration): the fast fused Metal encode kernels stay in use for + /// keys. Only the standard (both-K-and-V-quantized) symmetric schemes + /// ever set this, rawKeyMode/affineKeyMode keys are untouched. + private var keyCalibScale: MLXArray? + + /// Whether we've transitioned from raw → compressed + public private(set) var isCompressed = false + + private let step = 256 + + public init( + bits: Int = 4, keyBits: Int? = nil, valueBits: Int? = nil, seed: UInt64 = 42, + keyGroupSize: Int = 64 + ) { + self.bits = bits + self.keyBits = keyBits ?? bits + self.valueBits = valueBits ?? bits + self.rawKeyMode = (keyBits ?? bits) == 0 + self.affineKeyMode = (keyBits ?? bits) == 8 + self.seed = seed + self.keyGroupSize = keyGroupSize + super.init() + } + + /// Resolves (and, if needed, adjusts) `keyGroupSize` against the actual + /// key head dimension the first time an affine-K encode runs. Callers + /// that pre-resolve the group size (`maybeTurboQuantizeKVCache`) hit a + /// no-op here; a directly constructed cache with an incompatible head + /// dimension gets an actionable crash instead of an unhelpful failure + /// deep inside MLX's `quantized(...)`. + private func resolveAffineKeyGroupSize(headDim: Int) { + guard + let resolved = resolvedKVQuantizationGroupSize( + requested: keyGroupSize, keyHeadDim: headDim, valueHeadDim: headDim) + else { + fatalError( + "TurboQuant affine key quantization requires head dimensions divisible by one of the supported group sizes (32, 64, 128). Requested group size: \(keyGroupSize). Key head dim: \(headDim)." + ) + } + keyGroupSize = resolved + } + + override public var isTrimmable: Bool { true } + + // MARK: - Shared Codec Cache + + /// Shared codec cache: all layers with the same (dim, bits, seed) reuse the same codec. + /// Eliminates 56 redundant [128,128] rotation matrices (~7 MB) across 28 layers. + private static let codecLock = NSLock() + nonisolated(unsafe) private static var sharedCodecs: [String: MSECodec] = [:] + + private static func getOrCreateCodec(dim: Int, bits: Int, seed: UInt64) -> MSECodec { + let key = "\(dim)_\(bits)_\(seed)" + codecLock.lock() + if let cached = sharedCodecs[key] { + codecLock.unlock() + return cached + } + codecLock.unlock() + let codec = MSECodec(dim: dim, bits: bits, seed: seed) + codecLock.lock() + sharedCodecs[key] = codec + codecLock.unlock() + return codec + } + + /// Initialize codecs if needed. Uses shared cache to avoid duplicating rotation matrices. + /// In rawKeyMode, key codec is nil, keys stay at FP16, no rotation/quantization needed. + private func ensureCodecs(headDim: Int) { + guard valueMSECodec == nil else { return } + if !rawKeyMode && !affineKeyMode { + keyMSECodec = Self.getOrCreateCodec(dim: headDim, bits: keyBits, seed: seed) + } + valueMSECodec = Self.getOrCreateCodec(dim: headDim, bits: valueBits, seed: seed + 1) + } + + /// Dispatch to WHT or dense fused encode kernel based on codec configuration. + private func fusedEncodeDispatch( + input: MLXArray, codec: MSECodec, headDim: Int + ) -> (packed: MLXArray, norms: MLXArray) { + if codec.useWHT, let signs = codec.whtSigns { + return TurboQuantKernelOps.fusedEncodeWHT( + input: input, whtSigns: signs, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: codec.bits, dim: headDim + ) + } else { + return TurboQuantKernelOps.fusedEncode( + input: input, rotation: codec.rotation, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: codec.bits, dim: headDim + ) + } + } + + /// Calibrated key encode (per-dimension key calibration). The rotation + /// runs as MLX ops through `MSECodec.rotatedUnit`, the same ops + /// `MSECodec.encode(_:scale:)` uses, so the kernel quantizes + /// bit-identical rotated values; only the boundary compare, packing, and + /// norm correction run in the fused kernel. Rotating in-kernel instead + /// flips quantization ties on degenerate rows (attention sinks) whose + /// coordinates sit near boundaries, which measurably distorts attention. + private func fusedEncodeDispatchScaled( + input: MLXArray, scale: MLXArray, codec: MSECodec, headDim: Int + ) -> (packed: MLXArray, norms: MLXArray) { + let (rotated, norms) = codec.rotatedUnit(input.asType(.float32)) + return TurboQuantKernelOps.fusedQuantizePackScaled( + rotated: rotated, rawNorms: norms, scale: scale, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: codec.bits, dim: headDim + ) + } + + // MARK: - Phase 1: Raw Prefill + + /// Prefill update: store raw K/V, return raw. Zero encoding overhead. + /// Uses KVCacheSimple-style allocation with concatenated growth. + override public func update(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray) { + let previous = self.offset + + let reset = + if let currentKeys = self.rawKeys, (previous + keys.dim(2)) > currentKeys.dim(2) { + true + } else { + self.rawKeys == nil + } + if reset { + let B = keys.dim(0) + let kvHeads = keys.dim(1) + let kHeadDim = keys.dim(3) + let vHeadDim = values.dim(3) + + let nSteps = (step + keys.dim(2) - 1) / step + let kShape = [B, kvHeads, nSteps * step, kHeadDim] + let vShape = [B, kvHeads, nSteps * step, vHeadDim] + let newK = MLXArray.zeros(kShape, dtype: keys.dtype) + let newV = MLXArray.zeros(vShape, dtype: values.dtype) + + if var currentKeys = self.rawKeys, var currentValues = self.rawValues { + if previous % step != 0 { + currentKeys = currentKeys[.ellipsis, .. MLXArray? { + let dim = allKeys.dim(-1) + let sampleCount = allKeys.dim(0) * allKeys.dim(1) * allKeys.dim(2) + guard sampleCount >= 32 else { return nil } + + let keysF32 = allKeys.asType(.float32) + let norms = sqrt((keysF32 * keysF32).sum(axis: -1)) + let safeNorms = maximum(norms, MLXArray(Float(1e-8))) + let unit = keysF32 / expandedDimensions(safeNorms, axis: -1) + let rotated = matmul(unit, codec.rotationT).reshaped([-1, dim]) + + let perDimStd = std(rotated, axis: 0) + let meanStd = maximum(perDimStd.mean(), MLXArray(Float(1e-8))) + let clamped = clip(perDimStd / meanStd, min: Float(0.25), max: Float(4.0)) + eval(clamped) + + let maxDeviation = (clamped - MLXArray(Float(1.0))).abs().max().item(Float.self) + guard maxDeviation >= 0.1 else { return nil } + return clamped + } + + // MARK: - Transition: Compress Raw Cache + + /// Compress the entire raw K/V cache into packed format in one batch. + /// Called once when transitioning from prefill to decode. + /// + /// In rawKeyMode: only compress values. Keys stay as raw FP16 in rawKeys buffer. + /// This is the highest-quality TurboQuant+ mode, K precision dominates quality. + private func compressRawCache() { + guard !isCompressed, let rk = rawKeys, let rv = rawValues, offset > 0 else { return } + let allKeys = rk[.ellipsis, .. compressedAllocSteps { + let newAlloc = ((prev + numSteps + step - 1) / step) * step + let newKW = MLXArray.zeros([B, H, newAlloc, kw.dim(-1)], dtype: kw.dtype) + let newKS = MLXArray.zeros([B, H, newAlloc, ks.dim(-1)], dtype: ks.dtype) + let newKB = MLXArray.zeros([B, H, newAlloc, kb.dim(-1)], dtype: kb.dtype) + let newVP = MLXArray.zeros([B, H, newAlloc, vpw], dtype: .uint32) + let newVN = MLXArray.zeros([B, H, newAlloc]) + if prev > 0 { + newKW[.ellipsis, .. rawAllocSteps { + let newAlloc = ((prev + numSteps + step - 1) / step) * step + let newRK = MLXArray.zeros([B, H, newAlloc, headDim], dtype: keys.dtype) + if prev > 0, let rk = rawKeys { + newRK[.ellipsis, .. compressedAllocSteps { + let newAlloc = ((prev + numSteps + step - 1) / step) * step + let newVP = MLXArray.zeros([B, H, newAlloc, vpw], dtype: .uint32) + let newVN = MLXArray.zeros([B, H, newAlloc]) + if prev > 0 { + newVP[.ellipsis, .. compressedAllocSteps { + let newAlloc = ((prev + numSteps + step - 1) / step) * step + let newKP = MLXArray.zeros([B, H, newAlloc, kpw], dtype: .uint32) + let newKN = MLXArray.zeros([B, H, newAlloc]) + let newVP = MLXArray.zeros([B, H, newAlloc, vpw], dtype: .uint32) + let newVN = MLXArray.zeros([B, H, newAlloc]) + if prev > 0 { + newKP[.ellipsis, ..1 (prefill chunks): falls back to separate score → softmax → value kernels + /// since causal masking across multiple query positions requires the full score matrix. + /// + /// In rawKeyMode: uses standard matmul for Q*K scoring (raw FP16 keys, no rotation), + /// then compressed-domain Metal kernel for Attn*V (TurboQuant compressed values). + /// TurboFlash is NOT used, it assumes both K and V are packed. + public func compressedAttention( + queries: MLXArray, + keys newKeys: MLXArray, + values newValues: MLXArray, + scale: Float, + mask: MLXFast.ScaledDotProductAttentionMaskMode = .none + ) -> MLXArray { + let headDim = newKeys.dim(-1) + let B = queries.dim(0) + let nQHeads = queries.dim(1) + let nKVHeads = newKeys.dim(1) + let L = queries.dim(2) + let nRepeats = nQHeads / nKVHeads + + // Transition: compress raw cache on first decode call + if !isCompressed { + compressRawCache() + } + + // Phase A: Encode new token + encodeNewToken(keys: newKeys, values: newValues) + + guard let valueMSECodec else { + return queries + } + + let tokenCount = offset + + // Shared V slicing (used by all paths) + let flatValPacked = valPackedMSE![0..., 0..., .. 1 { + q = q.reshaped([B, nKVHeads, nRepeats, L, headDim]) + kwS = expandedDimensions(kwS, axis: -3) + ksS = expandedDimensions(ksS, axis: -3) + kbS = expandedDimensions(kbS, axis: -3) + } + scores = quantizedMM( + q, kwS, scales: ksS, biases: kbS, + transpose: true, groupSize: keyGroupSize, bits: 8) + if nRepeats > 1 { + scores = scores.reshaped([B, nQHeads, L, tokenCount]) + } + } else { + // Q*K scoring: standard matmul with raw FP16 keys + guard let rk = rawKeys else { return queries } + let allKeys = rk[.ellipsis, .. 1 { + let expanded = expandedDimensions(allKeys, axis: 2) + let tiledKeys = MLX.tiled(expanded, repetitions: [1, 1, nRepeats, 1, 1]) + expandedKeys = tiledKeys.reshaped([B, nQHeads, tokenCount, headDim]) + } else { + expandedKeys = allKeys + } + + // scores = Q * K^T * scale → [B, nQHeads, L, T], f32: the f16 + // dot product overflows on outlier-heavy families (Qwen2.5 + // layer-0/last K amax ~200-400 → q·k beyond f16 max). + scores = + matmul( + queries.asType(.float32), + expandedKeys.asType(.float32).transposed(0, 1, 3, 2)) * MLXArray(scale) + } + + // Mask + softmax + switch mask { + case .array(let maskArray): + if maskArray.dtype == .bool { + scores = MLX.where(maskArray, scores, MLXArray(-Float.greatestFiniteMagnitude)) + } else { + scores = scores + maskArray + } + case .causal: + // Build causal mask manually + let queryOffset = tokenCount - L + let causalMask = MLXArray.tri(L, m: tokenCount, k: queryOffset, type: Bool.self) + let expandedMask = expandedDimensions( + expandedDimensions(causalMask, axis: 0), axis: 0) + scores = MLX.where(expandedMask, scores, MLXArray(-Float.greatestFiniteMagnitude)) + case .arrays(let maskArrays): + if let maskArray = maskArrays.first { + if maskArray.dtype == .bool { + scores = MLX.where( + maskArray, scores, MLXArray(-Float.greatestFiniteMagnitude)) + } else { + scores = scores + maskArray + } + } + case .none: break + } + + let attnWeights = softmax(scores, axis: -1) + + // Attn*V: compressed-domain Metal kernel for weighted sum of TurboQuant values + let flatWeights = attnWeights.reshaped([B * nQHeads * L, tokenCount]) + let rotatedOutput = TurboQuantKernelOps.mseWeightedSum( + weights: flatWeights, packed: flatValPacked, norms: flatValNorms, + codebook: valueMSECodec.codebook, tokenCount: tokenCount, + repeatCount: nRepeats, bits: self.valueBits, dim: headDim, queryChunkLength: L + ) + + // Inverse value rotation (V was encoded in rotated space) + output = matmul( + rotatedOutput.reshaped([B, nQHeads, L, headDim]), + valueMSECodec.rotation + ) + } else { + // ═══ Standard TurboQuant path: both K and V compressed ═══ + guard let keyMSECodec else { return queries } + + // Pre-rotate query for compressed-domain K scoring. When key + // calibration is active, fold s into the query side (elementwise, + // post-rotation): score = rotate(q)*s · rotate(k)/s = rotate(q)·rotate(k). + // The flash/score kernels stay untouched, they only ever see a + // plain dot product between rotatedQueries and the packed K codebook. + var qRot = keyMSECodec.prepareQueries(queries) * MLXArray(scale) + if let keyCalibScale { + qRot = qRot * keyCalibScale + } + let flatQ = qRot.reshaped([B * nQHeads * L, headDim]) + + // K slicing + let flatKeyPacked = keyPackedMSE![0..., 0..., ..1) + let queryOffset = tokenCount - L + output = TurboQuantKernelOps.turboFlashAttentionCausal( + rotatedQueries: flatQ, + keyPacked: flatKeyPacked, keyNorms: flatKeyNorms, + keyCodebook: keyMSECodec.codebook, + valPacked: flatValPacked, valNorms: flatValNorms, + valCodebook: valueMSECodec.codebook, + tokenCount: tokenCount, repeatCount: nRepeats, + keyBits: self.keyBits, valueBits: self.valueBits, dim: headDim, + queryChunkLength: L, queryOffset: queryOffset, + valRotation: valRotation + ).reshaped([B, nQHeads, L, headDim]) + } else { + // Separated path (L>1, non-causal masks) + var scores = TurboQuantKernelOps.mseScore( + rotatedQueries: flatQ, packed: flatKeyPacked, norms: flatKeyNorms, + codebook: keyMSECodec.codebook, tokenCount: tokenCount, + repeatCount: nRepeats, bits: self.keyBits, dim: headDim, queryChunkLength: L + ).reshaped([B, nQHeads, L, tokenCount]) + + // Mask + softmax + switch mask { + case .array(let maskArray): + if maskArray.dtype == .bool { + scores = MLX.where( + maskArray, scores, MLXArray(-Float.greatestFiniteMagnitude)) + } else { + scores = scores + maskArray + } + case .arrays(let maskArrays): + if let maskArray = maskArrays.first { + if maskArray.dtype == .bool { + scores = MLX.where( + maskArray, scores, MLXArray(-Float.greatestFiniteMagnitude)) + } else { + scores = scores + maskArray + } + } + case .none: break + default: break + } + + let attnWeights = softmax(scores, axis: -1) + + // Metal value kernel + let flatWeights = attnWeights.reshaped([B * nQHeads * L, tokenCount]) + let rotatedOutput = TurboQuantKernelOps.mseWeightedSum( + weights: flatWeights, packed: flatValPacked, norms: flatValNorms, + codebook: valueMSECodec.codebook, tokenCount: tokenCount, + repeatCount: nRepeats, bits: self.valueBits, dim: headDim, queryChunkLength: L + ) + + // Inverse rotation + output = matmul( + rotatedOutput.reshaped([B, nQHeads, L, headDim]), + valueMSECodec.rotation + ) + } + } + + // Kernels emit f32; return in the activation dtype so a turbo layer + // does not promote the whole downstream stream (and every later + // layer's KV cache) to f32. + return output.dtype == queries.dtype ? output : output.asType(queries.dtype) + } + + // MARK: - Memory Reporting + + /// Actual memory footprint: compressed storage (packed indices + norms) for K and V, + /// plus any raw FP16 buffers if still in prefill phase or rawKeyMode. + /// Does NOT include codec overhead (rotation matrices, codebooks) which is shared across layers. + /// In rawKeyMode: rawKeys is always present (FP16 keys), no keyPackedMSE/keyNorms. + public var memoryBytes: Int { + var total = 0 + // Raw FP16 buffers (always present in rawKeyMode for keys, or during prefill) + if let rk = rawKeys { total += rk.shape.reduce(1, *) * rk.dtype.bytesPerElement } + if let rv = rawValues { total += rv.shape.reduce(1, *) * rv.dtype.bytesPerElement } + // Compressed storage (K only present when NOT rawKeyMode) + if let kw = affKeyW { total += kw.shape.reduce(1, *) * kw.dtype.bytesPerElement } + if let ks = affKeyScales { total += ks.shape.reduce(1, *) * ks.dtype.bytesPerElement } + if let kb = affKeyBiases { total += kb.shape.reduce(1, *) * kb.dtype.bytesPerElement } + if let kp = keyPackedMSE { total += kp.shape.reduce(1, *) * kp.dtype.bytesPerElement } + if let kn = keyNorms { total += kn.shape.reduce(1, *) * kn.dtype.bytesPerElement } + if let vp = valPackedMSE { total += vp.shape.reduce(1, *) * vp.dtype.bytesPerElement } + if let vn = valNorms { total += vn.shape.reduce(1, *) * vn.dtype.bytesPerElement } + if let kcs = keyCalibScale { total += kcs.shape.reduce(1, *) * kcs.dtype.bytesPerElement } + return total + } + + // MARK: - State / Trim + + override public var state: [MLXArray] { + get { + if isCompressed { + if affineKeyMode { + // Affine-K compressed: [kW, kScales, kBiases, valPacked, valNorms] + guard let kw = affKeyW, let ks = affKeyScales, let kb = affKeyBiases, + let vpm = valPackedMSE, let vn = valNorms, + offset > 0 + else { return [] } + return [ + kw[0..., 0..., .. 0 + else { return [] } + return [ + rk[0..., 0..., .. 0 + else { return [] } + var arrays = [ + kpm[0..., 0..., .. 0 else { return [] } + return [rk[0..., 0..., ..= 5, + let o = Int(newValue[0]) + else { return } + offset = o + } + } + + @discardableResult + override public func trim(_ n: Int) -> Int { + guard n > 0, offset > 0 else { return 0 } + let trimCount = min(n, offset) + offset -= trimCount + if offset == 0 { + rawKeys = nil + rawValues = nil + rawAllocSteps = 0 + keyPackedMSE = nil + keyNorms = nil + affKeyW = nil + affKeyScales = nil + affKeyBiases = nil + valPackedMSE = nil + valNorms = nil + keyCalibScale = nil + compressedAllocSteps = 0 + isCompressed = false + } + return trimCount + } +} + +// MARK: - kvScheme routing + +/// Resolve a TurboQuant kvScheme string to (keyBits, valueBits). +/// `keyBits == 0` means raw FP16 keys (only values are compressed). +/// Returns nil for non-TurboQuant schemes. +public func resolveTurboScheme(_ scheme: String?) -> (keyBits: Int, valueBits: Int)? { + switch scheme { + // Raw-key schemes: keys stay FP16, values are group-quantized. The + // teacher-forced PPL/KLD gate shows these are near-lossless (see the + // PR validation table); they are the recommended configurations. + case "turbo0v4": return (0, 4) + case "turbo0v3": return (0, 3) + case "turbo0v2": return (0, 2) + // Affine-K asymmetric family (keyBits == 8 → 8-bit affine keys, turbo + // values): the recommended ladder from the asymmetric-kv paper. q8-K is + // near-lossless at half the key bytes; f16-K buys nothing further. + case "turbo8v4": return (8, 4) + case "turbo8v3": return (8, 3) + case "turbo8v2": return (8, 2) + // Key-quantizing (symmetric) schemes: maximum compression. K sensitivity + // varies by model family, validate on your model; the asym family above + // is the recommended starting point. + case "turbo4": return (4, 4) + case "turbo4v2": return (4, 2) + case "turbo3": return (3, 3) + case "turbo2": return (2, 2) + + default: return nil + } +} + +/// Namespace holding the dedup state for the "rotating layers kept fp16" +/// notice. Mirrors the `nonisolated(unsafe) static var` + `NSLock` idiom +/// used elsewhere in this file (e.g. `TurboQuantCodebook.precomputed`). +private enum RotatingSkipNotice { + /// Once per process: the notice is informational and repeating it per + /// generation adds noise without new signal. Tracking instances instead + /// would grow without bound, `newCache(parameters:)` allocates fresh + /// caches for every generation. + nonisolated(unsafe) static var logged = false + static let lock = NSLock() +} + +/// Log, at most once per set of `RotatingKVCache` instances, that a +/// requested TurboQuant kvScheme is leaving sliding-window layers at fp16. +/// +/// TurboQuant only compresses `KVCacheSimple` layers (see +/// `maybeTurboQuantizeKVCache` below); `RotatingKVCache` (Gemma-style +/// sliding-window) layers are already memory-bounded by their window size, +/// and their rotating/eviction storage layout is not compatible with the +/// sequential-append compression path, so they are intentionally left +/// untouched. The notice tells a user which layers kept fp16 rotating +/// caches so a partially engaged scheme is visible. +private func logRotatingKVCacheSkipOnce(cache: [KVCache]) { + let rotatingIndices = cache.indices.filter { cache[$0] is RotatingKVCache } + guard !rotatingIndices.isEmpty else { return } + + RotatingSkipNotice.lock.lock() + defer { RotatingSkipNotice.lock.unlock() } + + guard !RotatingSkipNotice.logged else { return } + RotatingSkipNotice.logged = true + + let indexList = rotatingIndices.map(String.init).joined(separator: ", ") + print( + "[TurboQuant] kvScheme requested KV compression, but layer(s) at index \(indexList) " + + "use RotatingKVCache (sliding-window) and will stay fp16. TurboQuant only " + + "compresses non-rotating (global) KV cache layers." + ) +} + +/// Convert eligible `KVCacheSimple` layers to `TurboQuantKVCache` once their +/// offset passes `quantizedKVStart`, transferring the accumulated KV state. +/// Mamba / wrapper caches are left untouched. `RotatingKVCache` layers +/// (Gemma-style sliding-window) are also left untouched -- see +/// `logRotatingKVCacheSkipOnce` -- since they are already memory-bounded by +/// their window size and their rotating storage layout does not fit the +/// sequential-append compression path. Called from `maybeQuantizeKVCache` +/// when `kvScheme` names a TurboQuant scheme. +public func maybeTurboQuantizeKVCache( + cache: inout [KVCache], + keyBits: Int, + valueBits: Int, + quantizedKVStart: Int +) { + logRotatingKVCacheSkipOnce(cache: cache) + + guard + cache.contains(where: { $0 is KVCacheSimple && $0.offset > quantizedKVStart }) + else { return } + + // Boundary layer protection: the first and last attention layers are + // disproportionately sensitive to KV quantization, keeping 2 on each + // end at FP16 recovers 37-91% of the quality gap at minimal + // compression cost. Non-attention layers (Mamba/wrapper caches) are + // excluded from the rank count so hybrids protect the right layers. + let kvLayerIndices = cache.indices.filter { cache[$0] is KVCacheSimple } + // Only engage when the scheme is actually fragile: turbo-quantized keys + // (keyBits 2-4) or 2-bit values. Raw/affine-8 keys with >=3-bit values + // are near-lossless everywhere and protection just costs compression. + let fragile = (keyBits > 0 && keyBits < 8) || valueBits <= 2 + let boundaryLayers = fragile ? min(2, kvLayerIndices.count / 2) : 0 + let protected = Set( + kvLayerIndices.prefix(boundaryLayers) + kvLayerIndices.suffix(boundaryLayers)) + + for i in 0 ..< cache.count { + guard let simple = cache[i] as? KVCacheSimple, cache[i].offset > quantizedKVStart, + !(cache[i] is TurboQuantKVCache) + else { continue } + + let state = cache[i].innerState() + let headDims: (key: Int, value: Int)? = + state.count >= 2 ? (state[0].dim(3), state[1].dim(3)) : nil + + if protected.contains(i) { + // Boundary layers use 8-bit affine instead of the turbo scheme: + // near-lossless protection for the quantization-sensitive first + // and last layers at a quarter of the fp16 cost. Skip the + // conversion (leave fp16) rather than crash when neither head + // dimension divides one of the supported affine group sizes. + guard + let headDims, + resolvedKVQuantizationGroupSize( + requested: 64, keyHeadDim: headDims.key, valueHeadDim: headDims.value) != nil + else { continue } + cache[i] = simple.toQuantized(groupSize: 64, bits: 8) + continue + } + + // Affine-K mode (keyBits == 8) quantizes keys in groups, resolve the + // group size against this layer's head dimension up front so the + // TurboQuantKVCache instance never has to fall back at first encode. + // Skip layers whose head dimension divides none of the supported + // sizes rather than constructing a cache that would only crash on + // its first compress. Other modes (raw / symmetric turbo) don't + // group-quantize keys, so they carry the unused default unchanged. + var resolvedKeyGroupSize = 64 + if keyBits == 8 { + guard + let headDims, + let resolved = resolvedKVQuantizationGroupSize( + requested: 64, keyHeadDim: headDims.key, valueHeadDim: headDims.value) + else { continue } + resolvedKeyGroupSize = resolved + } + + let turbo = TurboQuantKVCache( + bits: max(keyBits, valueBits), keyBits: keyBits, valueBits: valueBits, + keyGroupSize: resolvedKeyGroupSize) + // Transfer existing KV data, trimmed to the live offset (the simple + // cache over-allocates in steps). + let offset = cache[i].offset + if state.count >= 2, offset > 0 { + let keys = state[0][.ellipsis, ..> shift); + + // Handle bits that spill across uint32 word boundary + int spill = (int)shift + (int)Bits - 32; + if (spill > 0) { + value |= (packed_ptr[word_idx + 1] << ((uint)Bits - (uint)spill)); + } + value &= MASK; + + // Codebook lookup + accumulate dot product + acc += q_ptr[d] * cb[value]; + } + + // SIMD reduction across 32 lanes + acc = simd_sum(acc); + + // Lane 0 writes final score (scaled by stored norm) + if (thread_index_in_simdgroup == 0) { + scores[q_idx * token_count + k_idx] = acc * norm_val; + } + """ + + /// Fused encode kernel: norm + rotate + quantize + pack + norm correction in ONE dispatch. + /// + /// For each input vector [D]: + /// 1. Compute L2 norm (SIMD reduction) + /// 2. Normalize to unit vector + /// 3. Rotate: y = Π · x_unit (shared memory matmul) + /// 4. Quantize: find codebook index via boundary comparison + /// 5. Pack bits into uint32 words (atomic OR) + /// 6. Norm correction: compute reconstruction norm, store original_norm / recon_norm + /// + /// Norm correction compensates for quantization error so that + /// centroid[idx] * corrected_norm more accurately reconstructs the original vector. + /// This is why TurboQuant beats q8_0 on perplexity. + /// + /// Grid: (Dim, numRows, 1), one threadgroup per vector + /// Threadgroup: (Dim, 1, 1), all D threads cooperate + /// + /// Template params: Bits, Dim, PackedWidth, NumBoundaries (= 2^Bits - 1) + static let fusedEncodeSource = """ + constexpr uint LEVELS = 1u << Bits; + + uint d = thread_position_in_threadgroup.x; // dimension index (0..Dim-1) + uint row = thread_position_in_grid.y; // vector index (B*H*T) + + // --- Step 1: Load input value --- + float val = input[row * Dim + d]; + + // --- Step 2: Compute L2 norm (SIMD reduction) --- + float sq = val * val; + float norm_sq = simd_sum(sq); + // For Dim > 32, need threadgroup reduction + threadgroup float shared_norm[32]; // up to 32 SIMD groups (dim <= 1024) + uint sg_id = d / 32; + if (d % 32 == 0) { + shared_norm[sg_id] = norm_sq; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float total_norm_sq = 0; + uint num_groups = (Dim + 31) / 32; + for (uint i = 0; i < num_groups; i++) { + total_norm_sq += shared_norm[i]; + } + float norm_val = sqrt(total_norm_sq); + float inv_norm = (norm_val > 1e-8f) ? (1.0f / norm_val) : 0.0f; + + // --- Step 3: Normalize --- + float unit_val = val * inv_norm; + + // --- Step 4: Rotate (y = Π · x_unit) via shared memory matmul --- + // Each thread d computes: y[d] = Σ_j rotation[d * Dim + j] * x_unit[j] + threadgroup float shared_unit[1024]; // max Dim = 1024 + shared_unit[d] = unit_val; + threadgroup_barrier(mem_flags::mem_threadgroup); + + float rotated = 0.0f; + for (uint j = 0; j < Dim; j++) { + rotated += rotation[d * Dim + j] * shared_unit[j]; + } + + // --- Step 5: Quantize via branchless boundary comparison --- + // V2.1 optimization: use arithmetic sum of comparisons instead of branching. + // Metal compiles (rotated > boundaries[b]) to a predicated 0/1, summing these + // is branchless and avoids SIMD lane divergence. + uint idx = 0; + for (uint b = 0; b < LEVELS - 1; b++) { + idx += (uint)(rotated > boundaries[b]); + } + + // --- Step 6: Pack bits into uint32 word (atomic OR) --- + uint bit_offset = d * Bits; + uint word_idx = bit_offset / 32; + uint shift = bit_offset % 32; + uint masked = idx & ((1u << Bits) - 1u); + + // Pack bits, use threadgroup shared memory to avoid atomic contention + // Each thread writes its index bits to shared, then thread 0 per word writes output + threadgroup uint shared_packed[64]; // max PackedWidth = 64 words + if (d < PackedWidth) shared_packed[d] = 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Each dimension contributes its bits via atomic OR on threadgroup memory + uint primary_val = masked << shift; + atomic_fetch_or_explicit((threadgroup atomic_uint*)&shared_packed[word_idx], primary_val, memory_order_relaxed); + + int spill = (int)shift + (int)Bits - 32; + if (spill > 0) { + uint spill_val = masked >> ((uint)Bits - (uint)spill); + atomic_fetch_or_explicit((threadgroup atomic_uint*)&shared_packed[word_idx + 1], spill_val, memory_order_relaxed); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Write packed words to output (one thread per word) + if (d < PackedWidth) { + packed_out[row * PackedWidth + d] = shared_packed[d]; + } + + // --- Step 7: Norm correction --- + // Compute reconstruction norm: ||codebook[idx]||₂ for the quantized unit vector. + // Store corrected_norm = original_norm / recon_norm so that + // decode(centroid[idx] * corrected_norm) better approximates the original vector. + float centroid_val = codebook[idx]; + float recon_sq = centroid_val * centroid_val; + float recon_norm_sq = simd_sum(recon_sq); + // Threadgroup reduction for Dim > 32 + if (d % 32 == 0) { + shared_norm[sg_id] = recon_norm_sq; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float total_recon_sq = 0; + for (uint i = 0; i < num_groups; i++) { + total_recon_sq += shared_norm[i]; + } + float recon_norm = sqrt(total_recon_sq); + float corrected_norm = (recon_norm > 1e-8f) ? (norm_val / recon_norm) : norm_val; + + if (d == 0) { + norms_out[row] = corrected_norm; + } + """ + + /// Fused WHT encode kernel: norm + WHT rotation + quantize + pack (NO norm correction). + /// + /// Same as fusedEncodeSource but replaces dense O(d²) matmul with Fast Walsh-Hadamard + /// Transform O(d log d) butterfly + random sign flip. 18× fewer ops for dim=128. + /// + /// Stores RAW norms (no reconstruction-norm correction), a measured choice, + /// not an omission: for values consumed by attention's weighted sums, raw + /// norms beat both matched-norm (attention cos 0.976 vs ≥0.980) and + /// projection-optimal (0.9798) rescaling at 4-bit; per-vector rescales bias + /// the decoded vectors toward zero while raw-norm errors stay unbiased and + /// average out across tokens. The dense fallback path keeps matched-norm + /// for canonical-formulation parity. + /// + /// WHT forward rotation: y = WHT(signs * x_unit) / sqrt(Dim) + /// The butterfly pattern: for each stage s in 0.. 1e-8f) ? (1.0f / norm_val) : 0.0f; + + // --- Step 3: Normalize + sign flip (fused) --- + // V2.1 optimization: pre-compute inv_norm * sign to eliminate one multiply per element. + // Instead of: unit_val = val * inv_norm; wht_val = sign * unit_val (2 muls) + // We do: wht_val = val * (inv_norm * sign) (1 mul + 1 FMA-friendly product) + float inv_norm_sign = inv_norm * wht_signs[d]; + float wht_val = val * inv_norm_sign; + + // --- Step 4: WHT rotation via cooperative SIMD shuffle --- + // V2.1 optimization: use simd_shuffle_xor for intra-SIMD butterfly stages + // (register-to-register, no shared memory or barriers needed for first 5 stages) + + // Phase 1: Intra-SIMD butterfly via simd_shuffle_xor (stages 0..min(LogDim,5)-1) + // Each stage s XORs lane indices at distance 2^s, effectively free on Apple GPU + uint log_dim_u = uint(LogDim); + uint simd_stages = min(log_dim_u, 5u); // 5 stages covers 32 lanes (2^5 = 32) + uint lane_in_simd = d % 32; + for (uint s = 0; s < simd_stages; s++) { + uint step = 1u << s; + float other = simd_shuffle_xor(wht_val, step); + wht_val = (lane_in_simd & step) ? (other - wht_val) : (other + wht_val); + } + + // Phase 2: Cross-SIMD-group butterfly via shared memory (stages 5..LogDim-1) + // Only needed when Dim > 32, these stages cross SIMD group boundaries + threadgroup float shared_buf[1024]; // max Dim = 1024 + if (log_dim_u > 5u) { + shared_buf[d] = wht_val; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint s = simd_stages; s < log_dim_u; s++) { + uint half_block = 1u << s; + uint block_size = half_block << 1; + uint block_id = d / block_size; + uint pos_in_block = d % block_size; + + float a, b; + if (pos_in_block < half_block) { + a = shared_buf[block_id * block_size + pos_in_block]; + b = shared_buf[block_id * block_size + pos_in_block + half_block]; + shared_buf[d] = a + b; + } else { + a = shared_buf[block_id * block_size + pos_in_block - half_block]; + b = shared_buf[block_id * block_size + pos_in_block]; + shared_buf[d] = a - b; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + wht_val = shared_buf[d]; + } + + // Normalize: WHT has scale factor sqrt(Dim) + float inv_sqrt_dim = 1.0f / sqrt((float)Dim); + float rotated = wht_val * inv_sqrt_dim; + + // --- Step 5: Quantize via branchless boundary comparison --- + // V2.1 optimization: arithmetic sum avoids SIMD lane divergence + uint idx = 0; + for (uint b = 0; b < LEVELS - 1; b++) { + idx += (uint)(rotated > boundaries[b]); + } + + // --- Step 6: Pack bits into uint32 word (atomic OR) --- + uint bit_offset = d * Bits; + uint word_idx = bit_offset / 32; + uint shift = bit_offset % 32; + uint masked = idx & ((1u << Bits) - 1u); + + threadgroup uint shared_packed[64]; + if (d < PackedWidth) shared_packed[d] = 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint primary_val = masked << shift; + atomic_fetch_or_explicit((threadgroup atomic_uint*)&shared_packed[word_idx], primary_val, memory_order_relaxed); + + int spill = (int)shift + (int)Bits - 32; + if (spill > 0) { + uint spill_val = masked >> ((uint)Bits - (uint)spill); + atomic_fetch_or_explicit((threadgroup atomic_uint*)&shared_packed[word_idx + 1], spill_val, memory_order_relaxed); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (d < PackedWidth) { + packed_out[row * PackedWidth + d] = shared_packed[d]; + } + + // --- Step 7: Store raw norm (WHT is orthogonal, no norm correction needed) --- + // WHT preserves norms: ||WHT(x)||₂ = ||x||₂. Reconstruction norm ≈ original norm, + // so the correction ratio ≈ 1.0. Skipping saves codebook lookup + norm + division. + if (d == 0) { + norms_out[row] = norm_val; + } + """ + + /// Quantize, pack, and norm-correct pre-rotated calibrated vectors. + /// + /// The rotation stays in MLX ops (MSECodec.rotatedUnit) so kernel and + /// codec quantize bit-identical values; this kernel only does the parts + /// that are slow as MLX ops: the boundary compare (which broadcasts a + /// [rows, dim, levels] intermediate) and low-bit packing. The stored + /// norm is norms_in / ||codebook[idx] * scale|| per row, matching + /// MSECodec.encode(_:scale:). + /// + /// Template params: Bits, Dim, PackedWidth. + static let fusedQuantizePackScaledSource = """ + constexpr uint LEVELS = 1u << Bits; + + uint d = thread_position_in_threadgroup.x; // dimension index (0..Dim-1) + uint row = thread_position_in_grid.y; // vector index (B*H*T) + + float calibrated = rotated[row * Dim + d] / scale[d]; + uint idx = 0; + for (uint b = 0; b < LEVELS - 1; b++) { + idx += (uint)(calibrated > boundaries[b]); + } + + uint bit_offset = d * Bits; + uint word_idx = bit_offset / 32; + uint shift = bit_offset % 32; + uint masked = idx & ((1u << Bits) - 1u); + + threadgroup uint shared_packed[64]; + if (d < PackedWidth) shared_packed[d] = 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint primary_val = masked << shift; + atomic_fetch_or_explicit((threadgroup atomic_uint*)&shared_packed[word_idx], primary_val, memory_order_relaxed); + + int spill = (int)shift + (int)Bits - 32; + if (spill > 0) { + uint spill_val = masked >> ((uint)Bits - (uint)spill); + atomic_fetch_or_explicit((threadgroup atomic_uint*)&shared_packed[word_idx + 1], spill_val, memory_order_relaxed); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (d < PackedWidth) { + packed_out[row * PackedWidth + d] = shared_packed[d]; + } + + // Norm correction against the calibrated reconstruction. + float centroid_val = codebook[idx] * scale[d]; + float recon_sq = centroid_val * centroid_val; + float recon_norm_sq = simd_sum(recon_sq); + threadgroup float shared_norm[32]; + uint sg_id = d / 32; + if (d % 32 == 0) { + shared_norm[sg_id] = recon_norm_sq; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float total_recon_sq = 0; + uint num_groups = (Dim + 31) / 32; + for (uint i = 0; i < num_groups; i++) { + total_recon_sq += shared_norm[i]; + } + float recon_norm = sqrt(total_recon_sq); + + if (d == 0) { + float raw_norm = norms_in[row]; + norms_out[row] = (recon_norm > 1e-8f) ? (raw_norm / recon_norm) : raw_norm; + } + """ + + /// TurboFlashAttention Pass 1: Per-block partial attention with online softmax. + /// + /// Parallelizes across both query heads AND token blocks. Each SIMD group (32 lanes) + /// handles one (query, block) pair, producing partial online softmax state (m, l, o[D]). + /// Pass 2 merges partials across blocks. + /// + /// Grid: (32, totalQueries, numBlocks) + /// Threadgroup: (32, 1, 1) + /// + /// Template params: KeyBits, ValueBits, Dim, KeyPackedWidth, ValuePackedWidth, + /// BlockSize, token_count, repeat_count, num_blocks + static let turboFlashPass1Source = """ + // Per-call values arrive via the `params` input buffer, NOT as template + // args: template values are baked into the kernel name, so a varying + // token_count would JIT-compile and permanently cache one Metal + // library per generated token (unbounded host-memory growth + a + // shader compile every step). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint num_blocks = params[2]; + constexpr uint KEY_MASK = (1u << KeyBits) - 1u; + constexpr uint KEY_LEVELS = 1u << KeyBits; + constexpr uint VAL_MASK = (1u << ValueBits) - 1u; + constexpr uint VAL_LEVELS = 1u << ValueBits; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; // SIMD lane (0-31) + uint q_idx = thread_position_in_grid.y; // query index (B*nQHeads*L) + uint block_idx = thread_position_in_grid.z; // token block index + uint kv_idx = q_idx / repeat_count; // map to KV head (GQA; L=1 only) + + // Token range for this block + uint t_start = block_idx * BlockSize; + uint t_end = t_start + BlockSize; + if (t_end > (uint)token_count) t_end = (uint)token_count; + + // Load key codebook into registers + float key_cb[KEY_LEVELS]; + for (uint i = 0; i < KEY_LEVELS; i++) { + key_cb[i] = key_codebook[i]; + } + + // Load value codebook into registers + float val_cb[VAL_LEVELS]; + for (uint i = 0; i < VAL_LEVELS; i++) { + val_cb[i] = val_codebook[i]; + } + + // Load query values for this lane's dimensions + float q_vals[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + q_vals[i] = (d < Dim) ? q_rot[q_idx * Dim + d] : 0.0f; + } + + // Online softmax state for this block + float m = -INFINITY; + float l = 0.0f; + float o[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) o[i] = 0.0f; + + // Process tokens in this block + for (uint t = t_start; t < t_end; t++) { + // --- Score: Q×K dot product --- + const device uint32_t* k_packed_ptr = key_packed + kv_idx * token_count * KeyPackedWidth + t * KeyPackedWidth; + float k_norm = key_norms[kv_idx * token_count + t]; + + float dot_partial = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + + uint k_bit_offset = d * KeyBits; + uint k_word_idx = k_bit_offset / 32; + uint k_shift = k_bit_offset % 32; + uint k_value = (k_packed_ptr[k_word_idx] >> k_shift); + int k_spill = (int)k_shift + (int)KeyBits - 32; + if (k_spill > 0) { + k_value |= (k_packed_ptr[k_word_idx + 1] << ((uint)KeyBits - (uint)k_spill)); + } + k_value &= KEY_MASK; + + dot_partial += q_vals[i] * key_cb[k_value]; + } + + float score = simd_sum(dot_partial) * k_norm; + + // --- Online softmax update + V accumulation --- + float new_m = max(m, score); + float exp_diff = exp(m - new_m); + float exp_score = exp(score - new_m); + + const device uint32_t* v_packed_ptr = val_packed + kv_idx * token_count * ValuePackedWidth + t * ValuePackedWidth; + float v_norm = val_norms[kv_idx * token_count + t]; + + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + + uint v_bit_offset = d * ValueBits; + uint v_word_idx = v_bit_offset / 32; + uint v_shift = v_bit_offset % 32; + uint v_value = (v_packed_ptr[v_word_idx] >> v_shift); + int v_spill = (int)v_shift + (int)ValueBits - 32; + if (v_spill > 0) { + v_value |= (v_packed_ptr[v_word_idx + 1] << ((uint)ValueBits - (uint)v_spill)); + } + v_value &= VAL_MASK; + + o[i] = o[i] * exp_diff + exp_score * (val_cb[v_value] * v_norm); + } + + l = l * exp_diff + exp_score; + m = new_m; + } + + // Write partial results: o[D], m, l + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o_partials[partial_base + d] = o[i]; + } + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = m; + l_partials[ml_idx] = l; + } + """ + + /// TurboFlashAttention Pass 1 (Causal): Per-block partial attention with causal masking. + /// + /// Same as turboFlashPass1Source but supports L>1 query chunks with causal masking. + /// Each query position q_within_L only attends to tokens where t <= q_offset + q_within_L. + /// Blocks that are entirely future-masked exit early. + /// + /// Grid: (32, totalQueries, numBlocks) where totalQueries = B * nQHeads * L + /// Threadgroup: (32, 1, 1) + /// + /// Additional template params: L (query chunk length), q_offset (absolute offset of first query) + /// Pass-1 variant: raw K scoring (KT template, matches the stored K + /// cache dtype), turbo-V accumulation. Same SIMD-parallel online-softmax + /// structure as the packed-K flash kernel, so the asymmetric family + /// gets the same decode throughput. + static let turboFlashPass1RawKSource = """ + // Per-call values arrive via the `params` input buffer, NOT as template + // args: template values are baked into the kernel name, so a varying + // token_count would JIT-compile and permanently cache one Metal + // library per generated token (unbounded host-memory growth + a + // shader compile every step). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint num_blocks = params[2]; + constexpr uint VAL_MASK = (1u << ValueBits) - 1u; + constexpr uint VAL_LEVELS = 1u << ValueBits; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; // SIMD lane (0-31) + uint q_idx = thread_position_in_grid.y; // query index (B*nQHeads*L) + uint block_idx = thread_position_in_grid.z; // token block index + uint kv_idx = q_idx / repeat_count; // map to KV head (GQA; L=1 only) + + // Token range for this block + uint t_start = block_idx * BlockSize; + uint t_end = t_start + BlockSize; + if (t_end > (uint)token_count) t_end = (uint)token_count; + + // Load value codebook into registers + float val_cb[VAL_LEVELS]; + for (uint i = 0; i < VAL_LEVELS; i++) { + val_cb[i] = val_codebook[i]; + } + + // Load query values for this lane's dimensions + float q_vals[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + q_vals[i] = (d < Dim) ? q_rot[q_idx * Dim + d] : 0.0f; + } + + // Online softmax state for this block + float m = -INFINITY; + float l = 0.0f; + float o[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) o[i] = 0.0f; + + // Process tokens in this block + for (uint t = t_start; t < t_end; t++) { + // --- Score: Q×K, K read raw f16 (no unpack, no rotation) --- + const device KT* k_raw_ptr = (const device KT*)k_raw + kv_idx * token_count * Dim + t * Dim; + float dot_partial = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + dot_partial += q_vals[i] * float(k_raw_ptr[d]); + } + float score = simd_sum(dot_partial); + + // --- Online softmax update + V accumulation --- + float new_m = max(m, score); + float exp_diff = exp(m - new_m); + float exp_score = exp(score - new_m); + + const device uint32_t* v_packed_ptr = val_packed + kv_idx * token_count * ValuePackedWidth + t * ValuePackedWidth; + float v_norm = val_norms[kv_idx * token_count + t]; + + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + + uint v_bit_offset = d * ValueBits; + uint v_word_idx = v_bit_offset / 32; + uint v_shift = v_bit_offset % 32; + uint v_value = (v_packed_ptr[v_word_idx] >> v_shift); + int v_spill = (int)v_shift + (int)ValueBits - 32; + if (v_spill > 0) { + v_value |= (v_packed_ptr[v_word_idx + 1] << ((uint)ValueBits - (uint)v_spill)); + } + v_value &= VAL_MASK; + + o[i] = o[i] * exp_diff + exp_score * (val_cb[v_value] * v_norm); + } + + l = l * exp_diff + exp_score; + m = new_m; + } + + // Write partial results: o[D], m, l + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o_partials[partial_base + d] = o[i]; + } + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = m; + l_partials[ml_idx] = l; + } + """ + + /// Pass-1 variant: 8-bit affine K scoring (inline dequant, scale/bias in + /// their native dtype via the KScaleT/KBiasT templates), turbo-V. + static let turboFlashPass1AffineKSource = """ + // Per-call values arrive via the `params` input buffer, NOT as template + // args: template values are baked into the kernel name, so a varying + // token_count would JIT-compile and permanently cache one Metal + // library per generated token (unbounded host-memory growth + a + // shader compile every step). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint num_blocks = params[2]; + constexpr uint VAL_MASK = (1u << ValueBits) - 1u; + constexpr uint VAL_LEVELS = 1u << ValueBits; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; // SIMD lane (0-31) + uint q_idx = thread_position_in_grid.y; // query index (B*nQHeads*L) + uint block_idx = thread_position_in_grid.z; // token block index + uint kv_idx = q_idx / repeat_count; // map to KV head (GQA; L=1 only) + + // Token range for this block + uint t_start = block_idx * BlockSize; + uint t_end = t_start + BlockSize; + if (t_end > (uint)token_count) t_end = (uint)token_count; + + // Load value codebook into registers + float val_cb[VAL_LEVELS]; + for (uint i = 0; i < VAL_LEVELS; i++) { + val_cb[i] = val_codebook[i]; + } + + // Load query values for this lane's dimensions + float q_vals[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + q_vals[i] = (d < Dim) ? q_rot[q_idx * Dim + d] : 0.0f; + } + + // Online softmax state for this block + float m = -INFINITY; + float l = 0.0f; + float o[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) o[i] = 0.0f; + + // Process tokens in this block + for (uint t = t_start; t < t_end; t++) { + // --- Score: Q×K, K dequantized inline from 8-bit affine --- + const device uint* k_w_ptr = k_weights + (kv_idx * token_count + t) * (Dim / 4); + const device KScaleT* k_s_ptr = + (const device KScaleT*)k_scales + (kv_idx * token_count + t) * (Dim / KGroup); + const device KBiasT* k_b_ptr = + (const device KBiasT*)k_biases + (kv_idx * token_count + t) * (Dim / KGroup); + float dot_partial = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + uint g = d / KGroup; + uint w8 = (k_w_ptr[d / 4] >> ((d % 4) * 8)) & 0xFFu; + dot_partial += q_vals[i] * (float(k_s_ptr[g]) * float(w8) + float(k_b_ptr[g])); + } + float score = simd_sum(dot_partial); + + // --- Online softmax update + V accumulation --- + float new_m = max(m, score); + float exp_diff = exp(m - new_m); + float exp_score = exp(score - new_m); + + const device uint32_t* v_packed_ptr = val_packed + kv_idx * token_count * ValuePackedWidth + t * ValuePackedWidth; + float v_norm = val_norms[kv_idx * token_count + t]; + + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + + uint v_bit_offset = d * ValueBits; + uint v_word_idx = v_bit_offset / 32; + uint v_shift = v_bit_offset % 32; + uint v_value = (v_packed_ptr[v_word_idx] >> v_shift); + int v_spill = (int)v_shift + (int)ValueBits - 32; + if (v_spill > 0) { + v_value |= (v_packed_ptr[v_word_idx + 1] << ((uint)ValueBits - (uint)v_spill)); + } + v_value &= VAL_MASK; + + o[i] = o[i] * exp_diff + exp_score * (val_cb[v_value] * v_norm); + } + + l = l * exp_diff + exp_score; + m = new_m; + } + + // Write partial results: o[D], m, l + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o_partials[partial_base + d] = o[i]; + } + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = m; + l_partials[ml_idx] = l; + } + """ + + static let turboFlashPass1CausalSource = """ + // Per-call values arrive via the `params` input buffer, NOT as template + // args: template values are baked into the kernel name, so a varying + // token_count would JIT-compile and permanently cache one Metal + // library per generated token (unbounded host-memory growth + a + // shader compile every step). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint num_blocks = params[2]; + const uint L = params[3]; + const uint q_offset = params[4]; + constexpr uint KEY_MASK = (1u << KeyBits) - 1u; + constexpr uint KEY_LEVELS = 1u << KeyBits; + constexpr uint VAL_MASK = (1u << ValueBits) - 1u; + constexpr uint VAL_LEVELS = 1u << ValueBits; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; // SIMD lane (0-31) + uint q_idx = thread_position_in_grid.y; // query index (B*nQHeads*L) + uint block_idx = thread_position_in_grid.z; // token block index + + // For L>1, queries are laid out as [B * nQHeads * L, D] from reshape of [B, nQHeads, L, D]. + // q_idx = b * (nQHeads * L) + h * L + l + // We need: l (position within chunk) and kv_head (for GQA mapping). + uint q_within_L = q_idx % L; + uint q_head_idx = q_idx / L; // index into [B * nQHeads] + uint kv_idx = q_head_idx / repeat_count; // map to KV head (GQA) + + // Causal boundary: this query can attend to tokens 0..q_abs (inclusive) + uint q_abs = q_offset + q_within_L; + + // Token range for this block + uint t_start = block_idx * BlockSize; + uint t_end = t_start + BlockSize; + if (t_end > (uint)token_count) t_end = (uint)token_count; + + // Early exit: entire block is future-masked + if (t_start > q_abs) { + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) o_partials[partial_base + d] = 0.0f; + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = -INFINITY; + l_partials[ml_idx] = 0.0f; + } + return; + } + + // Clamp t_end to causal boundary + if (t_end > q_abs + 1) t_end = q_abs + 1; + + // Load key codebook into registers + float key_cb[KEY_LEVELS]; + for (uint i = 0; i < KEY_LEVELS; i++) { + key_cb[i] = key_codebook[i]; + } + + // Load value codebook into registers + float val_cb[VAL_LEVELS]; + for (uint i = 0; i < VAL_LEVELS; i++) { + val_cb[i] = val_codebook[i]; + } + + // Load query values for this lane's dimensions + float q_vals[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + q_vals[i] = (d < Dim) ? q_rot[q_idx * Dim + d] : 0.0f; + } + + // Online softmax state for this block + float m = -INFINITY; + float l = 0.0f; + float o[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) o[i] = 0.0f; + + // Process tokens in this block (up to causal boundary) + for (uint t = t_start; t < t_end; t++) { + // --- Score: Q×K dot product --- + const device uint32_t* k_packed_ptr = key_packed + kv_idx * token_count * KeyPackedWidth + t * KeyPackedWidth; + float k_norm = key_norms[kv_idx * token_count + t]; + + float dot_partial = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + + uint k_bit_offset = d * KeyBits; + uint k_word_idx = k_bit_offset / 32; + uint k_shift = k_bit_offset % 32; + uint k_value = (k_packed_ptr[k_word_idx] >> k_shift); + int k_spill = (int)k_shift + (int)KeyBits - 32; + if (k_spill > 0) { + k_value |= (k_packed_ptr[k_word_idx + 1] << ((uint)KeyBits - (uint)k_spill)); + } + k_value &= KEY_MASK; + + dot_partial += q_vals[i] * key_cb[k_value]; + } + + float score = simd_sum(dot_partial) * k_norm; + + // --- Online softmax update + V accumulation --- + float new_m = max(m, score); + float exp_diff = exp(m - new_m); + float exp_score = exp(score - new_m); + + const device uint32_t* v_packed_ptr = val_packed + kv_idx * token_count * ValuePackedWidth + t * ValuePackedWidth; + float v_norm = val_norms[kv_idx * token_count + t]; + + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + + uint v_bit_offset = d * ValueBits; + uint v_word_idx = v_bit_offset / 32; + uint v_shift = v_bit_offset % 32; + uint v_value = (v_packed_ptr[v_word_idx] >> v_shift); + int v_spill = (int)v_shift + (int)ValueBits - 32; + if (v_spill > 0) { + v_value |= (v_packed_ptr[v_word_idx + 1] << ((uint)ValueBits - (uint)v_spill)); + } + v_value &= VAL_MASK; + + o[i] = o[i] * exp_diff + exp_score * (val_cb[v_value] * v_norm); + } + + l = l * exp_diff + exp_score; + m = new_m; + } + + // Write partial results: o[D], m, l + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o_partials[partial_base + d] = o[i]; + } + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = m; + l_partials[ml_idx] = l; + } + """ + + /// TurboFlashAttention Pass 1 NR0=2: Multi-row amortized KV dequant. + /// + /// Ported from llama.cpp V2.1 fused decode kernel concept. Each SIMD group processes + /// NR0=2 queries against one KV block simultaneously. The key win: packed index unpacking + /// + codebook lookup for K and V is done ONCE and reused across both queries. + /// + /// Register budget per thread (NR0=2, DIMS_PER_LANE=4 for dim=128): + /// - 2 × DIMS_PER_LANE q_vals = 8 floats (query data) + /// - 2 × 1 m/l = 4 floats (online softmax state) + /// - 2 × DIMS_PER_LANE o = 8 floats (value accumulators) + /// - codebook regs shared = KEY_LEVELS + VAL_LEVELS floats + /// Total: ~24 extra floats vs NR0=1. Well within Apple GPU register file. + /// + /// Zero threadgroup memory: all score computation + softmax + V accumulation happen + /// in SIMD registers. No shared memory needed in pass 1 (same as NR0=1 baseline). + /// Note: pass 2 still needs threadgroup memory for dim>32 (cross-SIMD gather for rotation). + /// See turboFlashPass2FusedRotSource comments for details. + /// + /// Grid: (32, totalQueries/NR0, numBlocks) + /// Threadgroup: (32, 1, 1) + /// + /// Template params: KeyBits, ValueBits, Dim, KeyPackedWidth, ValuePackedWidth, + /// BlockSize, token_count, repeat_count, num_blocks, NR0 + static let turboFlashPass1NR0Source = """ + // Per-call values arrive via the `params` input buffer, NOT as template + // args: template values are baked into the kernel name, so a varying + // token_count would JIT-compile and permanently cache one Metal + // library per generated token (unbounded host-memory growth + a + // shader compile every step). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint num_blocks = params[2]; + constexpr uint KEY_MASK = (1u << KeyBits) - 1u; + constexpr uint KEY_LEVELS = 1u << KeyBits; + constexpr uint VAL_MASK = (1u << ValueBits) - 1u; + constexpr uint VAL_LEVELS = 1u << ValueBits; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; // SIMD lane (0-31) + uint query_group = thread_position_in_grid.y; // which group of NR0 queries + uint block_idx = thread_position_in_grid.z; // which KV block + + // Token range for this block + uint t_start = block_idx * BlockSize; + uint t_end = t_start + BlockSize; + if (t_end > (uint)token_count) t_end = (uint)token_count; + + // Load key codebook into registers (shared across all NR0 queries) + float key_cb[KEY_LEVELS]; + for (uint i = 0; i < KEY_LEVELS; i++) { + key_cb[i] = key_codebook[i]; + } + + // Load value codebook into registers (shared across all NR0 queries) + float val_cb[VAL_LEVELS]; + for (uint i = 0; i < VAL_LEVELS; i++) { + val_cb[i] = val_codebook[i]; + } + + // Load query values for ALL NR0 rows, each row's dims interleaved in registers + float q_vals[NR0 * DIMS_PER_LANE]; + for (uint r = 0; r < NR0; r++) { + uint q_idx = query_group * NR0 + r; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + q_vals[r * DIMS_PER_LANE + i] = (d < Dim) ? q_rot[q_idx * Dim + d] : 0.0f; + } + } + + // Per-query KV head mapping (for GQA, each query may map to different KV head) + uint kv_indices[NR0]; + for (uint r = 0; r < NR0; r++) { + kv_indices[r] = (query_group * NR0 + r) / repeat_count; + } + + // Online softmax state, NR0 independent streams, all in registers + float m_state[NR0]; + float l_state[NR0]; + float o_state[NR0 * DIMS_PER_LANE]; + for (uint r = 0; r < NR0; r++) { + m_state[r] = -INFINITY; + l_state[r] = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + o_state[r * DIMS_PER_LANE + i] = 0.0f; + } + } + + // Process tokens in this block, KV dequant done ONCE, reused across NR0 queries + for (uint t = t_start; t < t_end; t++) { + // --- Dequant K for this token ONCE (amortized across NR0 queries) --- + // Each lane unpacks its dims' codebook values into registers + float k_decoded[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) { k_decoded[i] = 0.0f; continue; } + + uint k_bit_offset = d * KeyBits; + uint k_word_idx = k_bit_offset / 32; + uint k_shift = k_bit_offset % 32; + + // All NR0 rows share one KV head by construction: the Swift + // dispatcher only selects this kernel when repeat_count is a + // multiple of NR0, so an aligned group can never span a + // KV-head boundary and kv_indices[0] is exact. + const device uint32_t* k_packed_ptr = key_packed + kv_indices[0] * token_count * KeyPackedWidth + t * KeyPackedWidth; + + uint k_value = (k_packed_ptr[k_word_idx] >> k_shift); + int k_spill = (int)k_shift + (int)KeyBits - 32; + if (k_spill > 0) { + k_value |= (k_packed_ptr[k_word_idx + 1] << ((uint)KeyBits - (uint)k_spill)); + } + k_value &= KEY_MASK; + k_decoded[i] = key_cb[k_value]; + } + float k_norm = key_norms[kv_indices[0] * token_count + t]; + + // --- Dequant V for this token ONCE --- + float v_decoded[DIMS_PER_LANE]; + const device uint32_t* v_packed_ptr = val_packed + kv_indices[0] * token_count * ValuePackedWidth + t * ValuePackedWidth; + float v_norm = val_norms[kv_indices[0] * token_count + t]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) { v_decoded[i] = 0.0f; continue; } + + uint v_bit_offset = d * ValueBits; + uint v_word_idx = v_bit_offset / 32; + uint v_shift = v_bit_offset % 32; + uint v_value = (v_packed_ptr[v_word_idx] >> v_shift); + int v_spill = (int)v_shift + (int)ValueBits - 32; + if (v_spill > 0) { + v_value |= (v_packed_ptr[v_word_idx + 1] << ((uint)ValueBits - (uint)v_spill)); + } + v_value &= VAL_MASK; + v_decoded[i] = val_cb[v_value] * v_norm; + } + + // --- Score + softmax + V accumulate for each of NR0 queries --- + // K/V dequant above is the expensive part, this loop is cheap ALU + for (uint r = 0; r < NR0; r++) { + // Dot product: q[r] · k (both already in registers) + float dot_partial = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + dot_partial += q_vals[r * DIMS_PER_LANE + i] * k_decoded[i]; + } + float score = simd_sum(dot_partial) * k_norm; + + // Online softmax update + float new_m = max(m_state[r], score); + float exp_diff = exp(m_state[r] - new_m); + float exp_score = exp(score - new_m); + + // V accumulation (reusing pre-decoded values) + for (uint i = 0; i < DIMS_PER_LANE; i++) { + o_state[r * DIMS_PER_LANE + i] = o_state[r * DIMS_PER_LANE + i] * exp_diff + exp_score * v_decoded[i]; + } + + l_state[r] = l_state[r] * exp_diff + exp_score; + m_state[r] = new_m; + } + } + + // Write partial results for all NR0 queries + for (uint r = 0; r < NR0; r++) { + uint q_idx = query_group * NR0 + r; + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o_partials[partial_base + d] = o_state[r * DIMS_PER_LANE + i]; + } + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = m_state[r]; + l_partials[ml_idx] = l_state[r]; + } + } + """ + + /// TurboFlashAttention Pass 1 NR0 Causal: Multi-row amortized KV dequant with causal masking. + /// + /// Same as turboFlashPass1NR0Source but each query within the NR0 group has its own + /// causal boundary. For L>1 prefill, q_within_L differs per row so each row may attend + /// to a different number of tokens. We compute the conservative (minimum) causal boundary + /// across the NR0 group for the shared K/V dequant, then mask per-row in the score loop. + /// + /// Grid: (32, totalQueries/NR0, numBlocks) + /// Threadgroup: (32, 1, 1) + /// + /// Template params: KeyBits, ValueBits, Dim, KeyPackedWidth, ValuePackedWidth, + /// BlockSize, token_count, repeat_count, num_blocks, NR0, L, q_offset + static let turboFlashPass1NR0CausalSource = """ + // Per-call values arrive via the `params` input buffer, NOT as template + // args: template values are baked into the kernel name, so a varying + // token_count would JIT-compile and permanently cache one Metal + // library per generated token (unbounded host-memory growth + a + // shader compile every step). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint num_blocks = params[2]; + const uint L = params[3]; + const uint q_offset = params[4]; + constexpr uint KEY_MASK = (1u << KeyBits) - 1u; + constexpr uint KEY_LEVELS = 1u << KeyBits; + constexpr uint VAL_MASK = (1u << ValueBits) - 1u; + constexpr uint VAL_LEVELS = 1u << ValueBits; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; + uint query_group = thread_position_in_grid.y; + uint block_idx = thread_position_in_grid.z; + + // Token range for this block + uint t_start = block_idx * BlockSize; + uint t_end = t_start + BlockSize; + if (t_end > (uint)token_count) t_end = (uint)token_count; + + // Compute per-row causal boundaries and find the maximum (most permissive) + // for the shared token loop. Per-row masking happens inside the score loop. + uint q_abs[NR0]; + uint max_q_abs = 0; + for (uint r = 0; r < NR0; r++) { + uint q_idx = query_group * NR0 + r; + uint q_within_L = q_idx % L; + q_abs[r] = q_offset + q_within_L; + if (q_abs[r] > max_q_abs) max_q_abs = q_abs[r]; + } + + // Early exit: entire block is future-masked for ALL NR0 queries + if (t_start > max_q_abs) { + for (uint r = 0; r < NR0; r++) { + uint q_idx = query_group * NR0 + r; + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) o_partials[partial_base + d] = 0.0f; + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = -INFINITY; + l_partials[ml_idx] = 0.0f; + } + } + return; + } + + // Clamp t_end to the most permissive causal boundary + if (t_end > max_q_abs + 1) t_end = max_q_abs + 1; + + // Load codebooks (shared across all NR0 queries) + float key_cb[KEY_LEVELS]; + for (uint i = 0; i < KEY_LEVELS; i++) key_cb[i] = key_codebook[i]; + float val_cb[VAL_LEVELS]; + for (uint i = 0; i < VAL_LEVELS; i++) val_cb[i] = val_codebook[i]; + + // Load query values for all NR0 rows + float q_vals[NR0 * DIMS_PER_LANE]; + for (uint r = 0; r < NR0; r++) { + uint q_idx = query_group * NR0 + r; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + q_vals[r * DIMS_PER_LANE + i] = (d < Dim) ? q_rot[q_idx * Dim + d] : 0.0f; + } + } + + // KV head mapping (use first query's head, same assumption as non-causal NR0) + uint q_head_idx_0 = (query_group * NR0) / L; + uint kv_idx = q_head_idx_0 / repeat_count; + + // Online softmax state, NR0 independent streams + float m_state[NR0]; + float l_state[NR0]; + float o_state[NR0 * DIMS_PER_LANE]; + for (uint r = 0; r < NR0; r++) { + m_state[r] = -INFINITY; + l_state[r] = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) o_state[r * DIMS_PER_LANE + i] = 0.0f; + } + + // Process tokens, KV dequant once, score per-row with causal mask + for (uint t = t_start; t < t_end; t++) { + // Dequant K once + float k_decoded[DIMS_PER_LANE]; + const device uint32_t* k_packed_ptr = key_packed + kv_idx * token_count * KeyPackedWidth + t * KeyPackedWidth; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) { k_decoded[i] = 0.0f; continue; } + uint k_bit_offset = d * KeyBits; + uint k_word_idx = k_bit_offset / 32; + uint k_shift = k_bit_offset % 32; + uint k_value = (k_packed_ptr[k_word_idx] >> k_shift); + int k_spill = (int)k_shift + (int)KeyBits - 32; + if (k_spill > 0) { + k_value |= (k_packed_ptr[k_word_idx + 1] << ((uint)KeyBits - (uint)k_spill)); + } + k_value &= KEY_MASK; + k_decoded[i] = key_cb[k_value]; + } + float k_norm = key_norms[kv_idx * token_count + t]; + + // Dequant V once + float v_decoded[DIMS_PER_LANE]; + const device uint32_t* v_packed_ptr = val_packed + kv_idx * token_count * ValuePackedWidth + t * ValuePackedWidth; + float v_norm = val_norms[kv_idx * token_count + t]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) { v_decoded[i] = 0.0f; continue; } + uint v_bit_offset = d * ValueBits; + uint v_word_idx = v_bit_offset / 32; + uint v_shift = v_bit_offset % 32; + uint v_value = (v_packed_ptr[v_word_idx] >> v_shift); + int v_spill = (int)v_shift + (int)ValueBits - 32; + if (v_spill > 0) { + v_value |= (v_packed_ptr[v_word_idx + 1] << ((uint)ValueBits - (uint)v_spill)); + } + v_value &= VAL_MASK; + v_decoded[i] = val_cb[v_value] * v_norm; + } + + // Score + softmax + V for each query row (with per-row causal mask) + for (uint r = 0; r < NR0; r++) { + // Per-row causal: skip if this token is future for this specific query + if (t > q_abs[r]) continue; + + float dot_partial = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + dot_partial += q_vals[r * DIMS_PER_LANE + i] * k_decoded[i]; + } + float score = simd_sum(dot_partial) * k_norm; + + float new_m = max(m_state[r], score); + float exp_diff = exp(m_state[r] - new_m); + float exp_score = exp(score - new_m); + + for (uint i = 0; i < DIMS_PER_LANE; i++) { + o_state[r * DIMS_PER_LANE + i] = o_state[r * DIMS_PER_LANE + i] * exp_diff + exp_score * v_decoded[i]; + } + l_state[r] = l_state[r] * exp_diff + exp_score; + m_state[r] = new_m; + } + } + + // Write partial results for all NR0 queries + for (uint r = 0; r < NR0; r++) { + uint q_idx = query_group * NR0 + r; + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) o_partials[partial_base + d] = o_state[r * DIMS_PER_LANE + i]; + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = m_state[r]; + l_partials[ml_idx] = l_state[r]; + } + } + """ + + /// TurboFlashAttention Pass 2: Cross-block reduction. + /// + /// Merges partial online softmax states from pass 1 across token blocks. + /// Each SIMD group handles one query, iterating over all blocks to produce + /// the final normalized output. + /// + /// Grid: (32, totalQueries, 1) + /// Threadgroup: (32, 1, 1) + /// + /// Template params: Dim, num_blocks + static let turboFlashPass2Source = """ + // num_blocks via params buffer (varies per call; see pass-1 note). + const uint num_blocks = params[0]; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; + uint q_idx = thread_position_in_grid.y; + + float m = -INFINITY; + float l = 0.0f; + float o[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) o[i] = 0.0f; + + for (uint b = 0; b < (uint)num_blocks; b++) { + uint ml_idx = q_idx * num_blocks + b; + + // All lanes read the same m/l (broadcast read from device memory) + float block_m = m_partials[ml_idx]; + float block_l = l_partials[ml_idx]; + + // Skip empty blocks + if (block_l == 0.0f) continue; + + float new_m = max(m, block_m); + float exp_old = exp(m - new_m); + float exp_block = exp(block_m - new_m); + + uint partial_base = (q_idx * num_blocks + b) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o[i] = o[i] * exp_old + o_partials[partial_base + d] * exp_block; + } + } + + l = l * exp_old + block_l * exp_block; + m = new_m; + } + + // Write normalized output + float inv_l = (l > 0.0f) ? (1.0f / l) : 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + output[q_idx * Dim + d] = o[i] * inv_l; + } + } + """ + + /// TurboFlashAttention Pass 2 with fused output rotation. + /// + /// Same as turboFlashPass2Source but applies inverse value rotation (Π_val) in-kernel + /// after merging partials, eliminating a separate MLX matmul dispatch. + /// Uses threadgroup shared memory to gather the full output vector across SIMD lanes, + /// then each lane computes rotated output as dot product with rotation matrix rows. + /// + /// Grid: (32, totalQueries, 1) + /// Threadgroup: (32, 1, 1) + /// + /// Template params: Dim, num_blocks + static let turboFlashPass2FusedRotSource = """ + // num_blocks via params buffer (varies per call; see pass-1 note). + const uint num_blocks = params[0]; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; + uint q_idx = thread_position_in_grid.y; + + float m = -INFINITY; + float l = 0.0f; + float o[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) o[i] = 0.0f; + + for (uint b = 0; b < (uint)num_blocks; b++) { + uint ml_idx = q_idx * num_blocks + b; + + float block_m = m_partials[ml_idx]; + float block_l = l_partials[ml_idx]; + + if (block_l == 0.0f) continue; + + float new_m = max(m, block_m); + float exp_old = exp(m - new_m); + float exp_block = exp(block_m - new_m); + + uint partial_base = (q_idx * num_blocks + b) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o[i] = o[i] * exp_old + o_partials[partial_base + d] * exp_block; + } + } + + l = l * exp_old + block_l * exp_block; + m = new_m; + } + + // Normalize + float inv_l = (l > 0.0f) ? (1.0f / l) : 0.0f; + + // Gather normalized output into threadgroup shared memory for rotation + threadgroup float shared_out[Dim]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + shared_out[d] = o[i] * inv_l; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Apply inverse value rotation: output[d] = Σ_j shared_out[j] * Π_val[j][d] + // matmul(x, Π_val) reads column d of Π_val for output dimension d. + // Π_val is stored row-major [Dim, Dim], so column d = val_rotation[j * Dim + d] + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + float acc = 0.0f; + for (uint j = 0; j < Dim; j++) { + acc += shared_out[j] * val_rotation[j * Dim + d]; + } + output[q_idx * Dim + d] = acc; + } + } + """ + + /// Value aggregation kernel: weighted sum of codebook-quantized values. + /// + /// output[d] = Σ_t weights[t] * norm[t] * codebook[val_idx[t,d]] + /// Result is in rotated space, caller applies inverse rotation. + /// + /// Grid: (32, totalHeads, ceil(Dim/32)) + /// Threadgroup: (32, 1, 1) + static let valueKernelSource = """ + // token_count/repeat_count/L via params buffer (vary per call; template + // args would bake them into the kernel name = one compiled Metal + // library cached per token). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint chunk_len = params[2]; + constexpr uint MASK = (1u << Bits) - 1u; + constexpr uint LEVELS = 1u << Bits; + + uint lane = thread_position_in_grid.x; + uint head_idx = thread_position_in_grid.y; + uint dim_block = thread_position_in_grid.z; + + uint d = dim_block * 32 + lane; + if (d >= Dim) return; + + // Rows are flattened [B, heads, L]; divide out the chunk length + // before the GQA repeat mapping. + uint kv_head = (head_idx / chunk_len) / repeat_count; + + // Load codebook + float cb[LEVELS]; + for (uint i = 0; i < LEVELS; i++) { + cb[i] = codebook[i]; + } + + float acc = 0.0f; + for (uint t = 0; t < (uint)token_count; t++) { + float w = weights[head_idx * token_count + t]; + if (w < 1e-6f) continue; // Sparse V: skip negligible attention weights + + float norm_val = norms[kv_head * token_count + t]; + const device uint32_t* packed_ptr = packed + kv_head * token_count * PackedWidth + t * PackedWidth; + + uint bit_offset = d * Bits; + uint word_idx = bit_offset / 32; + uint shift = bit_offset % 32; + uint value = (packed_ptr[word_idx] >> shift); + + int spill = (int)shift + (int)Bits - 32; + if (spill > 0) { + value |= (packed_ptr[word_idx + 1] << ((uint)Bits - (uint)spill)); + } + value &= MASK; + + acc += w * norm_val * cb[value]; + } + + output[head_idx * Dim + d] = acc; + """ +} + +// MARK: - Kernel Dispatch Wrappers + +enum TurboQuantKernelOps { + + // Kernel caches + nonisolated(unsafe) private static var encodeKernels: [String: MLXFast.MLXFastKernel] = [:] + nonisolated(unsafe) private static var scoreKernels: [String: MLXFast.MLXFastKernel] = [:] + nonisolated(unsafe) private static var valueKernels: [String: MLXFast.MLXFastKernel] = [:] + private static let lock = NSLock() + + /// Fused encode: norm + rotate + quantize + pack + norm correction in single GPU dispatch. + /// + /// - Parameters: + /// - input: Raw vectors [numRows, D] float32 + /// - rotation: Rotation matrix Π [D, D] float32 + /// - boundaries: Codebook boundaries [2^bits - 1] float32 + /// - codebook: Centroids [2^bits] float32 (needed for norm correction) + /// - bits: Quantization bit-width + /// - dim: Vector dimension + /// - Returns: (packed: [numRows, PackedWidth] uint32, norms: [numRows] float32) + /// norms are norm-corrected: original_norm / reconstruction_norm + /// Kernel sources declare float buffers; normalize any half/bfloat input + /// so a f16/bf16 activation stream is not silently reinterpreted. + @inline(__always) + private static func f32(_ x: MLXArray) -> MLXArray { + x.dtype == .float32 ? x : x.asType(.float32) + } + + static func fusedEncode( + input: MLXArray, + rotation: MLXArray, + boundaries: MLXArray, + codebook: MLXArray, + bits: Int, + dim: Int + ) -> (packed: MLXArray, norms: MLXArray) { + let pw = TurboQuantPacking.packedWidth(count: dim, bits: bits) + let key = "encode_nc_\(bits)_\(dim)" // nc = norm-corrected + + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = encodeKernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_fused_encode_\(bits)_\(dim)", + inputNames: ["input", "rotation", "boundaries", "codebook"], + outputNames: ["packed_out", "norms_out"], + source: TurboQuantMetalKernels.fusedEncodeSource, + ensureRowContiguous: true + ) + lock.lock() + encodeKernels[key] = k + lock.unlock() + kernel = k + } + + let numRows = input.dim(0) + + let results = kernel( + [f32(input), f32(rotation), f32(boundaries), f32(codebook)], + template: [ + ("Bits", bits), ("Dim", dim), ("PackedWidth", pw), + ], + grid: (dim, numRows, 1), + threadGroup: (dim, 1, 1), + outputShapes: [[numRows, pw], [numRows]], + outputDTypes: [.uint32, .float32] + ) + + return (packed: results[0], norms: results[1]) + } + + /// Fused WHT encode: norm + WHT rotation + quantize + pack (raw norms, no correction). + /// + /// Same as fusedEncode but uses O(d log d) Walsh-Hadamard butterfly instead of + /// O(d²) dense matmul. Only works for power-of-2 dimensions. + /// + /// WHT is orthogonal so norms are preserved, no norm correction needed. + /// Codebook is NOT passed to the kernel (saves one buffer bind + GPU transfer). + /// + /// - Parameters: + /// - input: Raw vectors [numRows, D] float32 + /// - whtSigns: Random ±1 signs [D] float32 + /// - boundaries: Codebook boundaries [2^bits - 1] float32 + /// - codebook: Centroids [2^bits] float32 (unused by kernel, kept in API for caller convenience) + /// - bits: Quantization bit-width + /// - dim: Vector dimension (must be power of 2) + /// - Returns: (packed: [numRows, PackedWidth] uint32, norms: [numRows] float32, raw norms) + static func fusedEncodeWHT( + input: MLXArray, + whtSigns: MLXArray, + boundaries: MLXArray, + codebook: MLXArray, + bits: Int, + dim: Int + ) -> (packed: MLXArray, norms: MLXArray) { + let pw = TurboQuantPacking.packedWidth(count: dim, bits: bits) + let logDim = Int(log2(Double(dim))) + let key = "encode_wht_\(bits)_\(dim)" + + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = encodeKernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_fused_encode_wht_\(bits)_\(dim)", + inputNames: ["input", "wht_signs", "boundaries"], + outputNames: ["packed_out", "norms_out"], + source: TurboQuantMetalKernels.fusedEncodeWHTSource, + ensureRowContiguous: true + ) + lock.lock() + encodeKernels[key] = k + lock.unlock() + kernel = k + } + + let numRows = input.dim(0) + + // NOTE: codebook no longer passed, WHT kernel stores raw norms (no norm correction) + let results = kernel( + [f32(input), f32(whtSigns), f32(boundaries)], + template: [ + ("Bits", bits), ("Dim", dim), ("PackedWidth", pw), ("LogDim", logDim), + ], + grid: (dim, numRows, 1), + threadGroup: (dim, 1, 1), + outputShapes: [[numRows, pw], [numRows]], + outputDTypes: [.uint32, .float32] + ) + + return (packed: results[0], norms: results[1]) + } + + // Flash attention kernel caches + nonisolated(unsafe) private static var flashPass1Kernels: [String: MLXFast.MLXFastKernel] = [:] + nonisolated(unsafe) private static var flashPass1NR0Kernels: [String: MLXFast.MLXFastKernel] = + [:] + nonisolated(unsafe) private static var flashPass2Kernels: [String: MLXFast.MLXFastKernel] = [:] + + /// NR0: number of query rows processed per SIMD group in the multi-row amortized kernel. + /// + /// Ported from llama.cpp V2.1: each threadgroup loads K/V packed data once and reuses + /// it across NR0 queries. At NR0=2, the KV dequant cost is halved per query. + /// + /// NR0=2 is conservative, register pressure is ~24 extra floats per thread (for dim=128). + /// Apple M-series GPUs have 96 registers per thread (384 bytes), so this fits comfortably. + /// + /// Override via environment variable `TURBO_FLASH_NR0` (must be power of 2). + static let flashNR0: Int = { + if let envValue = ProcessInfo.processInfo.environment["TURBO_FLASH_NR0"], + let parsed = Int(envValue), parsed > 0, (parsed & (parsed - 1)) == 0 + { + return parsed + } + return 2 // default, conservative starting point + }() + + /// Default block size for TurboFlashAttention two-pass approach. + /// Each SIMD group processes this many tokens per block. + /// Tuned for M1 Max via sweep: B=64 wins or ties at all token counts (512-8192+). + /// Smaller blocks = more parallelism but more pass-2 merge work. + static let flashBlockSize = 64 + + /// Shared pass 1 dispatch, used by both causal and non-causal variants. + private static func dispatchFlashPass1( + source: String, cachePrefix: String, + rotatedQueries: MLXArray, + keyPacked: MLXArray, keyNorms: MLXArray, keyCodebook: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, + keyBits: Int, valueBits: Int, dim: Int, + blockSize: Int, queryChunkLength: Int = 1, queryOffset: Int = 0 + ) -> (oPartials: MLXArray, mPartials: MLXArray, lPartials: MLXArray) { + let kpw = TurboQuantPacking.packedWidth(count: dim, bits: keyBits) + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + + let pass1Key = "\(cachePrefix)_\(keyBits)_\(valueBits)_\(dim)" + let pass1Kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = flashPass1Kernels[pass1Key] { + pass1Kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_\(cachePrefix)_\(keyBits)_\(valueBits)_\(dim)", + inputNames: [ + "q_rot", "key_packed", "key_norms", "key_codebook", + "val_packed", "val_norms", "val_codebook", "params", + ], + outputNames: ["o_partials", "m_partials", "l_partials"], + source: source, + ensureRowContiguous: true + ) + lock.lock() + flashPass1Kernels[pass1Key] = k + lock.unlock() + pass1Kernel = k + } + + let template: [(String, Int)] = [ + ("KeyBits", keyBits), ("ValueBits", valueBits), + ("Dim", dim), ("KeyPackedWidth", kpw), ("ValuePackedWidth", vpw), + ("BlockSize", blockSize), + ] + // Varying values travel in a buffer, not templates, a template value + // is baked into the kernel name and would JIT + cache one Metal + // library per distinct token_count (i.e., per generated token). + let params = MLXArray( + [ + UInt32(tokenCount), UInt32(repeatCount), UInt32(numBlocks), + UInt32(queryChunkLength), UInt32(queryOffset), + ]) + + let partials = pass1Kernel( + [ + f32(rotatedQueries), keyPacked, f32(keyNorms), f32(keyCodebook), + valPacked, f32(valNorms), f32(valCodebook), params, + ], + template: template, + grid: (32, totalQ, numBlocks), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ * numBlocks, dim], [totalQ, numBlocks], [totalQ, numBlocks]], + outputDTypes: [.float32, .float32, .float32] + ) + + return (oPartials: partials[0], mPartials: partials[1], lPartials: partials[2]) + } + + /// NR0 multi-row pass 1 dispatch, processes NR0 queries per SIMD group. + /// + /// Each SIMD group loads K/V packed data once and computes scores for NR0 queries. + /// The grid Y dimension is totalQueries/NR0 instead of totalQueries. + /// Output shapes are the same as NR0=1 (partials indexed by original q_idx). + /// + /// Precondition: totalQueries must be divisible by NR0. + private static func dispatchFlashPass1NR0( + rotatedQueries: MLXArray, + keyPacked: MLXArray, keyNorms: MLXArray, keyCodebook: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, + keyBits: Int, valueBits: Int, dim: Int, + blockSize: Int, nr0: Int + ) -> (oPartials: MLXArray, mPartials: MLXArray, lPartials: MLXArray) { + let kpw = TurboQuantPacking.packedWidth(count: dim, bits: keyBits) + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + let queryGroups = totalQ / nr0 + + let pass1Key = "flash_p1_nr0_\(keyBits)_\(valueBits)_\(dim)_\(nr0)" + let pass1Kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = flashPass1NR0Kernels[pass1Key] { + pass1Kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_flash_p1_nr0_\(keyBits)_\(valueBits)_\(dim)_\(nr0)", + inputNames: [ + "q_rot", "key_packed", "key_norms", "key_codebook", + "val_packed", "val_norms", "val_codebook", "params", + ], + outputNames: ["o_partials", "m_partials", "l_partials"], + source: TurboQuantMetalKernels.turboFlashPass1NR0Source, + ensureRowContiguous: true + ) + lock.lock() + flashPass1NR0Kernels[pass1Key] = k + lock.unlock() + pass1Kernel = k + } + + let template: [(String, Int)] = [ + ("KeyBits", keyBits), ("ValueBits", valueBits), + ("Dim", dim), ("KeyPackedWidth", kpw), ("ValuePackedWidth", vpw), + ("BlockSize", blockSize), ("NR0", nr0), + ] + let params = MLXArray( + [UInt32(tokenCount), UInt32(repeatCount), UInt32(numBlocks), 1, 0] as [UInt32]) + + // Grid Y = queryGroups (totalQ / NR0), not totalQ + let partials = pass1Kernel( + [ + f32(rotatedQueries), keyPacked, f32(keyNorms), f32(keyCodebook), + valPacked, f32(valNorms), f32(valCodebook), params, + ], + template: template, + grid: (32, queryGroups, numBlocks), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ * numBlocks, dim], [totalQ, numBlocks], [totalQ, numBlocks]], + outputDTypes: [.float32, .float32, .float32] + ) + + return (oPartials: partials[0], mPartials: partials[1], lPartials: partials[2]) + } + + /// NR0 multi-row causal pass 1 dispatch, processes NR0 queries with per-row causal masking. + /// + /// Same as dispatchFlashPass1NR0 but supports causal masking for L>1 prefill. + /// Each query in the NR0 group has its own causal boundary. + /// + /// Precondition: totalQueries must be divisible by NR0. + private static func dispatchFlashPass1NR0Causal( + rotatedQueries: MLXArray, + keyPacked: MLXArray, keyNorms: MLXArray, keyCodebook: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, + keyBits: Int, valueBits: Int, dim: Int, + blockSize: Int, nr0: Int, + queryChunkLength: Int, queryOffset: Int + ) -> (oPartials: MLXArray, mPartials: MLXArray, lPartials: MLXArray) { + let kpw = TurboQuantPacking.packedWidth(count: dim, bits: keyBits) + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + let queryGroups = totalQ / nr0 + + let pass1Key = "flash_p1_nr0_causal_\(keyBits)_\(valueBits)_\(dim)_\(nr0)" + let pass1Kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = flashPass1NR0Kernels[pass1Key] { + pass1Kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_flash_p1_nr0_causal_\(keyBits)_\(valueBits)_\(dim)_\(nr0)", + inputNames: [ + "q_rot", "key_packed", "key_norms", "key_codebook", + "val_packed", "val_norms", "val_codebook", "params", + ], + outputNames: ["o_partials", "m_partials", "l_partials"], + source: TurboQuantMetalKernels.turboFlashPass1NR0CausalSource, + ensureRowContiguous: true + ) + lock.lock() + flashPass1NR0Kernels[pass1Key] = k + lock.unlock() + pass1Kernel = k + } + + let template: [(String, Int)] = [ + ("KeyBits", keyBits), ("ValueBits", valueBits), + ("Dim", dim), ("KeyPackedWidth", kpw), ("ValuePackedWidth", vpw), + ("BlockSize", blockSize), ("NR0", nr0), + ] + let params = MLXArray( + [ + UInt32(tokenCount), UInt32(repeatCount), UInt32(numBlocks), + UInt32(queryChunkLength), UInt32(queryOffset), + ]) + + let partials = pass1Kernel( + [ + f32(rotatedQueries), keyPacked, f32(keyNorms), f32(keyCodebook), + valPacked, f32(valNorms), f32(valCodebook), params, + ], + template: template, + grid: (32, queryGroups, numBlocks), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ * numBlocks, dim], [totalQ, numBlocks], [totalQ, numBlocks]], + outputDTypes: [.float32, .float32, .float32] + ) + + return (oPartials: partials[0], mPartials: partials[1], lPartials: partials[2]) + } + + /// Shared pass 2 dispatch, with optional fused output rotation. + /// + /// When `valRotation` is provided, the inverse value rotation (Π_val) is applied + /// in-kernel using threadgroup shared memory, eliminating a separate MLX matmul dispatch. + /// Output is in original (non-rotated) space. + /// + /// When `valRotation` is nil, output is in rotated V space (caller must apply inverse rotation). + private static func dispatchFlashPass2( + oPartials: MLXArray, mPartials: MLXArray, lPartials: MLXArray, + dim: Int, numBlocks: Int, totalQ: Int, + valRotation: MLXArray? = nil + ) -> MLXArray { + let fused = valRotation != nil + let pass2Key = fused ? "flash_p2_fused_\(dim)" : "flash_p2_\(dim)" + let pass2Kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = flashPass2Kernels[pass2Key] { + pass2Kernel = cached + lock.unlock() + } else { + lock.unlock() + let k: MLXFast.MLXFastKernel + if fused { + k = MLXFast.metalKernel( + name: "turbo_flash_p2_fused_\(dim)", + inputNames: [ + "o_partials", "m_partials", "l_partials", "val_rotation", "params", + ], + outputNames: ["output"], + source: TurboQuantMetalKernels.turboFlashPass2FusedRotSource, + ensureRowContiguous: true + ) + } else { + k = MLXFast.metalKernel( + name: "turbo_flash_p2_\(dim)", + inputNames: ["o_partials", "m_partials", "l_partials", "params"], + outputNames: ["output"], + source: TurboQuantMetalKernels.turboFlashPass2Source, + ensureRowContiguous: true + ) + } + lock.lock() + flashPass2Kernels[pass2Key] = k + lock.unlock() + pass2Kernel = k + } + + let params = MLXArray([UInt32(numBlocks)]) + let inputs: [MLXArray] = + fused + ? [oPartials, mPartials, lPartials, valRotation!, params] + : [oPartials, mPartials, lPartials, params] + + return pass2Kernel( + inputs, + template: [ + ("Dim", dim) + ], + grid: (32, totalQ, 1), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ, dim]], + outputDTypes: [.float32] + )[0] + } + + nonisolated(unsafe) private static var quantizePackKernels: [String: MLXFast.MLXFastKernel] = + [:] + + /// Quantize + pack + norm-correct pre-rotated calibrated vectors. + /// `rotated` is [rows, dim] from MSECodec.rotatedUnit, `rawNorms` [rows]. + static func fusedQuantizePackScaled( + rotated: MLXArray, rawNorms: MLXArray, scale: MLXArray, + boundaries: MLXArray, codebook: MLXArray, bits: Int, dim: Int + ) -> (packed: MLXArray, norms: MLXArray) { + let pw = TurboQuantPacking.packedWidth(count: dim, bits: bits) + let key = "qpack_scaled_\(bits)_\(dim)" + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = quantizePackKernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_qpack_scaled_\(bits)_\(dim)", + inputNames: ["rotated", "norms_in", "scale", "boundaries", "codebook"], + outputNames: ["packed_out", "norms_out"], + source: TurboQuantMetalKernels.fusedQuantizePackScaledSource, + ensureRowContiguous: true + ) + lock.lock() + quantizePackKernels[key] = k + lock.unlock() + kernel = k + } + let numRows = rotated.dim(0) + let results = kernel( + [f32(rotated), f32(rawNorms), f32(scale), f32(boundaries), f32(codebook)], + template: [("Bits", bits), ("Dim", dim), ("PackedWidth", pw)], + grid: (dim, numRows, 1), + threadGroup: (dim, 1, 1), + outputShapes: [[numRows, pw], [numRows]], + outputDTypes: [.uint32, .float32] + ) + return (packed: results[0], norms: results[1]) + } + + /// TurboFlashAttention: two-pass fused Score + Online Softmax + Value. + /// + /// Pass 1: Parallelizes across (query × token_block) pairs. Each SIMD group processes + /// BlockSize tokens, producing partial online softmax state (m, l, o[D]). + /// Pass 2: Merges partial states across blocks to produce final normalized output. + /// + /// Eliminates intermediate score and attention weight arrays entirely. + /// + /// - Parameter valRotation: Optional [D, D] inverse value rotation matrix. When provided, + /// rotation is fused into pass 2, eliminating a separate MLX matmul dispatch. + /// Output is in original space. When nil, output is in rotated V space. + /// - Parameter blockSize: Tokens per block (default: flashBlockSize). Smaller = more parallelism + /// but more pass-2 merge work. Must be > 0. + /// - Returns: Output [totalQ, D] float32 + /// Pass-1 dispatch, raw-fp16 K + turbo-V. Reuses the packed-K flash + /// structure and pass 2. + private static func dispatchFlashPass1RawK( + rotatedQueries: MLXArray, rawKeys: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, valueBits: Int, dim: Int, blockSize: Int + ) -> (oPartials: MLXArray, mPartials: MLXArray, lPartials: MLXArray) { + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + let key = "flash_p1_rawk_\(valueBits)_\(dim)_\(rawKeys.dtype)" + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = flashPass1Kernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_flash_p1_rawk_\(valueBits)_\(dim)_\(rawKeys.dtype)", + inputNames: [ + "q_rot", "k_raw", "val_packed", "val_norms", "val_codebook", "params", + ], + outputNames: ["o_partials", "m_partials", "l_partials"], + source: TurboQuantMetalKernels.turboFlashPass1RawKSource, + ensureRowContiguous: true + ) + lock.lock() + flashPass1Kernels[key] = k + lock.unlock() + kernel = k + } + let params = MLXArray( + [UInt32(tokenCount), UInt32(repeatCount), UInt32(numBlocks), 1, 0] as [UInt32]) + // rawKeys stays in its native dtype (KT template): a per-step cast + // would copy the whole growing cache every token and push bfloat16 + // values above the float16 range to infinity. + let partials = kernel( + [ + f32(rotatedQueries), rawKeys, + valPacked, f32(valNorms), f32(valCodebook), params, + ], + template: [ + ("ValueBits", valueBits), ("Dim", dim), ("ValuePackedWidth", vpw), + ("BlockSize", blockSize), ("KT", rawKeys.dtype), + ], + grid: (32, totalQ, numBlocks), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ * numBlocks, dim], [totalQ, numBlocks], [totalQ, numBlocks]], + outputDTypes: [.float32, .float32, .float32] + ) + return (partials[0], partials[1], partials[2]) + } + + /// Pass-1 dispatch, 8-bit affine K + turbo-V. + private static func dispatchFlashPass1AffineK( + rotatedQueries: MLXArray, + kWeights: MLXArray, kScales: MLXArray, kBiases: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, valueBits: Int, dim: Int, kGroup: Int, + blockSize: Int + ) -> (oPartials: MLXArray, mPartials: MLXArray, lPartials: MLXArray) { + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + let key = "flash_p1_affk_\(valueBits)_\(dim)_\(kGroup)_\(kScales.dtype)_\(kBiases.dtype)" + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = flashPass1Kernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: + "turbo_flash_p1_affk_\(valueBits)_\(dim)_\(kGroup)_\(kScales.dtype)_\(kBiases.dtype)", + inputNames: [ + "q_rot", "k_weights", "k_scales", "k_biases", + "val_packed", "val_norms", "val_codebook", "params", + ], + outputNames: ["o_partials", "m_partials", "l_partials"], + source: TurboQuantMetalKernels.turboFlashPass1AffineKSource, + ensureRowContiguous: true + ) + lock.lock() + flashPass1Kernels[key] = k + lock.unlock() + kernel = k + } + let params = MLXArray( + [UInt32(tokenCount), UInt32(repeatCount), UInt32(numBlocks), 1, 0] as [UInt32]) + // kScales/kBiases stay in their native dtype (KScaleT/KBiasT + // templates): they are cache-resident and grow with token_count, so + // a per-step cast would copy them every token. + let partials = kernel( + [ + f32(rotatedQueries), kWeights, kScales, kBiases, + valPacked, f32(valNorms), f32(valCodebook), params, + ], + template: [ + ("ValueBits", valueBits), ("Dim", dim), ("ValuePackedWidth", vpw), + ("BlockSize", blockSize), ("KGroup", kGroup), + ("KScaleT", kScales.dtype), ("KBiasT", kBiases.dtype), + ], + grid: (32, totalQ, numBlocks), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ * numBlocks, dim], [totalQ, numBlocks], [totalQ, numBlocks]], + outputDTypes: [.float32, .float32, .float32] + ) + return (partials[0], partials[1], partials[2]) + } + + /// Flash decode, raw-fp16 K + turbo-V (single decode step, L=1). + static func turboFlashRawK( + rotatedQueries: MLXArray, rawKeys: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, valueBits: Int, dim: Int, + valRotation: MLXArray? = nil, blockSize: Int? = nil + ) -> MLXArray { + let bs = blockSize ?? flashBlockSize + let numBlocks = (tokenCount + bs - 1) / bs + let totalQ = rotatedQueries.dim(0) + let (o, m, l) = dispatchFlashPass1RawK( + rotatedQueries: rotatedQueries, rawKeys: rawKeys, + valPacked: valPacked, valNorms: valNorms, valCodebook: valCodebook, + tokenCount: tokenCount, repeatCount: repeatCount, valueBits: valueBits, + dim: dim, blockSize: bs) + return dispatchFlashPass2( + oPartials: o, mPartials: m, lPartials: l, dim: dim, numBlocks: numBlocks, + totalQ: totalQ, valRotation: valRotation) + } + + /// Flash decode, 8-bit affine K + turbo-V (single decode step, L=1). + static func turboFlashAffineK( + rotatedQueries: MLXArray, + kWeights: MLXArray, kScales: MLXArray, kBiases: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, valueBits: Int, dim: Int, kGroup: Int, + valRotation: MLXArray? = nil, blockSize: Int? = nil + ) -> MLXArray { + let bs = blockSize ?? flashBlockSize + let numBlocks = (tokenCount + bs - 1) / bs + let totalQ = rotatedQueries.dim(0) + let (o, m, l) = dispatchFlashPass1AffineK( + rotatedQueries: rotatedQueries, + kWeights: kWeights, kScales: kScales, kBiases: kBiases, + valPacked: valPacked, valNorms: valNorms, valCodebook: valCodebook, + tokenCount: tokenCount, repeatCount: repeatCount, valueBits: valueBits, + dim: dim, kGroup: kGroup, blockSize: bs) + return dispatchFlashPass2( + oPartials: o, mPartials: m, lPartials: l, dim: dim, numBlocks: numBlocks, + totalQ: totalQ, valRotation: valRotation) + } + + static func turboFlashAttention( + rotatedQueries: MLXArray, + keyPacked: MLXArray, + keyNorms: MLXArray, + keyCodebook: MLXArray, + valPacked: MLXArray, + valNorms: MLXArray, + valCodebook: MLXArray, + tokenCount: Int, + repeatCount: Int, + keyBits: Int, + valueBits: Int, + dim: Int, + valRotation: MLXArray? = nil, + blockSize: Int? = nil + ) -> MLXArray { + let blockSize = blockSize ?? flashBlockSize + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + let nr0 = flashNR0 + + // Use NR0 multi-row kernel when totalQ is evenly divisible by NR0 and NR0 > 1. + // Falls back to NR0=1 (original kernel) for remainder queries or when NR0=1. + // NR0 groups read one kv head (kv_indices[0]) for all rows; only valid + // when the GQA repeat factor is a multiple of nr0 so aligned groups can + // never span a KV-head boundary. MHA and odd repeat factors take the + // per-row kernel. + let useNR0 = + nr0 > 1 && totalQ % nr0 == 0 && totalQ >= nr0 && repeatCount % nr0 == 0 + + let oPartials: MLXArray + let mPartials: MLXArray + let lPartials: MLXArray + + if useNR0 { + (oPartials, mPartials, lPartials) = dispatchFlashPass1NR0( + rotatedQueries: rotatedQueries, + keyPacked: keyPacked, keyNorms: keyNorms, keyCodebook: keyCodebook, + valPacked: valPacked, valNorms: valNorms, valCodebook: valCodebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim, + blockSize: blockSize, nr0: nr0 + ) + } else { + (oPartials, mPartials, lPartials) = dispatchFlashPass1( + source: TurboQuantMetalKernels.turboFlashPass1Source, + cachePrefix: "flash_p1", + rotatedQueries: rotatedQueries, + keyPacked: keyPacked, keyNorms: keyNorms, keyCodebook: keyCodebook, + valPacked: valPacked, valNorms: valNorms, valCodebook: valCodebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim, + blockSize: blockSize + ) + } + + return dispatchFlashPass2( + oPartials: oPartials, mPartials: mPartials, lPartials: lPartials, + dim: dim, numBlocks: numBlocks, totalQ: totalQ, + valRotation: valRotation + ) + } + + /// TurboFlashAttention with causal masking for L>1 prefill chunks. + /// + /// Same as turboFlashAttention but each query position only attends to tokens + /// where t <= queryOffset + q_within_L. Eliminates the need to materialize + /// the full [nQHeads, L, T] score matrix for causal masking. + /// + /// - Parameter queryChunkLength: Number of query positions in the chunk (L) + /// - Parameter queryOffset: Absolute position of the first query in the chunk + /// - Returns: Output [totalQ, D] float32 + static func turboFlashAttentionCausal( + rotatedQueries: MLXArray, + keyPacked: MLXArray, + keyNorms: MLXArray, + keyCodebook: MLXArray, + valPacked: MLXArray, + valNorms: MLXArray, + valCodebook: MLXArray, + tokenCount: Int, + repeatCount: Int, + keyBits: Int, + valueBits: Int, + dim: Int, + queryChunkLength: Int, + queryOffset: Int, + valRotation: MLXArray? = nil, + blockSize: Int? = nil + ) -> MLXArray { + let blockSize = blockSize ?? flashBlockSize + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + let nr0 = flashNR0 + + // NR0 groups read one kv head (kv_indices[0]) for all rows; only valid + // when the GQA repeat factor is a multiple of nr0 so aligned groups can + // never span a KV-head boundary. MHA and odd repeat factors take the + // per-row kernel. + let useNR0 = + nr0 > 1 && totalQ % nr0 == 0 && totalQ >= nr0 && repeatCount % nr0 == 0 + + let oPartials: MLXArray + let mPartials: MLXArray + let lPartials: MLXArray + + if useNR0 { + (oPartials, mPartials, lPartials) = dispatchFlashPass1NR0Causal( + rotatedQueries: rotatedQueries, + keyPacked: keyPacked, keyNorms: keyNorms, keyCodebook: keyCodebook, + valPacked: valPacked, valNorms: valNorms, valCodebook: valCodebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim, + blockSize: blockSize, nr0: nr0, + queryChunkLength: queryChunkLength, queryOffset: queryOffset + ) + } else { + (oPartials, mPartials, lPartials) = dispatchFlashPass1( + source: TurboQuantMetalKernels.turboFlashPass1CausalSource, + cachePrefix: "flash_p1_causal", + rotatedQueries: rotatedQueries, + keyPacked: keyPacked, keyNorms: keyNorms, keyCodebook: keyCodebook, + valPacked: valPacked, valNorms: valNorms, valCodebook: valCodebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim, + blockSize: blockSize, + queryChunkLength: queryChunkLength, queryOffset: queryOffset + ) + } + + return dispatchFlashPass2( + oPartials: oPartials, mPartials: mPartials, lPartials: lPartials, + dim: dim, numBlocks: numBlocks, totalQ: totalQ, + valRotation: valRotation + ) + } + + /// Compute Q×K attention scores from packed codebook indices. + /// + /// - Parameters: + /// - rotatedQueries: Pre-rotated queries [totalQ, D] (already scaled) + /// - packed: Packed key indices [totalKVHeads, T, PackedWidth] uint32 + /// - norms: Key norms [totalKVHeads, T] float32 + /// - codebook: Centroids [2^bits] float32 + /// - tokenCount: Number of cached tokens + /// - repeatCount: GQA repeat factor (nQHeads / nKVHeads) + /// - bits: MSE bit-width + /// - dim: Vector dimension + /// - Returns: Scores [totalQ, T] float32 + static func mseScore( + rotatedQueries: MLXArray, + packed: MLXArray, + norms: MLXArray, + codebook: MLXArray, + tokenCount: Int, + repeatCount: Int, + bits: Int, + dim: Int, queryChunkLength: Int = 1 + ) -> MLXArray { + let pw = TurboQuantPacking.packedWidth(count: dim, bits: bits) + let key = "\(bits)_\(dim)" + + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = scoreKernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_score_\(bits)_\(dim)", + inputNames: ["q_rot", "packed", "norms", "codebook", "params"], + outputNames: ["scores"], + source: TurboQuantMetalKernels.scoreKernelSource, + ensureRowContiguous: true + ) + lock.lock() + scoreKernels[key] = k + lock.unlock() + kernel = k + } + + let totalQ = rotatedQueries.dim(0) + + let params = MLXArray([UInt32(tokenCount), UInt32(repeatCount), UInt32(queryChunkLength)]) + return kernel( + [f32(rotatedQueries), packed, f32(norms), f32(codebook), params], + template: [ + ("Bits", bits), ("Dim", dim), ("PackedWidth", pw), + ], + grid: (32, totalQ, tokenCount), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ, tokenCount]], + outputDTypes: [.float32] + )[0] + } + + /// Compute weighted sum of values from packed codebook indices. + /// + /// Result is in ROTATED space, caller must apply inverse rotation. + /// + /// - Returns: [totalHeads, D] float32 (rotated space) + static func mseWeightedSum( + weights: MLXArray, + packed: MLXArray, + norms: MLXArray, + codebook: MLXArray, + tokenCount: Int, + repeatCount: Int, + bits: Int, + dim: Int, queryChunkLength: Int = 1 + ) -> MLXArray { + let pw = TurboQuantPacking.packedWidth(count: dim, bits: bits) + let key = "\(bits)_\(dim)" + + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = valueKernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_value_\(bits)_\(dim)", + inputNames: ["weights", "packed", "norms", "codebook", "params"], + outputNames: ["output"], + source: TurboQuantMetalKernels.valueKernelSource, + ensureRowContiguous: true + ) + lock.lock() + valueKernels[key] = k + lock.unlock() + kernel = k + } + + let totalHeads = weights.dim(0) + let dimBlocks = (dim + 31) / 32 + + let params = MLXArray([UInt32(tokenCount), UInt32(repeatCount), UInt32(queryChunkLength)]) + return kernel( + [f32(weights), packed, f32(norms), f32(codebook), params], + template: [ + ("Bits", bits), ("Dim", dim), ("PackedWidth", pw), + ], + grid: (32, totalHeads, dimBlocks), + threadGroup: (32, 1, 1), + outputShapes: [[totalHeads, dim]], + outputDTypes: [.float32] + )[0] + } +} diff --git a/Libraries/MLXVLM/Models/Gemma4.swift b/Libraries/MLXVLM/Models/Gemma4.swift index 85252c7f8..11245a1a5 100644 --- a/Libraries/MLXVLM/Models/Gemma4.swift +++ b/Libraries/MLXVLM/Models/Gemma4.swift @@ -9,6 +9,7 @@ import MLXNN private enum Gemma4Error: LocalizedError { case imageTokenCountMismatch(expectedVisionTokens: Int, actualPromptTokens: Int) case multimodalTokenCountMismatch(kind: String, featureTokens: Int, promptTokens: Int) + case imagePlaceholderMismatch(images: Int, placeholders: Int) var errorDescription: String? { switch self { @@ -18,6 +19,9 @@ private enum Gemma4Error: LocalizedError { case .multimodalTokenCountMismatch(let kind, let featureTokens, let promptTokens): return "Gemma4 \(kind) token count mismatch: encoder produced \(featureTokens) soft tokens, but the prompt contains \(promptTokens) \(kind) tokens." + case .imagePlaceholderMismatch(let images, let placeholders): + return + "Gemma4 image placeholder mismatch: the request has \(images) images but the prompt contains at least \(placeholders) image placeholders." } } } @@ -87,24 +91,6 @@ private func gemma4OneHot(_ indices: MLXArray, numClasses: Int) -> MLXArray { expandedDimensions(indices, axis: -1) .== MLXArray(0 ..< numClasses) } -/// Average-pool kernel for Gemma 4's vision pooler. -/// -/// The padded patch tensor has length -/// `paddedPatchCount = outputLength × pool²` where `pool` is the -/// model's `pooling_kernel_size`. Recovering `pool` from these -/// two values yields `floor(sqrt(paddedPatchCount / outputLength))`. -/// -/// Matches HuggingFace's reference image processor (see -/// `image_processing_gemma4.py`: `max_patches = max_soft_tokens * -/// pooling_kernel_size**2`). -internal func gemma4VisionPoolingKernel( - paddedPatchCount: Int, outputLength: Int -) -> Int { - let safeLength = max(outputLength, 1) - let ratio = max(1, paddedPatchCount / safeLength) - return Int(sqrt(Double(ratio))) -} - private func gemma4RotateHalf(_ x: MLXArray) -> MLXArray { let half = x.shape[x.shape.count - 1] / 2 let x1 = x[.ellipsis, .. MLXArray { - let paddingPositions = patchPositions[0..., 0..., 0] .< 0 - let pooledHiddenStates = MLX.where( - expandedDimensions(paddingPositions, axis: -1), - MLXArray(0.0, dtype: hiddenStates.dtype), - hiddenStates - ) - let length = outputLength ?? defaultOutputLength - if pooledHiddenStates.dim(1) <= length { - return pooledHiddenStates * MLXArray(rootHiddenSize, dtype: pooledHiddenStates.dtype) - } - - let actualPositions = patchPositions[0, ..bLd", weights, pooledHiddenStates[0..., ..bLd", weights, hiddenStates) + .asType(hiddenStates.dtype) + return output * scale } } @@ -1883,9 +1861,7 @@ private final class Gemma4VisionTransformerModel: Module { private final class Gemma4VisionModel: Module { let config: Gemma4VisionConfiguration let patchSize: Int - let defaultOutputLength: Int let poolingKernelSize: Int - let maxPatches: Int @ModuleInfo(key: "patch_embedder") var patchEmbedder: Gemma4VisionPatchEmbedder @ModuleInfo(key: "encoder") var encoder: Gemma4VisionTransformerModel @@ -1896,10 +1872,7 @@ private final class Gemma4VisionModel: Module { init(config: Gemma4VisionConfiguration) { self.config = config self.patchSize = config.patchSize - self.defaultOutputLength = config.defaultOutputLength self.poolingKernelSize = config.poolingKernelSize - self.maxPatches = - config.defaultOutputLength * config.poolingKernelSize * config.poolingKernelSize self._patchEmbedder.wrappedValue = Gemma4VisionPatchEmbedder(config: config) self._encoder.wrappedValue = Gemma4VisionTransformerModel(config: config) self._pooler.wrappedValue = Gemma4VisionPooler(config: config) @@ -1910,32 +1883,24 @@ private final class Gemma4VisionModel: Module { super.init() } - private func patchPositions(batch: Int, height: Int, width: Int) -> (MLXArray, Int) { - let patchesH = height / patchSize - let patchesW = width / patchSize - let realCount = patchesH * patchesW - let paddedCount = max(maxPatches - realCount, 0) - - var values = [Int32]() - values.reserveCapacity(batch * (realCount + paddedCount) * 2) - - for _ in 0 ..< batch { - for y in 0 ..< patchesH { - for x in 0 ..< patchesW { - values.append(Int32(x)) - values.append(Int32(y)) - } - } - for _ in 0 ..< paddedCount { - values.append(-1) - values.append(-1) - } - } - - let count = realCount + paddedCount - return (MLXArray(values, [batch, count, 2]), realCount) + private func patchPositions(batch: Int, patchesH: Int, patchesW: Int) -> MLXArray { + // .xy indexing makes x vary fastest, matching the row-major patch + // order the embedder and pooler expect. + let grids = meshGrid([ + MLXArray.arange(patchesW, dtype: .int32), + MLXArray.arange(patchesH, dtype: .int32), + ]) + let positions = stacked([grids[0].flattened(), grids[1].flattened()], axis: 1) + .reshaped(1, patchesH * patchesW, 2) + return batch == 1 + ? positions + : broadcast(positions, to: [batch, patchesH * patchesW, 2]) } + /// Encodes a batch of same-sized images. Every patch is real (callers + /// slice padded canvases down to each image's true size first), so + /// attention is dense and the pooled output length falls out of the + /// patch grid: numPatches / poolingKernelSize². func callAsFunction(_ pixelValues: MLXArray) -> MLXArray { let pixels = if pixelValues.ndim == 3 { @@ -1944,32 +1909,17 @@ private final class Gemma4VisionModel: Module { pixelValues } let batch = pixels.dim(0) - let height = pixels.dim(2) - let width = pixels.dim(3) - let (patchPositions, realCount) = patchPositions(batch: batch, height: height, width: width) - - let realPositions = patchPositions[0..., .. 0 { - let pad = MLXArray.zeros( - [batch, paddingCount, hiddenStates.dim(2)], dtype: hiddenStates.dtype) - hiddenStates = concatenated([hiddenStates, pad], axis: 1) - } - - let validMask = patchPositions[0..., 0..., 0] .>= 0 - var attentionMask = - expandedDimensions(validMask, axis: 1) * expandedDimensions(validMask, axis: 2) - attentionMask = MLX.where( - attentionMask, - MLXArray(0.0, dtype: hiddenStates.dtype), - MLXArray(-Float.infinity, dtype: hiddenStates.dtype) - ) - attentionMask = expandedDimensions(attentionMask, axis: 1) - - hiddenStates = encoder(hiddenStates, positions: patchPositions, mask: attentionMask) - hiddenStates = pooler(hiddenStates, patchPositions: patchPositions, validCount: realCount) + let patchesH = pixels.dim(2) / patchSize + let patchesW = pixels.dim(3) / patchSize + let numPatches = patchesH * patchesW + let outputLength = max(numPatches / (poolingKernelSize * poolingKernelSize), 1) + + let patchPositions = patchPositions(batch: batch, patchesH: patchesH, patchesW: patchesW) + var hiddenStates = patchEmbedder(pixels, patchPositions: patchPositions) + hiddenStates = encoder(hiddenStates, positions: patchPositions, mask: nil) + hiddenStates = pooler( + hiddenStates, patchPositions: patchPositions, patchesW: patchesW, + outputLength: outputLength) if let standardizationBias, let standardizationScale { hiddenStates = (hiddenStates - standardizationBias) * standardizationScale @@ -2028,7 +1978,7 @@ public final class Gemma4: Module, VLMModel, KVCacheDimensionProvider { private func getInputEmbeddings( inputIds: MLXArray, - pixelValues: MLXArray? = nil + image: LMInput.ProcessedImage? = nil ) throws -> (MLXArray, MLXArray?) { var inputsEmbeds = languageModel.model.embedTokens(inputIds) inputsEmbeds = @@ -2050,11 +2000,32 @@ public final class Gemma4: Module, VLMModel, KVCacheDimensionProvider { perLayerInputs = languageModel.model.getPerLayerInputs(perLayerTokens) } - guard let pixelValues else { + guard let image else { return (inputsEmbeds, perLayerInputs) } - var imageFeatures = visionTower(pixelValues) + // Images keep their own aspect-preserving sizes: the processor + // zero-pads them onto a shared canvas and records each real size in + // frames. Slice each image back out, run the tower on it alone, and + // concatenate the pooled tokens in placeholder order. + let pixels = + if image.pixels.ndim == 3 { + expandedDimensions(image.pixels, axis: 0) + } else { + image.pixels + } + let frames = + image.frames + ?? Array(repeating: THW(1, pixels.dim(2), pixels.dim(3)), count: pixels.dim(0)) + var perImageFeatures: [MLXArray] = [] + for (index, frame) in frames.enumerated() { + let imagePixels = pixels[index ..< index + 1, 0..., .. PrepareResult { let convertedCache = cache.map { $0 } - if let imagePixels = input.image?.pixels { + if let image = input.image { let (inputsEmbeds, perLayerInputs) = try getInputEmbeddings( - inputIds: input.text.tokens, pixelValues: imagePixels) + inputIds: input.text.tokens, image: image) let result = languageModel( nil, cache: convertedCache, @@ -2640,30 +2611,26 @@ public struct Gemma4Processor: UserInputProcessor { self.tokenizer = tokenizer } - public func preprocess(images: [CIImage], processing: UserInput.Processing?) throws -> ( + public func preprocess(image: CIImage, processing: UserInput.Processing?) throws -> ( MLXArray, THW ) { - var userProcessing = processing ?? UserInput.Processing() - let targetSize = config.fixedSize - userProcessing.resize = targetSize - - let processedImages = images.map { image in - let processedImage = MediaProcessing.apply(image, processing: userProcessing) - let srgbImage = MediaProcessing.inSRGBToneCurveSpace(processedImage) - let resizedImage = MediaProcessing.resampleBicubic(srgbImage, to: targetSize) - let finalImage = - if config.doNormalize { - MediaProcessing.normalize( - resizedImage, mean: config.imageMeanTuple, std: config.imageStdTuple) - } else { - resizedImage - } - return MediaProcessing.asMLXArray(finalImage) - } - - let pixelValues = concatenated(processedImages) + let processedImage = MediaProcessing.apply(image, processing: processing) + let srgbImage = MediaProcessing.inSRGBToneCurveSpace(processedImage) + let targetSize = config.aspectPreservingTargetSize(for: srgbImage.extent.size) + let resizedImage = + srgbImage.extent.size == targetSize + ? srgbImage + : MediaProcessing.resampleBicubic(srgbImage, to: targetSize) + let finalImage = + if config.doNormalize { + MediaProcessing.normalize( + resizedImage, mean: config.imageMeanTuple, std: config.imageStdTuple) + } else { + resizedImage + } + let pixelValues = MediaProcessing.asMLXArray(finalImage) - return (pixelValues, THW(images.count, Int(targetSize.height), Int(targetSize.width))) + return (pixelValues, THW(1, Int(targetSize.height), Int(targetSize.width))) } public func prepare(input: UserInput) async throws -> LMInput { @@ -2676,24 +2643,48 @@ public struct Gemma4Processor: UserInputProcessor { var processedImage: LMInput.ProcessedImage? if !input.images.isEmpty { let imagePixelsAndFrames = try input.images.map { - try preprocess(images: [$0.asCIImage()], processing: input.processing) + try preprocess(image: $0.asCIImage(), processing: input.processing) + } + let frames = imagePixelsAndFrames.map { $0.1 } + + // Each image keeps its own aspect-preserving size. ProcessedImage + // carries one array, so zero-pad every image onto the largest + // canvas in the request; the model slices the real regions back + // out using frames. + let maxHeight = frames.map(\.h).max() ?? 0 + let maxWidth = frames.map(\.w).max() ?? 0 + let paddedPixels = imagePixelsAndFrames.map { pixels, frame in + frame.h == maxHeight && frame.w == maxWidth + ? pixels + : MLX.padded( + pixels, + widths: [ + 0, 0, .init((0, maxHeight - frame.h)), .init((0, maxWidth - frame.w)), + ]) } - let imagePixelsConcatenated = concatenated(imagePixelsAndFrames.map { $0.0 }) processedImage = LMInput.ProcessedImage( - pixels: imagePixelsConcatenated, - frames: imagePixelsAndFrames.map { $0.1 } - ) + pixels: concatenated(paddedPixels), frames: frames) + // Expand the i-th image placeholder to that image's soft token + // count: numPatches / poolingKernelSize². + let softTokenCounts = frames.map { config.softTokenCount(height: $0.h, width: $0.w) } var expandedTokens: [Int] = [] + var imageIndex = 0 for token in promptTokens { if token == config.imageTokenId { + guard imageIndex < softTokenCounts.count else { + throw Gemma4Error.imagePlaceholderMismatch( + images: softTokenCounts.count, placeholders: imageIndex + 1) + } expandedTokens.append(config.boiTokenId) expandedTokens.append( contentsOf: Array( - repeating: config.imageTokenId, count: config.imageSeqLength)) + repeating: config.imageTokenId, + count: softTokenCounts[imageIndex])) if let eoiTokenId = config.eoiTokenId { expandedTokens.append(eoiTokenId) } + imageIndex += 1 } else { expandedTokens.append(token) } @@ -2713,19 +2704,47 @@ public struct Gemma4ProcessorConfiguration: Codable, Sendable { public let imageMean: [CGFloat] public let imageStd: [CGFloat] public let imageSeqLength: Int - public let size: Gemma3ProcessorConfiguration.ImageSize? + public let maxSoftTokens: Int + public let patchSize: Int + public let poolingKernelSize: Int public let imageTokenId: Int public let boiTokenId: Int public let eoiTokenId: Int? + /// Image keys nested under `image_processor` in processor_config.json. + /// Repos that ship a flat preprocessor_config.json put the same keys at + /// the top level, which wins when both are present. + private struct ImageProcessorConfiguration: Codable { + let doNormalize: Bool? + let imageMean: [CGFloat]? + let imageStd: [CGFloat]? + let imageSeqLength: Int? + let maxSoftTokens: Int? + let patchSize: Int? + let poolingKernelSize: Int? + + enum CodingKeys: String, CodingKey { + case doNormalize = "do_normalize" + case imageMean = "image_mean" + case imageStd = "image_std" + case imageSeqLength = "image_seq_length" + case maxSoftTokens = "max_soft_tokens" + case patchSize = "patch_size" + case poolingKernelSize = "pooling_kernel_size" + } + } + enum CodingKeys: String, CodingKey { case processorClass = "processor_class" case doNormalize = "do_normalize" case imageMean = "image_mean" case imageStd = "image_std" case imageSeqLength = "image_seq_length" - case size + case maxSoftTokens = "max_soft_tokens" + case patchSize = "patch_size" + case poolingKernelSize = "pooling_kernel_size" + case imageProcessor = "image_processor" case imageTokenId = "image_token_id" case boiTokenId = "boi_token_id" case eoiTokenId = "eoi_token_id" @@ -2733,20 +2752,50 @@ public struct Gemma4ProcessorConfiguration: Codable, Sendable { public init(from decoder: any Swift.Decoder) throws { let c = try decoder.container(keyedBy: CodingKeys.self) + let nested = try c.decodeIfPresent( + ImageProcessorConfiguration.self, forKey: CodingKeys.imageProcessor) processorClass = try c.decode(String.self, forKey: CodingKeys.processorClass) - doNormalize = try c.decodeIfPresent(Bool.self, forKey: CodingKeys.doNormalize) ?? false + doNormalize = + try c.decodeIfPresent(Bool.self, forKey: CodingKeys.doNormalize) + ?? nested?.doNormalize ?? false imageMean = - try c.decodeIfPresent([CGFloat].self, forKey: CodingKeys.imageMean) ?? [0.5, 0.5, 0.5] + try c.decodeIfPresent([CGFloat].self, forKey: CodingKeys.imageMean) + ?? nested?.imageMean ?? [0.5, 0.5, 0.5] imageStd = - try c.decodeIfPresent([CGFloat].self, forKey: CodingKeys.imageStd) ?? [0.5, 0.5, 0.5] - imageSeqLength = try c.decodeIfPresent(Int.self, forKey: CodingKeys.imageSeqLength) ?? 280 - size = try c.decodeIfPresent( - Gemma3ProcessorConfiguration.ImageSize.self, forKey: CodingKeys.size) + try c.decodeIfPresent([CGFloat].self, forKey: CodingKeys.imageStd) + ?? nested?.imageStd ?? [0.5, 0.5, 0.5] + imageSeqLength = + try c.decodeIfPresent(Int.self, forKey: CodingKeys.imageSeqLength) + ?? nested?.imageSeqLength ?? 280 + maxSoftTokens = + try c.decodeIfPresent(Int.self, forKey: CodingKeys.maxSoftTokens) + ?? nested?.maxSoftTokens ?? 280 + patchSize = + try c.decodeIfPresent(Int.self, forKey: CodingKeys.patchSize) + ?? nested?.patchSize ?? 16 + poolingKernelSize = + try c.decodeIfPresent(Int.self, forKey: CodingKeys.poolingKernelSize) + ?? nested?.poolingKernelSize ?? 3 imageTokenId = try c.decodeIfPresent(Int.self, forKey: CodingKeys.imageTokenId) ?? 258_880 boiTokenId = try c.decodeIfPresent(Int.self, forKey: CodingKeys.boiTokenId) ?? 255_999 eoiTokenId = try c.decodeIfPresent(Int.self, forKey: CodingKeys.eoiTokenId) ?? 258_882 } + public func encode(to encoder: any Swift.Encoder) throws { + var c = encoder.container(keyedBy: CodingKeys.self) + try c.encode(processorClass, forKey: CodingKeys.processorClass) + try c.encode(doNormalize, forKey: CodingKeys.doNormalize) + try c.encode(imageMean, forKey: CodingKeys.imageMean) + try c.encode(imageStd, forKey: CodingKeys.imageStd) + try c.encode(imageSeqLength, forKey: CodingKeys.imageSeqLength) + try c.encode(maxSoftTokens, forKey: CodingKeys.maxSoftTokens) + try c.encode(patchSize, forKey: CodingKeys.patchSize) + try c.encode(poolingKernelSize, forKey: CodingKeys.poolingKernelSize) + try c.encode(imageTokenId, forKey: CodingKeys.imageTokenId) + try c.encode(boiTokenId, forKey: CodingKeys.boiTokenId) + try c.encodeIfPresent(eoiTokenId, forKey: CodingKeys.eoiTokenId) + } + public var imageMeanTuple: (CGFloat, CGFloat, CGFloat) { (imageMean[0], imageMean[1], imageMean[2]) } @@ -2755,12 +2804,48 @@ public struct Gemma4ProcessorConfiguration: Codable, Sendable { (imageStd[0], imageStd[1], imageStd[2]) } - public var fixedSize: CGSize { - if let size { - return CGSize(width: size.width, height: size.height) + /// Soft tokens the vision tower produces for an image of the given + /// (already resized) pixel dimensions. + public func softTokenCount(height: Int, width: Int) -> Int { + ((height / patchSize) * (width / patchSize)) / (poolingKernelSize * poolingKernelSize) + } + + /// Port of the Python Gemma4ImageProcessor's aspect-ratio preserving + /// resize: the largest dimensions that (a) stay within the patch budget + /// maxSoftTokens * poolingKernelSize², and (b) keep both sides divisible + /// by poolingKernelSize * patchSize, so the pooling kernel is exact and + /// the pooled grid covers the image fully at any aspect ratio. + /// + /// Note the config's `size` entry is deliberately ignored, as in the + /// Python reference — models ship a vestigial 224x224 there. + public func aspectPreservingTargetSize(for imageSize: CGSize) -> CGSize { + let kernelArea = poolingKernelSize * poolingKernelSize + let maxPatches = maxSoftTokens * kernelArea + let sideMultiple = poolingKernelSize * patchSize + let height = Double(imageSize.height) + let width = Double(imageSize.width) + + let targetPixels = Double(maxPatches * patchSize * patchSize) + let factor = (targetPixels / max(height * width, 1)).squareRoot() + var targetHeight = Int((factor * height / Double(sideMultiple)).rounded(.down)) + var targetWidth = Int((factor * width / Double(sideMultiple)).rounded(.down)) + + // One side can floor to zero for extreme aspect ratios (both cannot: + // their product is pinned near maxPatches, far above 1). Clamp it to + // one pooling cell and cap the long side at the full token budget. + let maxSideLength = maxSoftTokens + if targetHeight == 0 { + targetHeight = 1 + targetWidth = min(Int((width / max(height, 1)).rounded(.down)), maxSideLength) + targetWidth = max(targetWidth, 1) + } else if targetWidth == 0 { + targetWidth = 1 + targetHeight = min(Int((height / max(width, 1)).rounded(.down)), maxSideLength) + targetHeight = max(targetHeight, 1) } - // 800x800 keeps the patch count under Gemma4's 280 * 3^2 vision budget. - return CGSize(width: 800, height: 800) + + return CGSize( + width: targetWidth * sideMultiple, height: targetHeight * sideMultiple) } } diff --git a/Tests/MLXLMTests/Gemma4PoolingTests.swift b/Tests/MLXLMTests/Gemma4PoolingTests.swift deleted file mode 100644 index 0eee9cfd6..000000000 --- a/Tests/MLXLMTests/Gemma4PoolingTests.swift +++ /dev/null @@ -1,84 +0,0 @@ -// Copyright © 2026 Apple Inc. - -import Foundation -import Testing - -@testable import MLXVLM - -/// Unit tests for `gemma4VisionPoolingKernel`. Pins the recovered -/// kernel to Gemma 4's `pooling_kernel_size` (3) for every -/// supported soft-token budget. -struct Gemma4VisionPoolingTests { - - /// `pooling_kernel_size` for Gemma 4 (per HF + the model - /// config's `pooling_kernel_size` field). - private static let expectedPoolingKernel = 3 - - @Test( - "Kernel matches pooling_kernel_size (3) for every supported budget", - arguments: [70, 140, 280, 560, 1120]) - func testKernelMatchesPoolingKernelSize(budget: Int) { - // Per HF: max_patches = max_soft_tokens * pool², so the - // padded patch tensor fed into the pooler has length - // budget * pool². The kernel recovered from those values - // must equal pool. - let pool = Self.expectedPoolingKernel - let paddedPatchCount = budget * pool * pool - - let kernel = gemma4VisionPoolingKernel( - paddedPatchCount: paddedPatchCount, outputLength: budget) - - #expect( - kernel == pool, - """ - Expected pooling kernel \(pool) for budget \(budget) \ - (padded patch count \(paddedPatchCount)), got \(kernel). \ - Kernel must match Gemma 4's `pooling_kernel_size` config field for \ - every supported soft-token budget. - """) - } - - @Test("Kernel is robust to degenerate inputs") - func testKernelHandlesDegenerateInputs() { - // A zero `outputLength` must not divide by zero — the - // function falls back to treating it as 1, matching the - // upstream pattern. The exact returned value at degenerate - // inputs is unspecified; the assertion is "does not crash". - _ = gemma4VisionPoolingKernel(paddedPatchCount: 0, outputLength: 0) - _ = gemma4VisionPoolingKernel(paddedPatchCount: 100, outputLength: 0) - - // A zero `paddedPatchCount` with a real `outputLength` - // yields kernel = 1 (no patches to pool yields the minimum - // kernel). - #expect(gemma4VisionPoolingKernel(paddedPatchCount: 0, outputLength: 280) == 1) - } - - /// Documents the prior bug. The original formula divided the - /// real (un-padded) patch count by `outputLength` and floored - /// the square root. For a 768×768 image — the aligned default - /// for Gemma 4 (divisible by `patch_size × pooling_kernel_size - /// = 48`) — this returns 2 even though `pooling_kernel_size` - /// is 3. - @Test("Pre-fix formula returns kernel=2 at the aligned default size") - func testPreFixFormulaReproducesBug() { - // 768×768 image at patch_size=16 → 48×48 = 2304 real patches. - let realPatchCount = 48 * 48 - let outputLength = 280 - - // The old formula, preserved here for documentation. - let safeLength = max(outputLength, 1) - let ratio = max(1, realPatchCount / safeLength) - let oldKernel = Int(sqrt(Double(ratio))) - - #expect( - oldKernel == 2, - "Sanity check on the pre-fix formula: should reproduce the bug.") - #expect( - oldKernel != Self.expectedPoolingKernel, - """ - The pre-fix kernel (\(oldKernel)) does not match Gemma 4's \ - documented `pooling_kernel_size` of \(Self.expectedPoolingKernel). \ - This is the bug `gemma4VisionPoolingKernel` was introduced to fix. - """) - } -} diff --git a/Tests/MLXLMTests/Gemma4ResizeTests.swift b/Tests/MLXLMTests/Gemma4ResizeTests.swift new file mode 100644 index 000000000..26cdf6f66 --- /dev/null +++ b/Tests/MLXLMTests/Gemma4ResizeTests.swift @@ -0,0 +1,113 @@ +// Copyright © 2026 Apple Inc. + +import CoreGraphics +import Foundation +import Testing + +@testable import MLXVLM + +/// Unit tests for `Gemma4ProcessorConfiguration.aspectPreservingTargetSize`. +/// +/// The resize must produce dimensions whose patch grid pools exactly: +/// both sides divisible by `pooling_kernel_size × patch_size`, and the +/// patch count within `max_soft_tokens × pooling_kernel_size²`. When +/// those invariants hold the pooler's kernel derivation is exact for +/// every aspect ratio and the pooled grid covers the image fully — +/// the bug class where real patches mapped past the one-hot budget +/// (and the bottom of the image was silently dropped) cannot occur. +struct Gemma4ResizeTests { + + /// Defaults: max_soft_tokens 280, patch_size 16, pooling_kernel_size 3. + private static let config: Gemma4ProcessorConfiguration = { + let json = #"{"processor_class": "Gemma4Processor"}"# + return try! JSONDecoder().decode( + Gemma4ProcessorConfiguration.self, from: Data(json.utf8)) + }() + + private static let sampleSizes: [(Int, Int)] = [ + (800, 800), // the size that previously truncated the pooled grid + (768, 768), + (960, 672), + (672, 960), + (1170, 2532), // portrait screenshot + (2532, 1170), + (30, 30), // tiny + (8000, 20), // extreme panorama + (20, 8000), + (1, 1), + (48, 48), + (4032, 3024), // camera photo + ] + + @Test("Both sides are divisible by pooling_kernel_size × patch_size", arguments: sampleSizes) + func testSideAlignment(width: Int, height: Int) { + let config = Self.config + let sideMultiple = config.poolingKernelSize * config.patchSize + let target = config.aspectPreservingTargetSize( + for: CGSize(width: width, height: height)) + + #expect(Int(target.width) % sideMultiple == 0) + #expect(Int(target.height) % sideMultiple == 0) + #expect(Int(target.width) >= sideMultiple) + #expect(Int(target.height) >= sideMultiple) + } + + @Test("Patch count stays within max_soft_tokens × kernel²", arguments: sampleSizes) + func testPatchBudget(width: Int, height: Int) { + let config = Self.config + let target = config.aspectPreservingTargetSize( + for: CGSize(width: width, height: height)) + let patches = + (Int(target.width) / config.patchSize) * (Int(target.height) / config.patchSize) + let kernelArea = config.poolingKernelSize * config.poolingKernelSize + + #expect(patches <= config.maxSoftTokens * kernelArea) + // Side alignment guarantees the pooling kernel divides the grid + // exactly, so every soft token pools kernel² real patches. + #expect(patches % kernelArea == 0) + #expect( + config.softTokenCount( + height: Int(target.height), width: Int(target.width)) <= config.maxSoftTokens) + } + + @Test("Near-square inputs use the full token budget") + func testFullBudgetAtSquare() { + let config = Self.config + let target = config.aspectPreservingTargetSize(for: CGSize(width: 800, height: 800)) + let softTokens = config.softTokenCount( + height: Int(target.height), width: Int(target.width)) + + // 800×800 → 768×768 → 48×48 patches → 256 soft tokens. The exact + // value matters less than the invariant pair: close to the budget, + // never over it. + #expect(softTokens > config.maxSoftTokens / 2) + #expect(softTokens <= config.maxSoftTokens) + } + + @Test("Aspect ratio is approximately preserved for ordinary images") + func testAspectRatioPreserved() { + let config = Self.config + let input = CGSize(width: 1170, height: 2532) + let target = config.aspectPreservingTargetSize(for: input) + + let inputRatio = input.width / input.height + let targetRatio = target.width / target.height + // Rounding to 48-pixel multiples bounds the ratio drift. + #expect(abs(inputRatio - targetRatio) / inputRatio < 0.25) + #expect(target.height > target.width) + } + + @Test("Extreme aspect ratios clamp to one pooling cell on the short side") + func testExtremeAspectClamps() { + let config = Self.config + let sideMultiple = config.poolingKernelSize * config.patchSize + + let panorama = config.aspectPreservingTargetSize(for: CGSize(width: 8000, height: 20)) + #expect(Int(panorama.height) == sideMultiple) + #expect(Int(panorama.width) <= config.maxSoftTokens * sideMultiple) + + let tower = config.aspectPreservingTargetSize(for: CGSize(width: 20, height: 8000)) + #expect(Int(tower.width) == sideMultiple) + #expect(Int(tower.height) <= config.maxSoftTokens * sideMultiple) + } +} From b3fa5d3dfc20162f64c22263a5f2b78fde98be94 Mon Sep 17 00:00:00 2001 From: TheTom Date: Thu, 16 Jul 2026 15:07:20 -0500 Subject: [PATCH 2/3] Add TurboQuant tests Codec round trips and pinned codebooks, rotation orthogonality, packing, kernel-vs-reference attention parity (GQA repeat factors 1/2/3, non-power-of-two head dims), calibrated-encode kernel exact index parity, prompt-cache round trips for all state layouts including calibration state, compressed partial trim with continued decode, speculative decoding with a turbo scheme, allocation-boundary growth, boundary-protection gating, mixed rotating/standard cache conversion, and end-to-end generation on tiny random-weight models. CI-runnable, no checkpoint downloads. --- Tests/MLXLMTests/TurboQuantTests.swift | 2145 ++++++++++++++++++++++++ 1 file changed, 2145 insertions(+) create mode 100644 Tests/MLXLMTests/TurboQuantTests.swift diff --git a/Tests/MLXLMTests/TurboQuantTests.swift b/Tests/MLXLMTests/TurboQuantTests.swift new file mode 100644 index 000000000..4b1edb539 --- /dev/null +++ b/Tests/MLXLMTests/TurboQuantTests.swift @@ -0,0 +1,2145 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import MLXLLM +import MLXNN +import Testing +import XCTest + +@testable import MLXLMCommon + +// MARK: - Codebook Tests + +@Suite("TurboQuant Codebook") +struct TurboQuantCodebookTests { + + @Test func codebookGeneration_3bit() { + let cb = TurboQuantCodebook.codebook(dim: 128, bits: 3) + #expect(cb.count == 8) + let vals = cb.asArray(Float.self) + for i in 0 ..< vals.count - 1 { + #expect(vals[i] <= vals[i + 1], "Codebook not sorted") + } + } + + @Test func codebookGeneration_4bit() { + let cb = TurboQuantCodebook.codebook(dim: 128, bits: 4) + #expect(cb.count == 16) + let vals = cb.asArray(Float.self) + for v in vals { + #expect(v >= -1.0 && v <= 1.0, "Centroid \(v) out of range") + } + } + + @Test func boundaryCount() { + let b = TurboQuantCodebook.boundaries(dim: 128, bits: 3) + #expect(b.count == 7) // 2^3 - 1 boundaries for 8 centroids + } + + @Test func codebookDeterminism() { + let cb1 = TurboQuantCodebook.codebook(dim: 64, bits: 3) + let cb2 = TurboQuantCodebook.codebook(dim: 64, bits: 3) + let v1 = cb1.asArray(Float.self) + let v2 = cb2.asArray(Float.self) + #expect(v1 == v2, "Codebook should be deterministic") + } +} + +// MARK: - Rotation Tests + +@Suite("TurboQuant Rotation") +struct TurboQuantRotationTests { + + @Test func rotationOrthogonality() { + let dim = 64 + let R = TurboQuantRotation.rotationMatrix(dim: dim, seed: 42) + #expect(R.shape == [dim, dim]) + let RRt = matmul(R, R.transposed()) + let identity = MLXArray.eye(dim) + let diff = MLX.abs(RRt - identity) + let maxDiff = diff.max().item(Float.self) + // MLX QR is f32 internally; Householder orthogonality error at d=64 + // is ~1e-3. The dense path's matched-norm scales absorb this. + #expect(maxDiff < 2e-3, "R @ R^T differs from I by \(maxDiff)") + } + + @Test func rotationDeterminism() { + let R1 = TurboQuantRotation.rotationMatrix(dim: 32, seed: 123) + let R2 = TurboQuantRotation.rotationMatrix(dim: 32, seed: 123) + let diff = MLX.abs(R1 - R2).max().item(Float.self) + #expect(diff < 1e-6, "Same seed should produce same rotation") + } +} + +// MARK: - Bit Packing Tests + +@Suite("TurboQuant Bit Packing") +struct TurboQuantPackingTests { + + @Test func packedWidth() { + #expect(TurboQuantPacking.packedWidth(count: 128, bits: 4) == 16) + #expect(TurboQuantPacking.packedWidth(count: 128, bits: 3) == 12) + #expect(TurboQuantPacking.packedWidth(count: 128, bits: 1) == 4) + } + + @Test func packUnpack_3bit() { + let indices = MLXArray([0, 1, 2, 3, 4, 5, 6, 7, 7, 6, 5, 4] as [UInt32]).reshaped([1, 12]) + let packed = TurboQuantPacking.packLowBit(indices, bits: 3) + let unpacked = TurboQuantPacking.unpackLowBit(packed, bits: 3, count: 12) + let orig = indices.asArray(UInt32.self) + let result = unpacked.asArray(UInt32.self) + #expect(orig == result, "3-bit pack/unpack round-trip failed") + } + + @Test func packUnpack_1bit() { + let indices = MLXArray([0, 1, 0, 1, 1, 0, 1, 0] as [UInt32]).reshaped([1, 8]) + let packed = TurboQuantPacking.packLowBit(indices, bits: 1) + let unpacked = TurboQuantPacking.unpackLowBit(packed, bits: 1, count: 8) + let orig = indices.asArray(UInt32.self) + let result = unpacked.asArray(UInt32.self) + #expect(orig == result, "1-bit pack/unpack round-trip failed") + } + + @Test func packUnpack_4bit() { + let indices = MLXArray([0, 5, 10, 15, 3, 7, 11, 14] as [UInt32]).reshaped([1, 8]) + let packed = TurboQuantPacking.packLowBit(indices, bits: 4) + let unpacked = TurboQuantPacking.unpackLowBit(packed, bits: 4, count: 8) + let orig = indices.asArray(UInt32.self) + let result = unpacked.asArray(UInt32.self) + #expect(orig == result, "4-bit pack/unpack round-trip failed") + } +} + +// MARK: - MSE Codec Tests + +@Suite("TurboQuant MSE Codec") +struct TurboQuantMSECodecTests { + + @Test func encodeDecodeRoundTrip_4bit() { + let codec = MSECodec(dim: 32, bits: 4, seed: 42) + let vectors = MLXRandom.normal([1, 2, 4, 32]) + eval(vectors) + + let state = codec.encode(vectors) + let decoded = codec.decode(state) + + #expect(state.tokenCount == 4) + #expect(state.dim == 32) + #expect(state.bits == 4) + + let mse = ((vectors - decoded) * (vectors - decoded)).mean().item(Float.self) + #expect(mse < 0.5, "4-bit MSE too high: \(mse)") + } + + @Test func encodeDecodeRoundTrip_3bit() { + let codec = MSECodec(dim: 32, bits: 3, seed: 42) + let vectors = MLXRandom.normal([1, 2, 4, 32]) + eval(vectors) + + let state = codec.encode(vectors) + let decoded = codec.decode(state) + + let mse = ((vectors - decoded) * (vectors - decoded)).mean().item(Float.self) + #expect(mse < 1.0, "3-bit MSE too high: \(mse)") + } + + @Test func cosineSimilarity_4bit() { + let codec = MSECodec(dim: 64, bits: 4, seed: 42) + let vectors = MLXRandom.normal([1, 1, 8, 64]) + eval(vectors) + + let decoded = codec.decode(codec.encode(vectors)) + let dot = (vectors * decoded).sum(axis: -1) + let normOrig = sqrt((vectors * vectors).sum(axis: -1)) + let normDec = sqrt((decoded * decoded).sum(axis: -1)) + let cosSim = dot / (normOrig * normDec + 1e-8) + let avgCosSim = cosSim.mean().item(Float.self) + + #expect(avgCosSim > 0.95, "4-bit cosine similarity too low: \(avgCosSim)") + } + + @Test func boundaryQuantizeMatchesArgmin() { + // Verify boundary-based quantization gives same result as argmin + let codec = MSECodec(dim: 16, bits: 3, seed: 42) + let vectors = MLXRandom.normal([1, 1, 4, 16]) + eval(vectors) + + let norms = sqrt((vectors * vectors).sum(axis: -1)) + let unit = vectors / expandedDimensions(maximum(norms, MLXArray(Float(1e-8))), axis: -1) + let rotated = matmul(unit, codec.rotationT) + + // Boundary quantize + let boundaryIdx = codec.boundaryQuantize(rotated) + + // Argmin quantize (naive but correct) + let expanded = expandedDimensions(rotated, axis: -1) + let cb = codec.codebook.reshaped([1, 1, 1, 1, -1]) + let distances = MLX.abs(expanded - cb) + let argminIdx = argMin(distances, axis: -1) + + let b = boundaryIdx.asArray(UInt32.self) + let a = argminIdx.asArray(UInt32.self) + #expect(b == a, "Boundary quantize should match argmin") + } + + @Test func fwhtRoundTrip() { + // FWHT forward → inverse should recover original vectors + let signs = TurboQuantRotation.whtSigns(dim: 128, seed: 42) + + let vectors = MLXRandom.normal([2, 4, 8, 128]) + eval(vectors) + + let rotated = TurboQuantRotation.fwhtForward(vectors, signs: signs) + let recovered = TurboQuantRotation.fwhtInverse(rotated, signs: signs) + eval(recovered) + + let diff = MLX.abs(vectors - recovered).max().item(Float.self) + #expect(diff < 1e-4, "FWHT round-trip max diff: \(diff)") + } + + @Test func whtRotationOrthogonality() { + // WHT rotation H*D/sqrt(d) should be orthogonal: Π·Π^T ≈ I + let codec = MSECodec(dim: 128, bits: 3, seed: 42) + #expect(codec.useWHT, "dim=128 should use WHT") + + let product = matmul(codec.rotation, codec.rotationT) + let identity = MLXArray.identity(128) + let diff = MLX.abs(product - identity).max().item(Float.self) + #expect(diff < 1e-4, "WHT rotation should be orthogonal, max diff: \(diff)") + } + + @Test func whtEncodeDecodeRoundTrip() { + // Verify encode/decode with WHT rotation produces reasonable reconstruction + let codec = MSECodec(dim: 128, bits: 4, seed: 42) + #expect(codec.useWHT, "dim=128 should use WHT") + + let vectors = MLXRandom.normal([1, 1, 8, 128]) + eval(vectors) + + let state = codec.encode(vectors) + let decoded = codec.decode(state) + + let cosDot = (vectors * decoded).sum(axis: -1) + let normOrig = sqrt((vectors * vectors).sum(axis: -1)) + let normDec = sqrt((decoded * decoded).sum(axis: -1)) + let cosSim = cosDot / (normOrig * normDec + 1e-8) + let avgCosSim = cosSim.mean().item(Float.self) + + #expect(avgCosSim > 0.90, "WHT 4-bit cosine similarity too low: \(avgCosSim)") + } + + @Test func normCorrectionImprovesNormAccuracy() { + // Norm correction ensures the reconstructed vector's L2 norm matches the original. + // This improves attention score accuracy (dot products), which is what matters + // for perplexity. Element-wise MSE may not improve, but norm accuracy should. + // dim 96 is not a power of two, so the codec takes the dense-rotation + // path, the only path that applies norm correction (WHT preserves + // norms exactly and skips it). + let codec = MSECodec(dim: 96, bits: 3, seed: 42) + let vectors = MLXRandom.normal([1, 1, 64, 96]) + eval(vectors) + + // Original norms + let origNorms = sqrt((vectors * vectors).sum(axis: -1)) + + // Encode with norm correction (current implementation) + let state = codec.encode(vectors) + let decodedCorrected = codec.decode(state) + let correctedNorms = sqrt((decodedCorrected * decodedCorrected).sum(axis: -1)) + + // Norm error with correction + let normErrorCorrected = MLX.abs(origNorms - correctedNorms).mean().item(Float.self) + + // Manually encode WITHOUT norm correction for comparison + let norms = sqrt((vectors * vectors).sum(axis: -1)) + let safeNorms = maximum(norms, MLXArray(Float(1e-8))) + let unit = vectors / expandedDimensions(safeNorms, axis: -1) + let rotated = matmul(unit, codec.rotationT) + let indices = codec.boundaryQuantize(rotated) + let packed = TurboQuantPacking.packLowBit(indices, bits: 3) + let uncorrectedState = MSECodecState( + norms: norms, packedIndices: packed, tokenCount: 64, dim: 96, bits: 3) + let decodedUncorrected = codec.decode(uncorrectedState) + let uncorrectedNorms = sqrt((decodedUncorrected * decodedUncorrected).sum(axis: -1)) + + // Norm error without correction + let normErrorUncorrected = MLX.abs(origNorms - uncorrectedNorms).mean().item(Float.self) + + #expect( + normErrorCorrected < normErrorUncorrected, + "Norm-corrected norm error (\(normErrorCorrected)) should be lower than uncorrected (\(normErrorUncorrected))" + ) + + // Also verify dot product accuracy improves (key for attention scoring) + // Dot product of original with itself = ||v||², so compare dot products + let dotCorrected = (vectors * decodedCorrected).sum(axis: -1) + let dotUncorrected = (vectors * decodedUncorrected).sum(axis: -1) + let dotOriginal = (vectors * vectors).sum(axis: -1) + let dotErrorCorrected = MLX.abs(dotOriginal - dotCorrected).mean().item(Float.self) + let dotErrorUncorrected = MLX.abs(dotOriginal - dotUncorrected).mean().item(Float.self) + + #expect( + dotErrorCorrected < dotErrorUncorrected, + "Norm-corrected dot product error (\(dotErrorCorrected)) should be lower than uncorrected (\(dotErrorUncorrected))" + ) + } + + /// Key calibration: on synthetic data with a 10x per-dimension variance + /// spread in the post-rotation space (the exact anisotropy calibration + /// targets), quantize(rotate(k) / s) with per-vector norm recomputed after + /// the scale divide should round-trip at least as well as the uncalibrated + /// quantize(rotate(k)). + @Test func keyCalibrationImprovesAnisotropicRoundTrip() { + let dim = 64 + let codec = MSECodec(dim: dim, bits: 4, seed: 42) + #expect(codec.useWHT, "dim=64 should use WHT") + + // Build vectors whose post-rotation coordinates have a 10x std spread: + // first half of dims sqrt(10)x, second half 1x. Map the anisotropic + // pattern back through the codec's own rotation so encode's internal + // rotate(unit(vectors)) recovers (up to per-row norm) this shape. + var stdVec = [Float](repeating: 1.0, count: dim) + for i in 0 ..< dim / 2 { stdVec[i] = Float(10.0).squareRoot() } + let stdArray = MLXArray(stdVec) + let rotatedSynthetic = MLXRandom.normal([1, 1, 512, dim], key: MLXRandom.key(7)) * stdArray + let vectors = matmul(rotatedSynthetic, codec.rotation) + eval(vectors) + + // Calibration scale, mirrors TurboQuantKVCache.computeKeyCalibrationScale: + // std per rotated dimension, normalized to mean 1, clamped to [0.25, 4.0]. + let norms = sqrt((vectors * vectors).sum(axis: -1)) + let unit = vectors / expandedDimensions(maximum(norms, MLXArray(Float(1e-8))), axis: -1) + let rotated = matmul(unit, codec.rotationT).reshaped([-1, dim]) + let perDimStd = std(rotated, axis: 0) + let meanStd = maximum(perDimStd.mean(), MLXArray(Float(1e-8))) + let scale = clip(perDimStd / meanStd, min: Float(0.25), max: Float(4.0)) + eval(scale) + + let uncalibratedDecoded = codec.decode(codec.encode(vectors)) + let uncalibratedMSE = + ((vectors - uncalibratedDecoded) * (vectors - uncalibratedDecoded)).mean().item( + Float.self) + + let calibratedDecoded = codec.decode(codec.encode(vectors, scale: scale), scale: scale) + let calibratedMSE = + ((vectors - calibratedDecoded) * (vectors - calibratedDecoded)).mean().item( + Float.self) + + #expect( + calibratedMSE <= uncalibratedMSE, + "calibrated MSE \(calibratedMSE) should not exceed uncalibrated \(uncalibratedMSE)") + } +} + +// MARK: - KV Cache Tests + +@Suite("TurboQuantKVCache") +struct TurboQuantKVCacheTests { + + @Test func cacheUpdate() { + let cache = TurboQuantKVCache(bits: 4) + let keys = MLXRandom.normal([1, 4, 8, 64]) + let values = MLXRandom.normal([1, 4, 8, 64]) + eval(keys, values) + + let (outKeys, outValues) = cache.update(keys: keys, values: values) + #expect(cache.offset == 8) + #expect(outKeys.shape == [1, 4, 8, 64]) + #expect(outValues.shape == [1, 4, 8, 64]) + } + + @Test func cacheIncrementalUpdate() { + let cache = TurboQuantKVCache(bits: 4) + + let k1 = MLXRandom.normal([1, 2, 4, 32]) + let v1 = MLXRandom.normal([1, 2, 4, 32]) + eval(k1, v1) + let (_, _) = cache.update(keys: k1, values: v1) + #expect(cache.offset == 4) + + let k2 = MLXRandom.normal([1, 2, 1, 32]) + let v2 = MLXRandom.normal([1, 2, 1, 32]) + eval(k2, v2) + let (outK, _) = cache.update(keys: k2, values: v2) + #expect(cache.offset == 5) + #expect(outK.shape == [1, 2, 5, 32]) + } + + @Test func cacheTrim() { + let cache = TurboQuantKVCache(bits: 4) + let keys = MLXRandom.normal([1, 2, 8, 32]) + let values = MLXRandom.normal([1, 2, 8, 32]) + eval(keys, values) + _ = cache.update(keys: keys, values: values) + #expect(cache.offset == 8) + + let trimmed = cache.trim(3) + #expect(trimmed == 3) + #expect(cache.offset == 5) + } + + @Test func cacheState() { + let cache = TurboQuantKVCache(bits: 4) + let keys = MLXRandom.normal([1, 2, 4, 32]) + let values = MLXRandom.normal([1, 2, 4, 32]) + eval(keys, values) + _ = cache.update(keys: keys, values: values) + + let state = cache.state + // In raw phase (prefill): 2 arrays (rawKeys, rawValues) + // In compressed phase: 4 arrays (keyPacked, keyNorms, valPacked, valNorms) + #expect( + state.count == 2 || state.count == 4, + "State should have 2 or 4 arrays, got \(state.count)") + } + + @Test func cacheIsTrimmable() { + let cache = TurboQuantKVCache(bits: 4) + #expect(cache.isTrimmable == true) + } + + /// Regression test: asymmetric bits (e.g. turbo3v2 = 3-bit K, 2-bit V) must encode + /// values with valueBits, not the legacy `bits` field. A mismatch causes packed width + /// errors during compressRawCache → reshape. + @Test func cacheAsymmetricCompression() { + // turbo3v2 config: keyBits=3, valueBits=2 + let cache = TurboQuantKVCache(bits: 3, keyBits: 3, valueBits: 2) + let B = 1 + let H = 4 // KV heads + let T = 32 // prefill tokens + let D = 128 + + // Phase 1: Prefill (raw FP16 storage) + let keys = MLXRandom.normal([B, H, T, D]) + let values = MLXRandom.normal([B, H, T, D]) + eval(keys, values) + let (_, _) = cache.update(keys: keys, values: values) + #expect(cache.offset == T) + + // Phase 2: First decode triggers compressRawCache() + let newKey = MLXRandom.normal([B, H, 1, D]) + let newVal = MLXRandom.normal([B, H, 1, D]) + let queries = MLXRandom.normal([B, H * 2, 1, D]) // nQHeads = 2 * nKVHeads (GQA) + eval(newKey, newVal, queries) + + let output = cache.compressedAttention( + queries: queries, keys: newKey, values: newVal, + scale: 1.0 / sqrt(Float(D)) + ) + eval(output) + #expect(output.shape == [B, H * 2, 1, D]) + #expect(cache.isCompressed) + + // Phase 3: Several more decode steps to confirm incremental encode also works + for _ in 0 ..< 5 { + let dk = MLXRandom.normal([B, H, 1, D]) + let dv = MLXRandom.normal([B, H, 1, D]) + let dq = MLXRandom.normal([B, H * 2, 1, D]) + eval(dk, dv, dq) + let out = cache.compressedAttention( + queries: dq, keys: dk, values: dv, + scale: 1.0 / sqrt(Float(D)) + ) + eval(out) + #expect(out.shape == [B, H * 2, 1, D]) + } + #expect(cache.offset == T + 6) + } + + @Test func cacheAllBitWidths() { + for bits in 2 ... 4 { + let cache = TurboQuantKVCache(bits: bits) + let keys = MLXRandom.normal([1, 2, 4, 32]) + let values = MLXRandom.normal([1, 2, 4, 32]) + eval(keys, values) + let (outKeys, outValues) = cache.update(keys: keys, values: values) + #expect(cache.offset == 4, "\(bits)-bit cache offset wrong") + #expect(outKeys.shape == [1, 2, 4, 32], "\(bits)-bit key shape wrong") + #expect(outValues.shape == [1, 2, 4, 32], "\(bits)-bit value shape wrong") + } + } +} + +// MARK: - TurboFlashAttention Tests + +@Suite("TurboFlashAttention") +struct TurboFlashAttentionTests { + + /// Validate that fused TurboFlashAttention produces the same output as + /// separated Score → Softmax → Value kernels. + @Test func flashMatchesSeparated() { + let dim = 128 + let keyBits = 4 + let valueBits = 4 + let nQHeads = 8 + let nKVHeads = 4 + let tokenCount = 64 + let repeatCount = nQHeads / nKVHeads + + let keyCodec = MSECodec(dim: dim, bits: keyBits, seed: 42) + let valCodec = MSECodec(dim: dim, bits: valueBits, seed: 43) + + // Generate random KV cache: encode random vectors + let rawKeys = MLXRandom.normal([nKVHeads, tokenCount, dim]) + let rawValues = MLXRandom.normal([nKVHeads, tokenCount, dim]) + eval(rawKeys, rawValues) + + // Encode K and V + let flatKeys = rawKeys.reshaped([nKVHeads * tokenCount, dim]) + let flatVals = rawValues.reshaped([nKVHeads * tokenCount, dim]) + + let (keyPacked, keyNorms) = TurboQuantKernelOps.fusedEncodeWHT( + input: flatKeys, whtSigns: keyCodec.whtSigns!, + boundaries: keyCodec.boundaries, codebook: keyCodec.codebook, + bits: keyBits, dim: dim) + let (valPacked, valNorms) = TurboQuantKernelOps.fusedEncodeWHT( + input: flatVals, whtSigns: valCodec.whtSigns!, + boundaries: valCodec.boundaries, codebook: valCodec.codebook, + bits: valueBits, dim: dim) + + let kpw = TurboQuantPacking.packedWidth(count: dim, bits: keyBits) + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let flatKeyPacked = keyPacked.reshaped([nKVHeads, tokenCount, kpw]) + let flatKeyNorms = keyNorms.reshaped([nKVHeads, tokenCount]) + let flatValPacked = valPacked.reshaped([nKVHeads, tokenCount, vpw]) + let flatValNorms = valNorms.reshaped([nKVHeads, tokenCount]) + + // Generate random queries (pre-rotated and scaled) + let scale: Float = 1.0 / sqrt(Float(dim)) + let queries = MLXRandom.normal([nQHeads, dim]) * MLXArray(scale) + eval(queries) + + // === Separated path: Score → Softmax → Value === + let scores = TurboQuantKernelOps.mseScore( + rotatedQueries: queries, packed: flatKeyPacked, norms: flatKeyNorms, + codebook: keyCodec.codebook, tokenCount: tokenCount, + repeatCount: repeatCount, bits: keyBits, dim: dim) + let attnWeights = softmax(scores, axis: -1) + let separatedOutput = TurboQuantKernelOps.mseWeightedSum( + weights: attnWeights, packed: flatValPacked, norms: flatValNorms, + codebook: valCodec.codebook, tokenCount: tokenCount, + repeatCount: repeatCount, bits: valueBits, dim: dim) + eval(separatedOutput) + + // === Fused path: TurboFlashAttention === + let fusedOutput = TurboQuantKernelOps.turboFlashAttention( + rotatedQueries: queries, + keyPacked: flatKeyPacked, keyNorms: flatKeyNorms, + keyCodebook: keyCodec.codebook, + valPacked: flatValPacked, valNorms: flatValNorms, + valCodebook: valCodec.codebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim) + eval(fusedOutput) + + // Compare outputs, should match within floating point tolerance + // Online softmax may have slightly different numerical behavior than + // materialized softmax, so allow a small tolerance + let diff = abs(separatedOutput - fusedOutput) + let maxDiff = diff.max().item(Float.self) + let meanDiff = mean(diff).item(Float.self) + + print("[TurboFlash] Max diff: \(maxDiff), Mean diff: \(meanDiff)") + print( + "[TurboFlash] Separated output range: [\(separatedOutput.min().item(Float.self)), \(separatedOutput.max().item(Float.self))]" + ) + print( + "[TurboFlash] Fused output range: [\(fusedOutput.min().item(Float.self)), \(fusedOutput.max().item(Float.self))]" + ) + + #expect(maxDiff < 1e-3, "Max diff \(maxDiff) exceeds tolerance 1e-3") + #expect(meanDiff < 1e-4, "Mean diff \(meanDiff) exceeds tolerance 1e-4") + } + + /// Test with asymmetric K/V bits (4-bit K, 2-bit V) + @Test func flashAsymmetricBits() { + let dim = 128 + let keyBits = 4 + let valueBits = 2 + let nQHeads = 4 + let nKVHeads = 2 + let tokenCount = 32 + let repeatCount = nQHeads / nKVHeads + + let keyCodec = MSECodec(dim: dim, bits: keyBits, seed: 42) + let valCodec = MSECodec(dim: dim, bits: valueBits, seed: 43) + + let rawKeys = MLXRandom.normal([nKVHeads, tokenCount, dim]) + let rawValues = MLXRandom.normal([nKVHeads, tokenCount, dim]) + eval(rawKeys, rawValues) + + let flatKeys = rawKeys.reshaped([nKVHeads * tokenCount, dim]) + let flatVals = rawValues.reshaped([nKVHeads * tokenCount, dim]) + + let (keyPacked, keyNorms) = TurboQuantKernelOps.fusedEncodeWHT( + input: flatKeys, whtSigns: keyCodec.whtSigns!, + boundaries: keyCodec.boundaries, codebook: keyCodec.codebook, + bits: keyBits, dim: dim) + let (valPacked, valNorms) = TurboQuantKernelOps.fusedEncodeWHT( + input: flatVals, whtSigns: valCodec.whtSigns!, + boundaries: valCodec.boundaries, codebook: valCodec.codebook, + bits: valueBits, dim: dim) + + let kpw = TurboQuantPacking.packedWidth(count: dim, bits: keyBits) + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let flatKeyPacked = keyPacked.reshaped([nKVHeads, tokenCount, kpw]) + let flatKeyNorms = keyNorms.reshaped([nKVHeads, tokenCount]) + let flatValPacked = valPacked.reshaped([nKVHeads, tokenCount, vpw]) + let flatValNorms = valNorms.reshaped([nKVHeads, tokenCount]) + + let scale: Float = 1.0 / sqrt(Float(dim)) + let queries = MLXRandom.normal([nQHeads, dim]) * MLXArray(scale) + eval(queries) + + // Separated + let scores = TurboQuantKernelOps.mseScore( + rotatedQueries: queries, packed: flatKeyPacked, norms: flatKeyNorms, + codebook: keyCodec.codebook, tokenCount: tokenCount, + repeatCount: repeatCount, bits: keyBits, dim: dim) + let attnWeights = softmax(scores, axis: -1) + let separatedOutput = TurboQuantKernelOps.mseWeightedSum( + weights: attnWeights, packed: flatValPacked, norms: flatValNorms, + codebook: valCodec.codebook, tokenCount: tokenCount, + repeatCount: repeatCount, bits: valueBits, dim: dim) + + // Fused + let fusedOutput = TurboQuantKernelOps.turboFlashAttention( + rotatedQueries: queries, + keyPacked: flatKeyPacked, keyNorms: flatKeyNorms, + keyCodebook: keyCodec.codebook, + valPacked: flatValPacked, valNorms: flatValNorms, + valCodebook: valCodec.codebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim) + + eval(separatedOutput, fusedOutput) + + let maxDiff = abs(separatedOutput - fusedOutput).max().item(Float.self) + print("[TurboFlash Asymmetric 4K/2V] Max diff: \(maxDiff)") + #expect(maxDiff < 1e-3, "Max diff \(maxDiff) exceeds tolerance for asymmetric bits") + } + + /// Microbenchmark: fused vs separated at various token counts + @Test func microbenchFlashVsSeparated() { + let dim = 128 + let keyBits = 4 + let valueBits = 2 + let nQHeads = 24 // Qwen3.5-2B query heads + let nKVHeads = 4 // Qwen3.5-2B KV heads + let repeatCount = nQHeads / nKVHeads + let iterations = 200 + let warmup = 50 + + let keyCodec = MSECodec(dim: dim, bits: keyBits, seed: 42) + let valCodec = MSECodec(dim: dim, bits: valueBits, seed: 43) + + for tokenCount in [128, 512, 1024, 2048, 4096, 8192] { + let rawKeys = MLXRandom.normal([nKVHeads, tokenCount, dim]) + let rawValues = MLXRandom.normal([nKVHeads, tokenCount, dim]) + eval(rawKeys, rawValues) + + let flatKeys = rawKeys.reshaped([nKVHeads * tokenCount, dim]) + let flatVals = rawValues.reshaped([nKVHeads * tokenCount, dim]) + + let (keyPacked, keyNorms) = TurboQuantKernelOps.fusedEncodeWHT( + input: flatKeys, whtSigns: keyCodec.whtSigns!, + boundaries: keyCodec.boundaries, codebook: keyCodec.codebook, + bits: keyBits, dim: dim) + let (valPacked, valNorms) = TurboQuantKernelOps.fusedEncodeWHT( + input: flatVals, whtSigns: valCodec.whtSigns!, + boundaries: valCodec.boundaries, codebook: valCodec.codebook, + bits: valueBits, dim: dim) + + let kpw = TurboQuantPacking.packedWidth(count: dim, bits: keyBits) + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let flatKeyPacked = keyPacked.reshaped([nKVHeads, tokenCount, kpw]) + let flatKeyNorms = keyNorms.reshaped([nKVHeads, tokenCount]) + let flatValPacked = valPacked.reshaped([nKVHeads, tokenCount, vpw]) + let flatValNorms = valNorms.reshaped([nKVHeads, tokenCount]) + + let scale: Float = 1.0 / sqrt(Float(dim)) + let queries = MLXRandom.normal([nQHeads, dim]) * MLXArray(scale) + eval(queries) + + // Warmup both paths + for _ in 0 ..< warmup { + let s = TurboQuantKernelOps.mseScore( + rotatedQueries: queries, packed: flatKeyPacked, norms: flatKeyNorms, + codebook: keyCodec.codebook, tokenCount: tokenCount, + repeatCount: repeatCount, bits: keyBits, dim: dim) + let w = softmax(s, axis: -1) + let v = TurboQuantKernelOps.mseWeightedSum( + weights: w, packed: flatValPacked, norms: flatValNorms, + codebook: valCodec.codebook, tokenCount: tokenCount, + repeatCount: repeatCount, bits: valueBits, dim: dim) + eval(v) + + let f = TurboQuantKernelOps.turboFlashAttention( + rotatedQueries: queries, + keyPacked: flatKeyPacked, keyNorms: flatKeyNorms, + keyCodebook: keyCodec.codebook, + valPacked: flatValPacked, valNorms: flatValNorms, + valCodebook: valCodec.codebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim) + eval(f) + } + + // Benchmark separated + let startSep = Date() + for _ in 0 ..< iterations { + let s = TurboQuantKernelOps.mseScore( + rotatedQueries: queries, packed: flatKeyPacked, norms: flatKeyNorms, + codebook: keyCodec.codebook, tokenCount: tokenCount, + repeatCount: repeatCount, bits: keyBits, dim: dim) + let w = softmax(s, axis: -1) + let v = TurboQuantKernelOps.mseWeightedSum( + weights: w, packed: flatValPacked, norms: flatValNorms, + codebook: valCodec.codebook, tokenCount: tokenCount, + repeatCount: repeatCount, bits: valueBits, dim: dim) + eval(v) + } + let sepMs = Date().timeIntervalSince(startSep) * 1000 / Double(iterations) + + // Benchmark fused (default block size) + let startFused = Date() + for _ in 0 ..< iterations { + let f = TurboQuantKernelOps.turboFlashAttention( + rotatedQueries: queries, + keyPacked: flatKeyPacked, keyNorms: flatKeyNorms, + keyCodebook: keyCodec.codebook, + valPacked: flatValPacked, valNorms: flatValNorms, + valCodebook: valCodec.codebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim) + eval(f) + } + let fusedMs = Date().timeIntervalSince(startFused) * 1000 / Double(iterations) + + let speedup = sepMs / fusedMs + print( + "[MICROBENCH] T=\(tokenCount): separated=\(String(format: "%.3f", sepMs))ms, fused(B=\(TurboQuantKernelOps.flashBlockSize))=\(String(format: "%.3f", fusedMs))ms, speedup=\(String(format: "%.2f", speedup))x" + ) + } + } + + /// Block size sweep: find optimal block size for each token count + @Test func microbenchBlockSizeSweep() { + let dim = 128 + let keyBits = 4 + let valueBits = 2 + let nQHeads = 24 + let nKVHeads = 4 + let repeatCount = nQHeads / nKVHeads + let iterations = 200 + let warmup = 30 + + let keyCodec = MSECodec(dim: dim, bits: keyBits, seed: 42) + let valCodec = MSECodec(dim: dim, bits: valueBits, seed: 43) + + let blockSizes = [32, 64, 128, 256, 512, 1024] + let tokenCounts = [512, 1024, 2048, 4096, 8192] + + for tokenCount in tokenCounts { + let rawKeys = MLXRandom.normal([nKVHeads, tokenCount, dim]) + let rawValues = MLXRandom.normal([nKVHeads, tokenCount, dim]) + eval(rawKeys, rawValues) + + let flatKeys = rawKeys.reshaped([nKVHeads * tokenCount, dim]) + let flatVals = rawValues.reshaped([nKVHeads * tokenCount, dim]) + + let (keyPacked, keyNorms) = TurboQuantKernelOps.fusedEncodeWHT( + input: flatKeys, whtSigns: keyCodec.whtSigns!, + boundaries: keyCodec.boundaries, codebook: keyCodec.codebook, + bits: keyBits, dim: dim) + let (valPacked, valNorms) = TurboQuantKernelOps.fusedEncodeWHT( + input: flatVals, whtSigns: valCodec.whtSigns!, + boundaries: valCodec.boundaries, codebook: valCodec.codebook, + bits: valueBits, dim: dim) + + let kpw = TurboQuantPacking.packedWidth(count: dim, bits: keyBits) + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let flatKeyPacked = keyPacked.reshaped([nKVHeads, tokenCount, kpw]) + let flatKeyNorms = keyNorms.reshaped([nKVHeads, tokenCount]) + let flatValPacked = valPacked.reshaped([nKVHeads, tokenCount, vpw]) + let flatValNorms = valNorms.reshaped([nKVHeads, tokenCount]) + + let scale: Float = 1.0 / sqrt(Float(dim)) + let queries = MLXRandom.normal([nQHeads, dim]) * MLXArray(scale) + eval(queries) + + // Separated baseline + for _ in 0 ..< warmup { + let s = TurboQuantKernelOps.mseScore( + rotatedQueries: queries, packed: flatKeyPacked, norms: flatKeyNorms, + codebook: keyCodec.codebook, tokenCount: tokenCount, + repeatCount: repeatCount, bits: keyBits, dim: dim) + let w = softmax(s, axis: -1) + let v = TurboQuantKernelOps.mseWeightedSum( + weights: w, packed: flatValPacked, norms: flatValNorms, + codebook: valCodec.codebook, tokenCount: tokenCount, + repeatCount: repeatCount, bits: valueBits, dim: dim) + eval(v) + } + let startSep = Date() + for _ in 0 ..< iterations { + let s = TurboQuantKernelOps.mseScore( + rotatedQueries: queries, packed: flatKeyPacked, norms: flatKeyNorms, + codebook: keyCodec.codebook, tokenCount: tokenCount, + repeatCount: repeatCount, bits: keyBits, dim: dim) + let w = softmax(s, axis: -1) + let v = TurboQuantKernelOps.mseWeightedSum( + weights: w, packed: flatValPacked, norms: flatValNorms, + codebook: valCodec.codebook, tokenCount: tokenCount, + repeatCount: repeatCount, bits: valueBits, dim: dim) + eval(v) + } + let sepMs = Date().timeIntervalSince(startSep) * 1000 / Double(iterations) + + var results: [(Int, Double)] = [] + for bs in blockSizes where bs <= tokenCount { + // Warmup + for _ in 0 ..< warmup { + let f = TurboQuantKernelOps.turboFlashAttention( + rotatedQueries: queries, + keyPacked: flatKeyPacked, keyNorms: flatKeyNorms, + keyCodebook: keyCodec.codebook, + valPacked: flatValPacked, valNorms: flatValNorms, + valCodebook: valCodec.codebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim, + blockSize: bs) + eval(f) + } + + let start = Date() + for _ in 0 ..< iterations { + let f = TurboQuantKernelOps.turboFlashAttention( + rotatedQueries: queries, + keyPacked: flatKeyPacked, keyNorms: flatKeyNorms, + keyCodebook: keyCodec.codebook, + valPacked: flatValPacked, valNorms: flatValNorms, + valCodebook: valCodec.codebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim, + blockSize: bs) + eval(f) + } + let ms = Date().timeIntervalSince(start) * 1000 / Double(iterations) + results.append((bs, ms)) + } + + let best = results.min(by: { $0.1 < $1.1 })! + let resultStr = results.map { "B=\($0.0):\(String(format: "%.2f", $0.1))ms" }.joined( + separator: " ") + print( + "[SWEEP] T=\(tokenCount): sep=\(String(format: "%.2f", sepMs))ms \(resultStr) BEST=B\(best.0)(\(String(format: "%.1f", sepMs/best.1))x)" + ) + } + } + + /// Validate that causal TurboFlashAttention matches per-position reference. + /// Computes reference by running non-causal flash on truncated KV for each query position. + @Test func flashCausalMatchesSeparated() { + let dim = 128 + let keyBits = 4 + let valueBits = 2 + let nQHeads = 8 + let nKVHeads = 4 + let L = 4 // query chunk length + let tokenCount = 32 // total KV cache length + let repeatCount = nQHeads / nKVHeads + let queryOffset = tokenCount - L // queries cover positions 28..31 + + let keyCodec = MSECodec(dim: dim, bits: keyBits, seed: 42) + let valCodec = MSECodec(dim: dim, bits: valueBits, seed: 43) + + let rawKeys = MLXRandom.normal([nKVHeads, tokenCount, dim]) + let rawValues = MLXRandom.normal([nKVHeads, tokenCount, dim]) + eval(rawKeys, rawValues) + + let flatKeys = rawKeys.reshaped([nKVHeads * tokenCount, dim]) + let flatVals = rawValues.reshaped([nKVHeads * tokenCount, dim]) + + let (keyPacked, keyNorms) = TurboQuantKernelOps.fusedEncodeWHT( + input: flatKeys, whtSigns: keyCodec.whtSigns!, + boundaries: keyCodec.boundaries, codebook: keyCodec.codebook, + bits: keyBits, dim: dim) + let (valPacked, valNorms) = TurboQuantKernelOps.fusedEncodeWHT( + input: flatVals, whtSigns: valCodec.whtSigns!, + boundaries: valCodec.boundaries, codebook: valCodec.codebook, + bits: valueBits, dim: dim) + + let kpw = TurboQuantPacking.packedWidth(count: dim, bits: keyBits) + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let flatKeyPacked = keyPacked.reshaped([nKVHeads, tokenCount, kpw]) + let flatKeyNorms = keyNorms.reshaped([nKVHeads, tokenCount]) + let flatValPacked = valPacked.reshaped([nKVHeads, tokenCount, vpw]) + let flatValNorms = valNorms.reshaped([nKVHeads, tokenCount]) + + let scale: Float = 1.0 / sqrt(Float(dim)) + // Queries in [nQHeads, L, dim] layout, flattened to [nQHeads * L, dim] + let queries = MLXRandom.normal([nQHeads, L, dim]) * MLXArray(scale) + let flatQueries = queries.reshaped([nQHeads * L, dim]) + eval(flatQueries) + + // === Reference: compute per-position using non-causal flash on truncated KV === + // For each query position l, attend only to tokens 0...(queryOffset + l) + var refOutputs: [MLXArray] = [] + for l in 0 ..< L { + let visibleTokens = queryOffset + l + 1 // causal: can see up to and including position + let truncKeyPacked = flatKeyPacked[0..., .. 0) + } + + /// Microbenchmark: WHT rotation fused encode kernel. + @Test func microbenchWHTEncode() { + let dim = 128 + let bits = 4 + let nKVHeads = 4 + let iterations = 500 + let warmup = 50 + + let codec = MSECodec(dim: dim, bits: bits, seed: 99) + guard codec.useWHT, let signs = codec.whtSigns else { + Issue.record("dim=128 should use WHT") + return + } + + let input = MLXRandom.normal([nKVHeads, dim]) + eval(input) + + // Warmup + for _ in 0 ..< warmup { + let (p, n) = TurboQuantKernelOps.fusedEncodeWHT( + input: input, whtSigns: signs, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + + // Timed iterations + let start = Date() + for _ in 0 ..< iterations { + let (p, n) = TurboQuantKernelOps.fusedEncodeWHT( + input: input, whtSigns: signs, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + let elapsed = Date().timeIntervalSince(start) * 1000 + let perCall = elapsed / Double(iterations) + + print( + "[MICROBENCH] WHT encode: \(iterations) iters, \(String(format: "%.1f", elapsed))ms total, \(String(format: "%.3f", perCall))ms/call" + ) + #expect(perCall > 0) + } + + /// Microbenchmark: Batch encode (prefill transition), 512 tokens at once. + @Test func microbenchBatchEncodeDense() { + let dim = 128 + let bits = 4 + let batchSize = 512 * 4 // 512 tokens × 4 KV heads + let iterations = 100 + let warmup = 10 + + let rotation = TurboQuantRotation.rotationMatrix(dim: dim, seed: 99) + let codec = MSECodec(dim: dim, bits: bits, seed: 99) + let input = MLXRandom.normal([batchSize, dim]) + eval(input) + + for _ in 0 ..< warmup { + let (p, n) = TurboQuantKernelOps.fusedEncode( + input: input, rotation: rotation, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + + let start = Date() + for _ in 0 ..< iterations { + let (p, n) = TurboQuantKernelOps.fusedEncode( + input: input, rotation: rotation, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + let elapsed = Date().timeIntervalSince(start) * 1000 + let perCall = elapsed / Double(iterations) + + print( + "[MICROBENCH] Dense batch encode (512×4): \(iterations) iters, \(String(format: "%.1f", elapsed))ms total, \(String(format: "%.3f", perCall))ms/call" + ) + #expect(perCall > 0) + } + + /// Microbenchmark: Batch encode WHT, 512 tokens at once. + @Test func microbenchBatchEncodeWHT() { + let dim = 128 + let bits = 4 + let batchSize = 512 * 4 + let iterations = 100 + let warmup = 10 + + let codec = MSECodec(dim: dim, bits: bits, seed: 99) + guard codec.useWHT, let signs = codec.whtSigns else { + Issue.record("dim=128 should use WHT") + return + } + let input = MLXRandom.normal([batchSize, dim]) + eval(input) + + for _ in 0 ..< warmup { + let (p, n) = TurboQuantKernelOps.fusedEncodeWHT( + input: input, whtSigns: signs, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + + let start = Date() + for _ in 0 ..< iterations { + let (p, n) = TurboQuantKernelOps.fusedEncodeWHT( + input: input, whtSigns: signs, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + let elapsed = Date().timeIntervalSince(start) * 1000 + let perCall = elapsed / Double(iterations) + + print( + "[MICROBENCH] WHT batch encode dim=128 (512×4): \(iterations) iters, \(String(format: "%.1f", elapsed))ms total, \(String(format: "%.3f", perCall))ms/call" + ) + #expect(perCall > 0) + } + + // --- dim=256 variants (Qwen3.5-27B full attention head_dim) --- + + @Test func microbenchDenseEncode256() { + let dim = 256 + let bits = 4 + let nKVHeads = 4 + let iterations = 500 + let warmup = 50 + + let rotation = TurboQuantRotation.rotationMatrix(dim: dim, seed: 99) + let codec = MSECodec(dim: dim, bits: bits, seed: 99) + let input = MLXRandom.normal([nKVHeads, dim]) + eval(input) + + for _ in 0 ..< warmup { + let (p, n) = TurboQuantKernelOps.fusedEncode( + input: input, rotation: rotation, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + + let start = Date() + for _ in 0 ..< iterations { + let (p, n) = TurboQuantKernelOps.fusedEncode( + input: input, rotation: rotation, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + let elapsed = Date().timeIntervalSince(start) * 1000 + let perCall = elapsed / Double(iterations) + + print( + "[MICROBENCH] Dense encode dim=256: \(iterations) iters, \(String(format: "%.1f", elapsed))ms total, \(String(format: "%.3f", perCall))ms/call" + ) + #expect(perCall > 0) + } + + @Test func microbenchWHTEncode256() { + let dim = 256 + let bits = 4 + let nKVHeads = 4 + let iterations = 500 + let warmup = 50 + + let codec = MSECodec(dim: dim, bits: bits, seed: 99) + guard codec.useWHT, let signs = codec.whtSigns else { + Issue.record("dim=256 should use WHT") + return + } + let input = MLXRandom.normal([nKVHeads, dim]) + eval(input) + + for _ in 0 ..< warmup { + let (p, n) = TurboQuantKernelOps.fusedEncodeWHT( + input: input, whtSigns: signs, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + + let start = Date() + for _ in 0 ..< iterations { + let (p, n) = TurboQuantKernelOps.fusedEncodeWHT( + input: input, whtSigns: signs, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + let elapsed = Date().timeIntervalSince(start) * 1000 + let perCall = elapsed / Double(iterations) + + print( + "[MICROBENCH] WHT encode dim=256: \(iterations) iters, \(String(format: "%.1f", elapsed))ms total, \(String(format: "%.3f", perCall))ms/call" + ) + #expect(perCall > 0) + } + + @Test func microbenchBatchEncodeDense256() { + let dim = 256 + let bits = 4 + let batchSize = 512 * 4 + let iterations = 100 + let warmup = 10 + + let rotation = TurboQuantRotation.rotationMatrix(dim: dim, seed: 99) + let codec = MSECodec(dim: dim, bits: bits, seed: 99) + let input = MLXRandom.normal([batchSize, dim]) + eval(input) + + for _ in 0 ..< warmup { + let (p, n) = TurboQuantKernelOps.fusedEncode( + input: input, rotation: rotation, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + + let start = Date() + for _ in 0 ..< iterations { + let (p, n) = TurboQuantKernelOps.fusedEncode( + input: input, rotation: rotation, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + let elapsed = Date().timeIntervalSince(start) * 1000 + let perCall = elapsed / Double(iterations) + + print( + "[MICROBENCH] Dense batch encode dim=256 (512×4): \(iterations) iters, \(String(format: "%.1f", elapsed))ms total, \(String(format: "%.3f", perCall))ms/call" + ) + #expect(perCall > 0) + } + + @Test func microbenchBatchEncodeWHT256() { + let dim = 256 + let bits = 4 + let batchSize = 512 * 4 + let iterations = 100 + let warmup = 10 + + let codec = MSECodec(dim: dim, bits: bits, seed: 99) + guard codec.useWHT, let signs = codec.whtSigns else { + Issue.record("dim=256 should use WHT") + return + } + let input = MLXRandom.normal([batchSize, dim]) + eval(input) + + for _ in 0 ..< warmup { + let (p, n) = TurboQuantKernelOps.fusedEncodeWHT( + input: input, whtSigns: signs, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + + let start = Date() + for _ in 0 ..< iterations { + let (p, n) = TurboQuantKernelOps.fusedEncodeWHT( + input: input, whtSigns: signs, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(p, n) + } + let elapsed = Date().timeIntervalSince(start) * 1000 + let perCall = elapsed / Double(iterations) + + print( + "[MICROBENCH] WHT batch encode dim=256 (512×4): \(iterations) iters, \(String(format: "%.1f", elapsed))ms total, \(String(format: "%.3f", perCall))ms/call" + ) + #expect(perCall > 0) + } +} + +// MARK: - kvScheme routing + integration (XCTest) + +final class TurboQuantIntegrationTests: XCTestCase { + + func testResolveTurboScheme() { + XCTAssertNil(resolveTurboScheme(nil)) + XCTAssertNil(resolveTurboScheme("affine4")) + XCTAssertNil(resolveTurboScheme("nonsense")) + + // Raw-key schemes are the validated surface; keyBits 0 = FP16 keys. + XCTAssertEqual(resolveTurboScheme("turbo0v4")?.keyBits, 0) + XCTAssertEqual(resolveTurboScheme("turbo0v4")?.valueBits, 4) + XCTAssertEqual(resolveTurboScheme("turbo0v2")?.valueBits, 2) + // Key-quantizing schemes are exposed on the canonical codec. + XCTAssertEqual(resolveTurboScheme("turbo4")?.keyBits, 4) + XCTAssertEqual(resolveTurboScheme("turbo4v2")?.valueBits, 2) + // Affine-K asymmetric ladder: keyBits 8 = 8-bit affine keys. + XCTAssertEqual(resolveTurboScheme("turbo8v4")?.keyBits, 8) + XCTAssertEqual(resolveTurboScheme("turbo8v3")?.valueBits, 3) + XCTAssertEqual(resolveTurboScheme("turbo8v2")?.valueBits, 2) + XCTAssertEqual(resolveTurboScheme("turbo0v3")?.valueBits, 3) + XCTAssertNil(resolveTurboScheme("turbo8")) + } + + func testMaybeTurboQuantizeConvertsEligibleLayers() { + let (b, h, d, t) = (1, 2, 64, 8) + let simple = KVCacheSimple() + let keys = MLXRandom.normal([b, h, t, d], key: MLXRandom.key(9)) + let values = MLXRandom.normal([b, h, t, d], key: MLXRandom.key(10)) + _ = simple.update(keys: keys, values: values) + + var caches: [KVCache] = [simple] + maybeTurboQuantizeKVCache(cache: &caches, keyBits: 4, valueBits: 4, quantizedKVStart: 0) + + XCTAssertTrue(caches[0] is TurboQuantKVCache, "eligible layer not converted") + XCTAssertEqual(caches[0].offset, t, "offset lost in conversion") + + // Below the start threshold: untouched. + let early = KVCacheSimple() + _ = early.update( + keys: MLXRandom.normal([b, h, 2, d], key: MLXRandom.key(11)), + values: MLXRandom.normal([b, h, 2, d], key: MLXRandom.key(12))) + var earlyCaches: [KVCache] = [early] + maybeTurboQuantizeKVCache( + cache: &earlyCaches, keyBits: 4, valueBits: 4, quantizedKVStart: 100) + XCTAssertTrue(earlyCaches[0] is KVCacheSimple, "layer below threshold was converted") + } + + // MARK: - Mixed rotating/standard cache lists (Gemma-style sliding window) + + func testMixedRotatingAndStandardCacheLeavesRotatingLayersUntouched() { + // Gemma-family layouts interleave RotatingKVCache (sliding window) + // with KVCacheSimple (global) layers. TurboQuant should only ever + // touch the KVCacheSimple layers: RotatingKVCache's eviction/rotation + // storage layout is not compatible with the sequential-append + // compression path, and its window is already memory-bounded. + let (b, h, d, t) = (1, 2, 64, 8) + + func filledStandard() -> KVCacheSimple { + let c = KVCacheSimple() + _ = c.update( + keys: MLXRandom.normal([b, h, t, d], key: MLXRandom.key(61)), + values: MLXRandom.normal([b, h, t, d], key: MLXRandom.key(62))) + return c + } + + func filledRotating() -> RotatingKVCache { + let c = RotatingKVCache(maxSize: 16, keep: 0) + _ = c.update( + keys: MLXRandom.normal([b, h, t, d], key: MLXRandom.key(63)), + values: MLXRandom.normal([b, h, t, d], key: MLXRandom.key(64))) + return c + } + + // turbo8v4 is a near-lossless (non-fragile) scheme, so boundary-layer + // protection does not engage and every eligible layer converts -- + // isolating the rotating-vs-standard split this test targets. + var caches: [KVCache] = [ + filledRotating(), filledStandard(), filledRotating(), filledStandard(), + ] + maybeTurboQuantizeKVCache(cache: &caches, keyBits: 8, valueBits: 4, quantizedKVStart: 0) + + XCTAssertTrue(caches[0] is RotatingKVCache, "rotating layer 0 was converted") + XCTAssertTrue(caches[1] is TurboQuantKVCache, "standard layer 1 was not converted") + XCTAssertTrue(caches[2] is RotatingKVCache, "rotating layer 2 was converted") + XCTAssertTrue(caches[3] is TurboQuantKVCache, "standard layer 3 was not converted") + XCTAssertEqual(caches[0].offset, t, "rotating cache offset disturbed by conversion pass") + XCTAssertEqual(caches[2].offset, t, "rotating cache offset disturbed by conversion pass") + } + + func testAllRotatingCacheIsNoOp() { + // Edge case: every layer is a RotatingKVCache (no global layers at + // all). There is nothing eligible to convert; the scheme should + // no-op cleanly rather than crash or silently drop state. + let (b, h, d, t) = (1, 2, 64, 8) + + func filledRotating(_ seed: UInt64) -> RotatingKVCache { + let c = RotatingKVCache(maxSize: 16, keep: 0) + _ = c.update( + keys: MLXRandom.normal([b, h, t, d], key: MLXRandom.key(seed)), + values: MLXRandom.normal([b, h, t, d], key: MLXRandom.key(seed + 1))) + return c + } + + var caches: [KVCache] = [filledRotating(71), filledRotating(73)] + maybeTurboQuantizeKVCache(cache: &caches, keyBits: 4, valueBits: 4, quantizedKVStart: 0) + + XCTAssertTrue( + caches.allSatisfy { $0 is RotatingKVCache }, "all-rotating cache list should no-op") + XCTAssertEqual(caches[0].offset, t) + XCTAssertEqual(caches[1].offset, t) + } + + func testEndToEndGenerationWithTurboScheme() throws { + // Tiny random-weight Llama: prefill, then a short greedy decode loop + // with the turbo4 scheme applied after each step, the same call + // pattern as TokenIterator.step(). Catches integration breaks + // (cache conversion, attention routing, NaN logits) without a + // checkpoint. + let config = LlamaConfiguration( + hiddenSize: 128, hiddenLayers: 6, intermediateSize: 128, attentionHeads: 8, + rmsNormEps: 1e-5, vocabularySize: 100, kvHeads: 4) + let model = LlamaModel(config) + quantize(model: model, groupSize: 64, bits: 4) + eval(model) + + func decode(scheme: String?) -> [Int] { + var cache = model.newCache(parameters: nil) + let prompt = MLXArray([1, 7, 3, 12, 5, 9, 2, 8])[.newAxis, .ellipsis] + var logits = model(prompt, cache: cache) + var tokens: [Int] = [] + for _ in 0 ..< 8 { + maybeQuantizeKVCache( + cache: &cache, kvBits: nil, quantizedKVStart: 0, kvScheme: scheme) + let next = argMax(logits[0..., -1, 0...], axis: -1) + let token = next.item(Int.self) + tokens.append(token) + logits = model(next[.newAxis, .ellipsis], cache: cache) + let bad = MLX.isNaN(logits).any().item(Bool.self) + XCTAssertFalse(bad, "NaN logits during \(scheme ?? "fp16") decode") + } + if scheme != nil { + // turbo0v4 is a non-fragile scheme: no boundary protection, + // every layer converts. + XCTAssertTrue( + cache.allSatisfy { $0 is TurboQuantKVCache }, + "caches not converted to TurboQuantKVCache") + } + return tokens + } + + let turbo = decode(scheme: "turbo0v4") + let reference = decode(scheme: nil) + XCTAssertEqual(turbo.count, 8) + // The first decode token is computed from the still-raw cache, so it + // must match the FP16 run exactly; later tokens may diverge within + // 4-bit quantization error. + XCTAssertEqual(turbo[0], reference[0], "first decode token diverged before compression") + } + + func testEndToEndGemmaLikeMixedCacheDecode() throws { + // Tiny random-weight Gemma3Text: alternating sliding-window + // (RotatingKVCache) and global (KVCacheSimple) layers via + // slidingWindowPattern: 2. Prefill, then a short greedy decode loop + // applying the turbo8v4 scheme after each step, the same call + // pattern as TokenIterator.step(). Only the global layers should + // ever convert; the sliding-window layers must keep decoding + // correctly as plain fp16 RotatingKVCache throughout. + // + // headDim 64 (not a smaller tiny value) is required here: turbo8v4's + // 8-bit affine key path uses a fixed groupSize of 64 internally, so + // headDim must be a multiple of it. + let config = Gemma3TextConfiguration( + modelType: "gemma3_text", hiddenSize: 64, hiddenLayers: 4, intermediateSize: 128, + attentionHeads: 4, headDim: 64, rmsNormEps: 1e-5, vocabularySize: 50, kvHeads: 2, + ropeTheta: 1_000_000, ropeLocalBaseFreq: 10_000, ropeTraditional: false, + queryPreAttnScalar: 64, slidingWindow: 8, slidingWindowPattern: 2, + maxPositionEmbeddings: 128) + let model = Gemma3TextModel(config) + eval(model) + + var cache = model.newCache(parameters: nil) + // isGlobalLayer = (i % slidingWindowPattern == slidingWindowPattern - 1), + // so with pattern 2: layers 0, 2 are sliding-window; 1, 3 are global. + XCTAssertTrue(cache[0] is RotatingKVCache, "layer 0 should be sliding-window") + XCTAssertTrue(cache[1] is KVCacheSimple, "layer 1 should be global") + XCTAssertTrue(cache[2] is RotatingKVCache, "layer 2 should be sliding-window") + XCTAssertTrue(cache[3] is KVCacheSimple, "layer 3 should be global") + + let prompt = MLXArray([1, 7, 3, 12, 5])[.newAxis, .ellipsis] + var logits = model(prompt, cache: cache) + eval(logits) + XCTAssertFalse(MLX.isNaN(logits).any().item(Bool.self), "NaN logits during prefill") + + for _ in 0 ..< 6 { + maybeQuantizeKVCache( + cache: &cache, kvBits: nil, quantizedKVStart: 0, kvScheme: "turbo8v4") + let next = argMax(logits[0..., -1, 0...], axis: -1) + logits = model(next[.newAxis, .ellipsis], cache: cache) + eval(logits) + XCTAssertFalse(MLX.isNaN(logits).any().item(Bool.self), "NaN logits during decode") + } + + XCTAssertTrue(cache[0] is RotatingKVCache, "sliding-window layer 0 was converted") + XCTAssertTrue(cache[1] is TurboQuantKVCache, "global layer 1 was not converted") + XCTAssertTrue(cache[2] is RotatingKVCache, "sliding-window layer 2 was converted") + XCTAssertTrue(cache[3] is TurboQuantKVCache, "global layer 3 was not converted") + } + + func testPromptCacheSaveRestoreRoundTrip() throws { + // Regression test: TurboQuant caches used to serialize into a + // prompt cache and restore as plain KVCacheSimple, silently + // dropping the compressed (packed indices + scales) state. + let (b, hq, hkv, d) = (1, 4, 2, 64) + let prefill = 16 + + let cache = TurboQuantKVCache(bits: 4, keyBits: 4, valueBits: 2, seed: 123) + let keys = MLXRandom.normal([b, hkv, prefill, d], key: MLXRandom.key(31)) * 0.3 + let values = MLXRandom.normal([b, hkv, prefill, d], key: MLXRandom.key(32)) * 0.3 + _ = cache.update(keys: keys, values: values) + + let newK = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(33)) * 0.3 + let newV = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(34)) * 0.3 + let q = MLXRandom.normal([b, hq, 1, d], key: MLXRandom.key(35)) * 0.3 + let scale = 1.0 / Float(d).squareRoot() + + // Force compression so the saved state exercises the compressed + // (packed + scales) branch, not the raw-prefill branch. + _ = cache.compressedAttention(queries: q, keys: newK, values: newV, scale: scale) + XCTAssertTrue(cache.isCompressed) + XCTAssertEqual(cache.offset, prefill + 1) + + let url = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString) + .appendingPathExtension("safetensors") + + try savePromptCache(url: url, cache: [cache], metadata: [:]) + let (loaded, _) = try loadPromptCache(url: url) + + XCTAssertEqual(loaded.count, 1) + guard let restored = loaded[0] as? TurboQuantKVCache else { + XCTFail("expected TurboQuantKVCache, got \(type(of: loaded[0]))") + return + } + + XCTAssertEqual(restored.offset, cache.offset) + XCTAssertEqual(restored.metaState, cache.metaState) + XCTAssertEqual(restored.keyBits, cache.keyBits) + XCTAssertEqual(restored.valueBits, cache.valueBits) + + let originalState = cache.state + let restoredState = restored.state + XCTAssertEqual(originalState.count, restoredState.count) + for (i, (orig, rest)) in zip(originalState, restoredState).enumerated() { + XCTAssertEqual(orig.shape, rest.shape, "state[\(i)] shape mismatch") + XCTAssertTrue(allClose(orig, rest).item(Bool.self), "state[\(i)] values diverged") + } + + // Restored cache must still be functional for further decode steps. + let nextK = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(36)) * 0.3 + let nextV = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(37)) * 0.3 + let nextQ = MLXRandom.normal([b, hq, 1, d], key: MLXRandom.key(38)) * 0.3 + let out = restored.compressedAttention( + queries: nextQ, keys: nextK, values: nextV, scale: scale) + + XCTAssertEqual(out.shape, [b, hq, 1, d]) + XCTAssertEqual(restored.offset, prefill + 2) + XCTAssertFalse(MLX.isNaN(out).any().item(Bool.self), "post-restore decode produced NaN") + } + + /// Key calibration scale must round-trip through save/restore: build + /// anisotropic keys (same 10x post-rotation spread as the codec-level + /// test) so compressRawCache actually engages calibration, then verify + /// the extra 5th state array (keyCalibScale) survives a save/load cycle. + func testKeyCalibrationPromptCacheRoundTrip() throws { + let (b, hkv, hq, d) = (1, 2, 4, 64) + let prefill = 64 + let seed: UInt64 = 42 + + // Same seed the cache will use internally for its key codec, so the + // rotation matrix here matches the one compressRawCache rotates by. + let refCodec = MSECodec(dim: d, bits: 4, seed: seed) + var stdVec = [Float](repeating: 1.0, count: d) + for i in 0 ..< d / 2 { stdVec[i] = Float(10.0).squareRoot() } + let stdArray = MLXArray(stdVec) + let rotatedSynthetic = + MLXRandom.normal([b, hkv, prefill, d], key: MLXRandom.key(71)) * stdArray + let keys = matmul(rotatedSynthetic, refCodec.rotation) + let values = MLXRandom.normal([b, hkv, prefill, d], key: MLXRandom.key(72)) * 0.3 + eval(keys, values) + + let cache = TurboQuantKVCache(bits: 4, keyBits: 4, valueBits: 4, seed: seed) + _ = cache.update(keys: keys, values: values) + + let newK = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(73)) * 0.3 + let newV = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(74)) * 0.3 + let q = MLXRandom.normal([b, hq, 1, d], key: MLXRandom.key(75)) * 0.3 + let scale = 1.0 / Float(d).squareRoot() + _ = cache.compressedAttention(queries: q, keys: newK, values: newV, scale: scale) + XCTAssertTrue(cache.isCompressed) + + let originalState = cache.state + XCTAssertEqual( + originalState.count, 5, + "anisotropic key data should engage key calibration (5-array state)") + + let url = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString) + .appendingPathExtension("safetensors") + try savePromptCache(url: url, cache: [cache], metadata: [:]) + let (loaded, _) = try loadPromptCache(url: url) + let restored = try XCTUnwrap(loaded[0] as? TurboQuantKVCache) + + let restoredState = restored.state + XCTAssertEqual(originalState.count, restoredState.count) + for (i, (orig, rest)) in zip(originalState, restoredState).enumerated() { + XCTAssertEqual(orig.shape, rest.shape, "state[\(i)] shape mismatch") + XCTAssertTrue(allClose(orig, rest).item(Bool.self), "state[\(i)] values diverged") + } + + // Restored cache must still decode correctly using the restored scale. + let nextK = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(76)) * 0.3 + let nextV = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(77)) * 0.3 + let nextQ = MLXRandom.normal([b, hq, 1, d], key: MLXRandom.key(78)) * 0.3 + let out = restored.compressedAttention( + queries: nextQ, keys: nextK, values: nextV, scale: scale) + XCTAssertEqual(out.shape, [b, hq, 1, d]) + XCTAssertFalse(MLX.isNaN(out).any().item(Bool.self), "post-restore decode produced NaN") + } + + func testBoundaryProtectionGating() { + // Fragile schemes (turbo-quantized keys) protect the first/last 2 + // attention layers; near-lossless schemes convert everything. + func layers(_ n: Int) -> [KVCache] { + (0 ..< n).map { _ in + let c = KVCacheSimple() + _ = c.update( + keys: MLXArray.ones([1, 2, 8, 64], dtype: .float16), + values: MLXArray.ones([1, 2, 8, 64], dtype: .float16)) + return c + } + } + + var fragile = layers(6) + maybeTurboQuantizeKVCache( + cache: &fragile, keyBits: 4, valueBits: 4, quantizedKVStart: 0) + // Boundary layers get near-lossless q8 affine, interior layers turbo. + XCTAssertTrue(fragile[0] is QuantizedKVCache && fragile[5] is QuantizedKVCache) + XCTAssertTrue(fragile[2] is TurboQuantKVCache && fragile[3] is TurboQuantKVCache) + + var safe = layers(6) + maybeTurboQuantizeKVCache( + cache: &safe, keyBits: 8, valueBits: 4, quantizedKVStart: 0) + XCTAssertTrue(safe.allSatisfy { $0 is TurboQuantKVCache }) + + var aggressiveV = layers(6) + maybeTurboQuantizeKVCache( + cache: &aggressiveV, keyBits: 0, valueBits: 2, quantizedKVStart: 0) + XCTAssertTrue(aggressiveV[0] is QuantizedKVCache, "2-bit V should re-engage protection") + XCTAssertTrue(aggressiveV[2] is TurboQuantKVCache) + } + + /// Head dimension 80 divides none of the supported affine-quantization + /// group sizes (32, 64, 128). Regular-path affine-K conversion + /// (keyBits == 8, non-fragile so boundary protection doesn't engage) + /// must skip such layers, leaving them fp16, rather than crashing deep + /// inside `quantized()`. + func testAffineKeyModeSkipsIncompatibleHeadDim() { + let (b, h, d, t) = (1, 2, 80, 8) + let simple = KVCacheSimple() + _ = simple.update( + keys: MLXRandom.normal([b, h, t, d], key: MLXRandom.key(81)), + values: MLXRandom.normal([b, h, t, d], key: MLXRandom.key(82))) + + var caches: [KVCache] = [simple] + // turbo8v4 is non-fragile (keyBits == 8, valueBits > 2), so this + // exercises the regular affine-K path, not boundary protection. + maybeTurboQuantizeKVCache(cache: &caches, keyBits: 8, valueBits: 4, quantizedKVStart: 0) + + XCTAssertTrue( + caches[0] is KVCacheSimple, + "head dim 80 is incompatible with every affine group size, should stay fp16") + XCTAssertEqual(caches[0].offset, t, "offset lost while skipping conversion") + } + + /// Same head-dim-80 incompatibility, but on a fragile scheme so the + /// first/last layers route through boundary protection (always 8-bit + /// affine for both K and V, regardless of the requested scheme). + func testBoundaryProtectionSkipsIncompatibleHeadDim() { + let (b, h, d, t) = (1, 2, 80, 8) + func filled(_ seed: UInt64) -> KVCacheSimple { + let c = KVCacheSimple() + _ = c.update( + keys: MLXRandom.normal([b, h, t, d], key: MLXRandom.key(seed)), + values: MLXRandom.normal([b, h, t, d], key: MLXRandom.key(seed + 1))) + return c + } + + var caches: [KVCache] = [filled(83), filled(85), filled(87), filled(89)] + // turbo2 is fragile (valueBits <= 2): every layer in this 4-layer + // list lands in the protected boundary set. + maybeTurboQuantizeKVCache(cache: &caches, keyBits: 2, valueBits: 2, quantizedKVStart: 0) + + XCTAssertTrue( + caches.allSatisfy { $0 is KVCacheSimple }, + "head dim 80 is incompatible with every affine group size, boundary layers should stay fp16" + ) + } + + // MARK: - Affine-K asymmetric mode (turbo8vN) + + func testAffineKeyModeMatchesReference() throws { + // q8-affine keys + turbo4 values vs exact attention over raw K/V. + let (b, hq, hkv, t, d) = (1, 4, 2, 256, 128) + let scale = 1.0 / Float(d).squareRoot() + let keys = MLXRandom.normal([b, hkv, t, d], key: MLXRandom.key(51)) * 0.5 + let values = MLXRandom.normal([b, hkv, t, d], key: MLXRandom.key(52)) * 0.5 + let q = MLXRandom.normal([b, hq, 1, d], key: MLXRandom.key(53)) * 0.5 + let newK = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(54)) * 0.5 + let newV = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(55)) * 0.5 + + let cache = TurboQuantKVCache(bits: 4, keyBits: 8, valueBits: 4, seed: 42) + XCTAssertTrue(cache.affineKeyMode) + _ = cache.update(keys: keys, values: values) + let out = cache.compressedAttention( + queries: q, keys: newK, values: newV, scale: scale) + + let nRep = hq / hkv + let fullK = concatenated([keys, newK], axis: 2) + let fullV = concatenated([values, newV], axis: 2) + let ek = MLX.tiled(expandedDimensions(fullK, axis: 2), repetitions: [1, 1, nRep, 1, 1]) + .reshaped([b, hq, t + 1, d]) + let ev = MLX.tiled(expandedDimensions(fullV, axis: 2), repetitions: [1, 1, nRep, 1, 1]) + .reshaped([b, hq, t + 1, d]) + let ref = matmul( + softmax(matmul(q, ek.transposed(0, 1, 3, 2)) * scale, axis: -1), ev) + + let a = out.asType(.float32).reshaped([-1]) + let r = ref.asType(.float32).reshaped([-1]) + let cos = ((a * r).sum() / (sqrt((a * a).sum()) * sqrt((r * r).sum()) + 1e-9)) + .item(Float.self) + // 4-bit V + q8 K at T=256 lands at cos ~0.976-0.980; the exact value + // shifts a few 1e-4 across mlx-swift patch versions. 0.97 catches real + // breakage (bugs push this below 0.9) without pinning numerics. + XCTAssertGreaterThan(cos, 0.97, "affine-K attention diverged from reference (cos \(cos))") + XCTAssertFalse(MLX.isNaN(out).any().item(Bool.self)) + } + + /// rawKeyMode (keyBits: 0) + turbo4 values, bfloat16 K cache: the raw-K + /// flash kernel reads K in its native dtype (KT template) instead of + /// casting the whole cache to float16 every decode step. One cached key + /// row is pushed above float16 max (65504) to prove that no longer + /// turns finite bfloat16 activations into inf. + func testRawKeyModeBFloat16MatchesReference() throws { + let (b, hq, hkv, t, d) = (1, 4, 2, 256, 128) + let scale = 1.0 / Float(d).squareRoot() + let keys = (MLXRandom.normal([b, hkv, t, d], key: MLXRandom.key(71)) * 0.5) + .asType(.bfloat16) + let values = (MLXRandom.normal([b, hkv, t, d], key: MLXRandom.key(72)) * 0.5) + .asType(.bfloat16) + let q = MLXRandom.normal([b, hq, 1, d], key: MLXRandom.key(73)) * 0.5 + let newK = (MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(74)) * 0.5) + .asType(.bfloat16) + let newV = (MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(75)) * 0.5) + .asType(.bfloat16) + + // Beyond float16's max finite value: forcing this cache through + // asType(.float16) would turn it into inf and NaN the softmax. + keys[0..., 0..., 0, 0] = MLXArray(Float(70000)).asType(.bfloat16) + eval(keys) + + let cache = TurboQuantKVCache(bits: 4, keyBits: 0, valueBits: 4, seed: 42) + XCTAssertTrue(cache.rawKeyMode) + _ = cache.update(keys: keys, values: values) + let out = cache.compressedAttention( + queries: q, keys: newK, values: newV, scale: scale) + + XCTAssertFalse(MLX.isNaN(out).any().item(Bool.self)) + XCTAssertFalse(MLX.isInf(out).any().item(Bool.self)) + + let nRep = hq / hkv + let fullK = concatenated([keys, newK], axis: 2).asType(.float32) + let fullV = concatenated([values, newV], axis: 2).asType(.float32) + let ek = MLX.tiled(expandedDimensions(fullK, axis: 2), repetitions: [1, 1, nRep, 1, 1]) + .reshaped([b, hq, t + 1, d]) + let ev = MLX.tiled(expandedDimensions(fullV, axis: 2), repetitions: [1, 1, nRep, 1, 1]) + .reshaped([b, hq, t + 1, d]) + let ref = matmul( + softmax(matmul(q.asType(.float32), ek.transposed(0, 1, 3, 2)) * scale, axis: -1), ev) + + let a = out.asType(.float32).reshaped([-1]) + let r = ref.asType(.float32).reshaped([-1]) + let cos = ((a * r).sum() / (sqrt((a * a).sum()) * sqrt((r * r).sum()) + 1e-9)) + .item(Float.self) + // Same cos-similarity style/threshold as testAffineKeyModeMatchesReference: + // 4-bit V at T=256 plus bfloat16 K rounding lands well above 0.97; + // a real regression (e.g. inf from an f16 cast) collapses this. + XCTAssertGreaterThan( + cos, 0.97, "raw-K bfloat16 attention diverged from reference (cos \(cos))") + } + + func testAffineKeyModePromptCacheRoundTrip() throws { + let (b, hkv, hq, d) = (1, 2, 4, 64) + let cache = TurboQuantKVCache(bits: 4, keyBits: 8, valueBits: 4, seed: 7) + let keys = MLXRandom.normal([b, hkv, 16, d], key: MLXRandom.key(61)) + let values = MLXRandom.normal([b, hkv, 16, d], key: MLXRandom.key(62)) + _ = cache.update(keys: keys, values: values) + let q = MLXRandom.normal([b, hq, 1, d], key: MLXRandom.key(63)) + let nk = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(64)) + let nv = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(65)) + _ = cache.compressedAttention( + queries: q, keys: nk, values: nv, scale: 0.125) + XCTAssertTrue(cache.isCompressed) + XCTAssertEqual(cache.state.count, 5) + + let url = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString) + .appendingPathExtension("safetensors") + try savePromptCache(url: url, cache: [cache], metadata: [:]) + let (loaded, _) = try loadPromptCache(url: url) + let restored = try XCTUnwrap(loaded[0] as? TurboQuantKVCache) + XCTAssertTrue(restored.affineKeyMode) + XCTAssertEqual(restored.offset, cache.offset) + + let out = restored.compressedAttention( + queries: q, keys: nk, values: nv, scale: 0.125) + XCTAssertEqual(restored.offset, 18) + XCTAssertFalse(MLX.isNaN(out).any().item(Bool.self)) + } + + // MARK: - Persistence for the remaining state layouts + + func testRawKeyModePromptCacheRoundTrip() throws { + // turbo0vN (the recommended family) serializes as [rawKeys, valPacked, + // valNorms], 3 arrays. + let (b, hq, hkv, d) = (1, 4, 2, 64) + let cache = TurboQuantKVCache(bits: 4, keyBits: 0, valueBits: 4, seed: 9) + _ = cache.update( + keys: MLXRandom.normal([b, hkv, 16, d], key: MLXRandom.key(71)), + values: MLXRandom.normal([b, hkv, 16, d], key: MLXRandom.key(72))) + let q = MLXRandom.normal([b, hq, 1, d], key: MLXRandom.key(73)) + let nk = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(74)) + let nv = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(75)) + _ = cache.compressedAttention(queries: q, keys: nk, values: nv, scale: 0.125) + XCTAssertTrue(cache.isCompressed) + XCTAssertEqual(cache.state.count, 3) + + let url = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString).appendingPathExtension("safetensors") + try savePromptCache(url: url, cache: [cache], metadata: [:]) + let (loaded, _) = try loadPromptCache(url: url) + let restored = try XCTUnwrap(loaded[0] as? TurboQuantKVCache) + XCTAssertTrue(restored.rawKeyMode) + XCTAssertEqual(restored.offset, cache.offset) + let out = restored.compressedAttention(queries: q, keys: nk, values: nv, scale: 0.125) + XCTAssertEqual(restored.offset, 18) + XCTAssertFalse(MLX.isNaN(out).any().item(Bool.self)) + } + + func testRawPhasePromptCacheRoundTrip() throws { + // Saving straight after prefill (before any decode) is the primary + // prompt-cache use case: state is the raw 2-array layout. + let (b, hq, hkv, d) = (1, 4, 2, 64) + let cache = TurboQuantKVCache(bits: 4, keyBits: 0, valueBits: 4, seed: 9) + _ = cache.update( + keys: MLXRandom.normal([b, hkv, 16, d], key: MLXRandom.key(81)), + values: MLXRandom.normal([b, hkv, 16, d], key: MLXRandom.key(82))) + XCTAssertFalse(cache.isCompressed) + XCTAssertEqual(cache.state.count, 2) + + let url = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString).appendingPathExtension("safetensors") + try savePromptCache(url: url, cache: [cache], metadata: [:]) + let (loaded, _) = try loadPromptCache(url: url) + let restored = try XCTUnwrap(loaded[0] as? TurboQuantKVCache) + XCTAssertFalse(restored.isCompressed) + XCTAssertEqual(restored.offset, 16) + + // Decode through the compression transition on the restored cache. + let q = MLXRandom.normal([b, hq, 1, d], key: MLXRandom.key(83)) + let nk = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(84)) + let nv = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(85)) + let out = restored.compressedAttention(queries: q, keys: nk, values: nv, scale: 0.125) + XCTAssertTrue(restored.isCompressed) + XCTAssertEqual(restored.offset, 17) + XCTAssertFalse(MLX.isNaN(out).any().item(Bool.self)) + } + + // MARK: - GQA repeat factors the NR0 kernel must not mis-map + + func testStandardAttentionRepeatFactors() throws { + // rep=1 (MHA) and rep=3 (odd) previously aliased KV heads through the + // NR0 fast path; the dispatcher now excludes them. Compare against + // exact attention over the dequantized cache contents at loose 4-bit + // tolerance. + for (hq, hkv) in [(4, 4), (6, 2), (4, 2)] { + let (b, t, d) = (1, 192, 64) + let scale = 1.0 / Float(d).squareRoot() + let k = MLXRandom.normal([b, hkv, t, d], key: MLXRandom.key(91)) * 0.5 + let v = MLXRandom.normal([b, hkv, t, d], key: MLXRandom.key(92)) * 0.5 + let q = MLXRandom.normal([b, hq, 1, d], key: MLXRandom.key(93)) * 0.5 + let nk = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(94)) * 0.5 + let nv = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(95)) * 0.5 + + let cache = TurboQuantKVCache(bits: 4, keyBits: 4, valueBits: 4, seed: 5) + _ = cache.update(keys: k, values: v) + let out = cache.compressedAttention(queries: q, keys: nk, values: nv, scale: scale) + + let nRep = hq / hkv + let fullK = concatenated([k, nk], axis: 2) + let fullV = concatenated([v, nv], axis: 2) + let ek = MLX.tiled( + expandedDimensions(fullK, axis: 2), repetitions: [1, 1, nRep, 1, 1] + ).reshaped([b, hq, t + 1, d]) + let ev = MLX.tiled( + expandedDimensions(fullV, axis: 2), repetitions: [1, 1, nRep, 1, 1] + ).reshaped([b, hq, t + 1, d]) + let ref = matmul( + softmax(matmul(q, ek.transposed(0, 1, 3, 2)) * scale, axis: -1), ev) + + let a = out.asType(.float32).reshaped([-1]) + let r = ref.asType(.float32).reshaped([-1]) + let cos = ((a * r).sum() / (sqrt((a * a).sum()) * sqrt((r * r).sum()) + 1e-9)) + .item(Float.self) + XCTAssertGreaterThan(cos, 0.95, "rep=\(nRep): cos \(cos)") + } + } + + // MARK: - Trim on a compressed cache, then keep decoding + + func testCompressedTrimThenDecode() throws { + for (kb, vb) in [(0, 4), (8, 4), (4, 4)] { + let (b, hq, hkv, d) = (1, 4, 2, 64) + let cache = TurboQuantKVCache(bits: 4, keyBits: kb, valueBits: vb, seed: 3) + _ = cache.update( + keys: MLXRandom.normal([b, hkv, 32, d], key: MLXRandom.key(101)), + values: MLXRandom.normal([b, hkv, 32, d], key: MLXRandom.key(102))) + let q = MLXRandom.normal([b, hq, 1, d], key: MLXRandom.key(103)) + let nk = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(104)) + let nv = MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(105)) + _ = cache.compressedAttention(queries: q, keys: nk, values: nv, scale: 0.125) + XCTAssertTrue(cache.isCompressed) + XCTAssertEqual(cache.offset, 33) + + let trimmed = cache.trim(5) + XCTAssertEqual(trimmed, 5) + XCTAssertEqual(cache.offset, 28) + + // Two more decode steps after the partial trim must stay finite + // and advance the offset from the trimmed position. + for i in 0 ..< 2 { + let out = cache.compressedAttention( + queries: q, + keys: MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(UInt64(110 + i))), + values: MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(UInt64(120 + i))), + scale: 0.125) + XCTAssertFalse( + MLX.isNaN(out).any().item(Bool.self), "kb=\(kb) vb=\(vb) step \(i)") + } + XCTAssertEqual(cache.offset, 30, "kb=\(kb) vb=\(vb)") + } + } + + // MARK: - Speculative decoding × turbo caches + + func testSpeculativeDecodeWithTurboScheme() throws { + // Speculative rounds trim the cache every round and verify with L>1 + // chunks, the composition that previously hit the memory-unsafe raw + // update() path on compressed caches. + let config = LlamaConfiguration( + hiddenSize: 128, hiddenLayers: 6, intermediateSize: 128, attentionHeads: 8, + rmsNormEps: 1e-5, vocabularySize: 100, kvHeads: 4) + let main = LlamaModel(config) + let draft = LlamaModel(config) + quantize(model: main, groupSize: 64, bits: 4) + quantize(model: draft, groupSize: 64, bits: 4) + eval(main, draft) + + let prompt = MLXArray([1, 7, 3, 12, 5, 9, 2, 8, 4, 11]) + var iterator = try SpeculativeTokenIterator( + input: .init(text: .init(tokens: prompt)), + mainModel: main, draftModel: draft, + parameters: GenerateParameters(maxTokens: 12, kvScheme: "turbo0v4", temperature: 0), + numDraftTokens: 2) + + var produced = 0 + while iterator.next() != nil { + produced += 1 + // Mirror the generate() wrapper's per-round conversion. + var cache = iterator.mainCache + maybeQuantizeKVCache( + cache: &cache, kvBits: nil, quantizedKVStart: 0, kvScheme: "turbo0v4") + } + XCTAssertGreaterThan(produced, 4) + XCTAssertTrue( + iterator.mainCache.contains { $0 is TurboQuantKVCache }, + "main cache never converted") + } + + // MARK: - Dense (non-pow2) encode kernel parity with the reference codec + + func testDenseEncodeKernelMatchesCodecDim80() throws { + let d = 80 + let codec = MSECodec(dim: d, bits: 4, seed: 42) + let vectors = MLXRandom.normal([1, 1, 48, d], key: MLXRandom.key(131)) + let ref = codec.encode(vectors.asType(.float32)) + let (packed, norms) = TurboQuantKernelOps.fusedEncode( + input: vectors.reshaped([48, d]).asType(.float32), + rotation: codec.rotation, boundaries: codec.boundaries, + codebook: codec.codebook, bits: 4, dim: d) + // Exact bit parity is not achievable: the kernel's per-thread dot + // product and MLX's matmul reduce in different orders, flipping + // borderline quantization ties. Assert equivalence where it matters: + // reconstructions agree and only a small fraction of indices differ. + let unpackedK = TurboQuantPacking.unpackLowBit(packed, bits: 4, count: d) + let unpackedR = TurboQuantPacking.unpackLowBit( + ref.packedIndices.reshaped([48, -1]), bits: 4, count: d) + let mismatch = MLX.notEqual(unpackedK, unpackedR).asType(.float32).mean() + .item(Float.self) + XCTAssertLessThan(mismatch, 0.02, "index mismatch fraction \(mismatch)") + + let recK = matmul( + codec.codebook[unpackedK.asType(.int32)] * expandedDimensions(norms, axis: -1), + codec.rotation) + let recR = codec.decode(ref).reshaped([48, d]) + let a = recK.reshaped([-1]) + let r = recR.asType(.float32).reshaped([-1]) + let cos = ((a * r).sum() / (sqrt((a * a).sum()) * sqrt((r * r).sum()) + 1e-9)) + .item(Float.self) + XCTAssertGreaterThan(cos, 0.995, "dense kernel reconstruction diverges (cos \(cos))") + } + + // MARK: - Scaled (key-calibrated) encode kernel parity with the reference codec + // + // F-85 follow-up: per-dimension key calibration forced calibrated keys + // through MSECodec.encode(_:scale:)'s MLX-ops path, dropping decode + // throughput 124 -> 27.5 tps on Qwen3-1.7B turbo4. These verify the + // fused Metal kernel variants (fusedEncodeWHTScaled / fusedEncodeScaled) + // reproduce that math so calibrated keys can go back through the kernel. + + /// The calibrated encode path rotates via MSECodec.rotatedUnit (MLX ops) + /// and quantizes in the fused kernel; indices must match + /// MSECodec.encode(_:scale:) exactly since both quantize bit-identical + /// rotated values with the same boundary predicate. Norms carry a + /// reduce-order tolerance only. + func testScaledQuantizePackKernelMatchesCodecExactly() throws { + for dim in [64, 80, 128] { + let bits = 4 + let n = 128 + let codec = MSECodec(dim: dim, bits: bits, seed: 42) + var scaleVals = [Float](repeating: 1.0, count: dim) + for i in 0 ..< dim { + scaleVals[i] = 0.25 + 3.75 * Float(i) / Float(dim - 1) + } + let scale = MLXArray(scaleVals) + let vectors = MLXRandom.normal([1, 1, n, dim], key: MLXRandom.key(211)) + .asType(.float32) + let ref = codec.encode(vectors, scale: scale) + + let flat = vectors.reshaped([n, dim]) + let (rotated, norms) = codec.rotatedUnit(flat) + let (packed, kernelNorms) = TurboQuantKernelOps.fusedQuantizePackScaled( + rotated: rotated, rawNorms: norms, scale: scale, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: bits, dim: dim) + eval(packed, kernelNorms) + + let refIdx = TurboQuantPacking.unpackLowBit( + ref.packedIndices.reshaped([n, -1]), bits: bits, count: dim) + let kernelIdx = TurboQuantPacking.unpackLowBit(packed, bits: bits, count: dim) + let mismatches = (refIdx .!= kernelIdx).asType(.float32).sum().item(Float.self) + XCTAssertEqual(mismatches, 0, "dim=\(dim): indices must match exactly") + + let normRatio = kernelNorms / ref.norms.reshaped([n]) + let maxDev = (normRatio - MLXArray(Float(1.0))).abs().max().item(Float.self) + XCTAssertLessThan(maxDev, 1e-4, "dim=\(dim): norm deviation \(maxDev)") + } + } + + // MARK: - Buffer growth across the step allocation boundary + + func testCacheGrowthAcrossStepBoundary() throws { + // step = 256: prefill 250 then decode 12 steps so compressed storage + // re-allocates mid-decode for every layout. + for (kb, vb) in [(0, 4), (8, 4), (4, 4)] { + let (b, hq, hkv, d) = (1, 2, 1, 64) + let cache = TurboQuantKVCache(bits: 4, keyBits: kb, valueBits: vb, seed: 11) + _ = cache.update( + keys: MLXRandom.normal([b, hkv, 250, d], key: MLXRandom.key(141)), + values: MLXRandom.normal([b, hkv, 250, d], key: MLXRandom.key(142))) + let q = MLXRandom.normal([b, hq, 1, d], key: MLXRandom.key(143)) + for i in 0 ..< 12 { + let out = cache.compressedAttention( + queries: q, + keys: MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(UInt64(150 + i))), + values: MLXRandom.normal([b, hkv, 1, d], key: MLXRandom.key(UInt64(170 + i))), + scale: 0.125) + XCTAssertFalse( + MLX.isNaN(out).any().item(Bool.self), + "kb=\(kb) vb=\(vb) offset \(cache.offset)") + } + XCTAssertEqual(cache.offset, 262, "kb=\(kb) vb=\(vb)") + } + } + +} From b64793cf053ff291e88bb53cdee86621e8577af5 Mon Sep 17 00:00:00 2001 From: TheTom Date: Thu, 16 Jul 2026 15:07:20 -0500 Subject: [PATCH 3/3] Add KV cache quantization documentation Scheme reference table, the asymmetric-first selection ladder, measured family sensitivity and calibration notes, performance guidance, and limitations. --- .../Documentation.docc/Documentation.md | 1 + .../kv-cache-quantization.md | 109 ++++++++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 Libraries/MLXLMCommon/Documentation.docc/kv-cache-quantization.md diff --git a/Libraries/MLXLMCommon/Documentation.docc/Documentation.md b/Libraries/MLXLMCommon/Documentation.docc/Documentation.md index c9dd22b6b..f1d480d1f 100644 --- a/Libraries/MLXLMCommon/Documentation.docc/Documentation.md +++ b/Libraries/MLXLMCommon/Documentation.docc/Documentation.md @@ -7,6 +7,7 @@ Common language model code. - - - +- ## Other MLX Libraries Packages diff --git a/Libraries/MLXLMCommon/Documentation.docc/kv-cache-quantization.md b/Libraries/MLXLMCommon/Documentation.docc/kv-cache-quantization.md new file mode 100644 index 000000000..5338e932b --- /dev/null +++ b/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 `turbov`; `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 ; 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).