Aligned weight loading and ring KV cache for sliding-window layers (1.2.1)

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## 1.2.1
+
+- Fixes the GPU out-of-memory error when running a large model such as Gemma 4 12B on a 16 GB Mac. MLX writes `safetensors` without the 16-byte alignment the kernels need, so every tensor was copied into its own Metal buffer while the mapped file stayed referenced: the GPU held two copies of the weights (18 GB for a 9 GB model). Shards are now copied once, in file order, into a page-aligned buffer and the mapping is released. Loading is faster too: Gemma 4 12B 156 s → 25 s from disk, 11 s from the page cache; Mistral 7B 3.3 → 2.9 s.
+- Sliding-window layers (Gemma 2, 3 and 4) keep a ring of window + 512 K/V rows instead of the whole context: Gemma 4 12B at 8K context needs 0.5 GB of cache instead of 2.7 GB, and runs at 32K on 16 GB. A request whose shared prefix with the previous one ends more than 512 tokens before that request's end no longer reuses the cache (its rows are gone); it is processed from the start.
+- A GPU out-of-memory failure is reported as such, with the advice to lower the context length, instead of as a kernel compilation error; the load-time memory check includes the scratch buffers.
+
 ## 1.2.0
 
 Gemma support (`gemma`, `gemma2`, `gemma3` / `gemma3_text`, `gemma4` / `gemma4_text`, `gemma4_unified`):

README.md

@@ -9,7 +9,7 @@ A native macOS app for discovering, downloading and chatting with **MLX models f
 ![MLX Runner](site/assets/screenshot-chat.jpg)
 
 - **Zero third-party dependencies.** The project uses the Swift standard library and Apple's system frameworks only: SwiftUI, Foundation, Metal, Network and CryptoKit. It has no SwiftPM packages, no vendored code and no Python.
-- **Native MLX format.** It reads MLX-community repos directly: memory-mapped `safetensors` plus MLX affine quantization at 2, 3, 4, 5, 6 or 8 bits, with group sizes of 32, 64 or 128.
+- **Native MLX format.** It reads MLX-community repos directly: `safetensors` plus MLX affine quantization at 2, 3, 4, 5, 6 or 8 bits, with group sizes of 32, 64 or 128.
 - **Fast.** Hand-written Metal kernels cover every layer:
   - fused decode matvecs
   - simdgroup-matrix GEMMs for prefill
@@ -101,7 +101,7 @@ Each checkpoint above passes end-to-end generation and the prefill-versus-decode
 ## Limitations
 
 - **Unsupported model types:** mixture-of-experts (Mixtral, Qwen3-MoE, Gemma 4 26B-A4B), Gemma 3n, and the vision and audio parts of multimodal models (Gemma 3 and Gemma 4 exports run as text-only). Discover marks unsupported types as not runnable.
-- **Gemma:** no tool calling and no Gemma 4 thinking mode (Gemma 4 runs with thinking off, its template's default). Sliding-window layers keep a full-length KV cache, so long contexts cost as much memory as on a model without sliding windows.
+- **Gemma:** no tool calling and no Gemma 4 thinking mode (Gemma 4 runs with thinking off, its template's default). Sliding-window layers keep a ring of window + 512 K/V rows, so a chat whose shared prefix with the previous request ends more than 512 tokens before that request's end is processed from the start instead of reusing the cache.
 - **Chat templates:** the prompt format is detected (ChatML, Llama 3, Mistral, Phi-3, Gemma, Gemma 4) rather than executed as Jinja, so a model with an unusual template may get a near-match format.
 - **Prefill speed:** about 165 tok/s on a 7B model on an M2 Pro, because the prefill GEMM reaches about 35% of fp32 peak.
 - **Speculative decoding is opt-in.** Prompt-lookup drafting with exact verification is implemented and tested (greedy output is identical), but checking 4 drafted tokens on a 7B model still costs about 3 decode steps. So it only helps when most drafts are accepted (copying or editing text), and it is off by default (`SamplingParameters.speculative`, or `mlxr run --spec`).
@@ -118,10 +118,12 @@ Measured with `mlxr bench` on an **M2 Pro (16 GB)**. Decode is measured with gre
 | Llama 3.2 1B Instruct | 695 MB | 1,180 tok/s | 203 tok/s | 188 tok/s |
 | Phi-3.5 mini | 2.15 GB | 320 tok/s | 70 tok/s | 61 tok/s |
 | Mistral 7B v0.3 | 4.08 GB | 166 tok/s | 41.0 tok/s | 39.5 tok/s |
+| Gemma 3 1B (QAT) | 732 MB | 1,450 tok/s | 195 tok/s | 178 tok/s |
+| Gemma 3 4B | 2.56 GB | 345 tok/s | 67 tok/s | 62 tok/s |
 
 **The 7B ceiling.** Decoding one token reads every weight once, so speed is capped by memory bandwidth. This M2 Pro measures 197 GB/s with a plain read kernel, which caps a 4.08 GB model at about 48 tok/s. MLX Runner reaches 41.0 tok/s, about 85% of that cap. Reaching 50+ tok/s on a 4-bit 7B model therefore needs one of: a Mac with more bandwidth (M-series Max chips have about 400 GB/s), a 3-bit model (which would raise the ceiling to about 60 tok/s, but only 4-bit has an optimized decode kernel so far), or speculative decoding (see Limitations).
 
-Load times stay short because the weights are `mmap`ed and wrapped in Metal buffers with no copy: about 0.5 s for the small models, 1.6 s for Phi-3.5 and 3.3 s for Mistral 7B. Those figures include a warm-up pass that pages the weights in.
+Loading copies the weights once, in file order, into a page-aligned Metal buffer (MLX writes `safetensors` without the 16-byte alignment the kernels' vector loads need), then releases the file mapping, so the GPU holds exactly one copy: about 0.5 s for the small models, 3 s for Mistral 7B and 11 s for the 9 GB Gemma 4 12B from the page cache (25 s from disk). Those figures include a warm-up pass.
 
 ## Build and run
 
@@ -156,7 +158,7 @@ mlxr tokenize <model> "some text"
 Sources/
   MLXEngine/            Inference engine (library)
     Kernels.swift       Metal shader source, compiled at runtime
-    GPU.swift           Device and pipelines; zero-copy safetensors loader
+    GPU.swift           Device and pipelines; safetensors loader (one aligned copy per shard)
     Transformer.swift   Model graph: one command buffer per step, fused kernels
     Tokenizer.swift     BPE from tokenizer.json (byte-level and SentencePiece)
     ChatTemplate.swift  ChatML / Llama 3 / Mistral / Phi-3 prompt formats
@@ -172,7 +174,7 @@ Tests/MLXEngineTests/   Swift Testing suite
 How the engine works:
 
 - **Decode** (one token per step):
-  - q/k/v and gate/up are each one fused quantized matvec over several zero-copy weight matrices, with bias folded in.
+  - q/k/v and gate/up are each one fused quantized matvec over several weight matrices, with bias folded in.
   - Per-head RMSNorm, RoPE and the KV-cache append run as a single kernel.
   - The residual add is fused with the next RMSNorm.
   - Attention uses online softmax across 16 to 32 simdgroups.

Sources/MLXEngine/Engine.swift

@@ -398,6 +398,8 @@ public final class InferenceEngine: @unchecked Sendable {
         var common = 0
         let limit = min(cachedTokens.count, tokens.count - 1)
         while common < limit && cachedTokens[common] == tokens[common] { common += 1 }
+        common = model.reusablePrefix(common, previousLength: cachedTokens.count)
+        if common == 0 { model.resetWriteTracking() }
         stats.cachedPromptTokens = common
 
         let vocab = model.vocabSize
@@ -591,8 +593,9 @@ extension InferenceEngine {
         }
         let cfg = try ModelConfig(directory: directory)
         let ctx = min(contextLength, max(cfg.maxPositionEmbeddings, 256))
-        let kv = cfg.kvBytesPerToken * ctx
-        let need = UInt64(weights + kv)
+        let kv = cfg.kvCacheBytes(context: ctx)
+        let scratch = Transformer.scratchBytes(config: cfg)
+        let need = UInt64(weights + kv + scratch)
         let budget = GPU.shared.device.recommendedMaxWorkingSetSize
         guard need < budget else {
             let f = ByteCountFormatter()

Sources/MLXEngine/GPU.swift

@@ -179,12 +179,24 @@ extension MTLComputeCommandEncoder {
     }
 }
 
-// MARK: - Safetensors (zero-copy via mmap)
+// MARK: - Safetensors
 
+/// One safetensors shard in Metal buffers. The kernels use 16-byte vector loads, so every
+/// tensor must start on a 16-byte boundary. A file whose tensors all do is wrapped in place
+/// (mmap, zero-copy). MLX-written files never do (the header is padded to 8 bytes, and 2-byte
+/// tensors such as Gemma 4's `layer_scalar` shift everything after them), so their tensors are
+/// copied once, in file order, each into its own buffer, and the mapping is released: the GPU
+/// then holds exactly one copy of the weights. (Copying while the mapping stayed referenced
+/// doubled the working set — a 9 GB model needed 18 GB. One arena for all tensors measured
+/// 6% slower at 1K context on the smallest model, so each tensor keeps its own buffer.)
 final class SafeTensorFile {
-    let buffer: MTLBuffer
+    /// The mapped file (zero-copy) or, for copied files, a buffer per tensor.
+    private let buffer: MTLBuffer?
+    private var buffers: [String: MTLBuffer] = [:]
     private(set) var infos: [String: (dtype: DType, shape: [Int], start: Int, end: Int)] = [:]
     let url: URL
+    /// Bytes of weights held in Metal buffers.
+    private(set) var byteCount = 0
 
     init(url: URL, device: MTLDevice) throws {
         self.url = url
@@ -201,13 +213,8 @@ final class SafeTensorFile {
         }
         // Ask the kernel to start paging weights in now.
         madvise(ptr, mapped, MADV_WILLNEED)
-        guard let buf = device.makeBuffer(bytesNoCopy: ptr, length: mapped, options: .storageModeShared, deallocator: { p, len in
-            munmap(p, len)
-        }) else {
-            munmap(ptr, mapped)
-            throw EngineError.invalidFile("cannot wrap \(url.lastPathComponent) in a Metal buffer")
-        }
-        buffer = buf
+        var mappingOwned = true   // released below unless the mapping becomes the buffer
+        defer { if mappingOwned { munmap(ptr, mapped) } }
         let base = ptr.assumingMemoryBound(to: UInt8.self)
         let file = url.lastPathComponent
         guard size >= 8 else { throw EngineError.invalidFile("\(file) is truncated") }
@@ -250,17 +257,32 @@ final class SafeTensorFile {
             }
             infos[name] = (dtype, shape.map { Int($0) }, dataStart + Int(offs[0]), dataStart + Int(offs[1]))
         }
-    }
 
-    func tensor(_ name: String) -> Tensor? {
-        guard let i = infos[name] else { return nil }
-        // Kernels use vector loads; copy the rare misaligned tensor into its own buffer.
-        if i.start % 16 != 0 {
+        if infos.values.allSatisfy({ $0.start % 16 == 0 }) {
+            guard let buf = device.makeBuffer(bytesNoCopy: ptr, length: mapped, options: .storageModeShared, deallocator: { p, len in
+                munmap(p, len)
+            }) else { throw EngineError.invalidFile("cannot wrap \(file) in a Metal buffer") }
+            buffer = buf
+            byteCount = size
+            mappingOwned = false
+            return
+        }
+        // Copy in file order (sequential read), one buffer per tensor.
+        buffer = nil
+        for (name, i) in infos.sorted(by: { $0.value.start < $1.value.start }) {
             let len = i.end - i.start
-            let copy = buffer.device.makeBuffer(bytes: buffer.contents() + i.start, length: max(len, 16), options: .storageModeShared)!
-            return Tensor(buffer: copy, offset: 0, dtype: i.dtype, shape: i.shape)
+            guard let b = device.makeBuffer(bytes: base + i.start, length: max(len, 16), options: .storageModeShared) else {
+                throw EngineError.outOfMemory("could not allocate \(ByteCountFormatter.string(fromByteCount: Int64(len), countStyle: .file)) for \(name) in \(file)")
+            }
+            buffers[name] = b
+            byteCount += len
+            infos[name] = (i.dtype, i.shape, 0, len)
         }
-        return Tensor(buffer: buffer, offset: i.start, dtype: i.dtype, shape: i.shape)
+    }
+
+    func tensor(_ name: String) -> Tensor? {
+        guard let i = infos[name], let b = buffer ?? buffers[name] else { return nil }
+        return Tensor(buffer: b, offset: i.start, dtype: i.dtype, shape: i.shape)
     }
 }
 
@@ -321,5 +343,5 @@ final class WeightStore {
         return t
     }
 
-    var totalBytes: Int { files.reduce(0) { $0 + $1.buffer.length } }
+    var totalBytes: Int { files.reduce(0) { $0 + $1.byteCount } }
 }

Sources/MLXEngine/Kernels.swift

@@ -14,6 +14,9 @@ struct NormParams { int dim; int headsPerRow; int rowStride; float eps; };
 struct AttnParams { int D; int nH; int nKV; int qStride; int maxSeq; int offset; float scale; int outStride; int T; int window; float softcap; };
 // Sliding window (window > 0): the query at position i sees keys i - window < j <= i.
 inline int attn_first_key(constant AttnParams& p, int pos) { return p.window > 0 ? max(0, pos - p.window + 1) : 0; }
+// K/V cache row of position j: caches hold `len` rows (AttnParams / QKVPostParams.maxSeq) and
+// a sliding-window layer keeps only the last ones, position j living at row j mod len.
+inline int kv_slot(int j, int len) { return j < len ? j : j % len; }
 // Gemma 2 attention-score soft-capping: s = tanh(s / cap) * cap. A template parameter of the
 // attention kernels (`_cap` instantiations) so the other models' inner loops do not carry it.
 template <bool CAP>
@@ -730,6 +733,7 @@ kernel void qkv_post(device float*       qkv     [[buffer(0)]],
     const int h = int(tg.x), t = int(tg.y), i = int(tid);
     device float* base = qkv + t * p.rowStride + h * p.D;
     const int pos = p.offset + t;
+    const int row = kv_slot(pos, p.maxSeq);
     const bool isV = h >= p.nH + p.nKV;
     float v = (isV && p.kEqV) ? base[i - p.nKV * p.D] : base[i];
     threadgroup float buf[1024];
@@ -743,7 +747,7 @@ kernel void qkv_post(device float*       qkv     [[buffer(0)]],
             for (uint s = 0; s < nsg; s++) tot += red[s];
             v *= rsqrt(tot / float(p.D) + p.eps);
         }
-        vc[(ulong(h - p.nH - p.nKV) * ulong(p.maxSeq) + ulong(pos)) * ulong(p.D) + i] = half(v);
+        vc[(ulong(h - p.nH - p.nKV) * ulong(p.maxSeq) + ulong(row)) * ulong(p.D) + i] = half(v);
         return;
     }
     const bool isQ = h < p.nH;
@@ -768,7 +772,7 @@ kernel void qkv_post(device float*       qkv     [[buffer(0)]],
         out = first ? v * c - buf[partner] * s : v * c + buf[partner] * s;
     }
     if (isQ) base[i] = out;
-    else kc[(ulong(h - p.nH) * ulong(p.maxSeq) + ulong(pos)) * ulong(p.D) + i] = half(out);
+    else kc[(ulong(h - p.nH) * ulong(p.maxSeq) + ulong(row)) * ulong(p.D) + i] = half(out);
 }
 
 // h += r; y = RMSNorm(h) * w   (residual add fused with the next norm)
@@ -917,7 +921,7 @@ kernel void attention(device const float* qkv [[buffer(0)]],
     for (int j0 = attn_first_key(p, L - 1) + int(sg); j0 < L; j0 += ANSG * AU) {
         half kh[AU][E], vh[AU][E];
         for (int u = 0; u < AU; u++) {
-            const int j = min(j0 + u * ANSG, L - 1);   // clamped: rows past L are loaded but unused
+            const int j = kv_slot(min(j0 + u * ANSG, L - 1), p.maxSeq);   // clamped: rows past L are loaded but unused
             device const half* kp = kb + ulong(j) * ulong(p.D);
             device const half* vp = vb + ulong(j) * ulong(p.D);
             for (int e = 0; e < E; e++) { kh[u][e] = kp[e * 32 + lane]; vh[u][e] = vp[e * 32 + lane]; }
@@ -969,7 +973,7 @@ kernel void attention_split(device const float* qkv  [[buffer(0)]],
     for (int j0 = attn_first_key(p, L - 1) + v; j0 < L; j0 += ASPLIT_V * AU) {
         half kh[AU][E], vh[AU][E];
         for (int u = 0; u < AU; u++) {
-            const int j = min(j0 + u * ASPLIT_V, L - 1);
+            const int j = kv_slot(min(j0 + u * ASPLIT_V, L - 1), p.maxSeq);
             device const half* kp = kb + ulong(j) * ulong(p.D);
             device const half* vp = vb + ulong(j) * ulong(p.D);
             for (int e = 0; e < E; e++) { kh[u][e] = kp[e * 32 + lane]; vh[u][e] = vp[e * 32 + lane]; }
@@ -1030,8 +1034,8 @@ kernel void attention_prefill(device const float* qkv [[buffer(0)]],
     const int jlo = attn_first_key(p, p.offset + t0);   // the first query's first key
     for (int j = jlo + int(sg); j < L; j += NSG) {
         float kv[E], vv[E];
-        device const half* kp = kb + ulong(j) * ulong(p.D);
-        device const half* vp = vb + ulong(j) * ulong(p.D);
+        device const half* kp = kb + ulong(kv_slot(j, p.maxSeq)) * ulong(p.D);
+        device const half* vp = vb + ulong(kv_slot(j, p.maxSeq)) * ulong(p.D);
         for (int e = 0; e < E; e++) { kv[e] = float(kp[e * 32 + lane]); vv[e] = float(vp[e * 32 + lane]); }
         for (int qi = 0; qi < QB; qi++) {
             if (qi >= nq || j > p.offset + t0 + qi) continue;   // causal mask (uniform per simdgroup)
@@ -1122,11 +1126,14 @@ kernel void flash_prefill(device const float* qkv [[buffer(0)]],
     const int qFirst = attn_first_key(p, p.offset + t0);       // first key of the block's first query
     for (int j0 = qFirst / BK * BK; j0 < L; j0 += BK) {
         threadgroup_barrier(mem_flags::mem_threadgroup);
+        // D is a power of two: row / column from i are shifts; the ring slot costs a modulo only
+        // for positions past the cache length.
         for (int i = int(tid); i < BK * D; i += 128) {
-            const int j = j0 + i / D;
+            const int r = i / D, j = j0 + r;
             const bool ok = j < L;
-            Ks[i] = ok ? float(kb[ulong(j0) * D + i]) : 0.0f;
-            Vs[i] = ok ? float(vb[ulong(j0) * D + i]) : 0.0f;
+            const ulong src = ulong(kv_slot(j, p.maxSeq)) * D + (i - r * D);
+            Ks[i] = ok ? float(kb[src]) : 0.0f;
+            Vs[i] = ok ? float(vb[src]) : 0.0f;
         }
         threadgroup_barrier(mem_flags::mem_threadgroup);
         // Skip key blocks that are entirely in the future for all 8 rows of this simdgroup.

Sources/MLXEngine/ModelConfig.swift

@@ -97,8 +97,21 @@ public struct ModelConfig {
     public var perLayerInputSize = 0
     public var valueNorm = false
 
-    /// KV cache bytes per token of context (f16 K and V, layers with their own cache).
-    public var kvBytesPerToken: Int { layers.reduce(0) { $0 + ($1.hasOwnKV ? 2 * $1.numKVHeads * $1.headDim * 2 : 0) } }
+    /// Largest batch of positions one forward pass writes (the engine's scratch buffers and
+    /// the ring caches are sized for it).
+    public static let maxBatch = 512
+
+    /// Rows a layer's K/V cache holds for a context: all of it, or for a sliding-window layer
+    /// the window plus one batch of new positions (see `TransformerLayer.cacheLen`).
+    static func cacheRows(_ spec: LayerSpec, context: Int, window: Int) -> Int {
+        spec.sliding && window > 0 ? min(context, window + maxBatch) : context
+    }
+    func cacheRows(_ spec: LayerSpec, context: Int) -> Int { Self.cacheRows(spec, context: context, window: slidingWindow) }
+
+    /// KV cache bytes for a context (f16 K and V, layers with their own cache).
+    public func kvCacheBytes(context: Int) -> Int {
+        layers.reduce(0) { $0 + ($1.hasOwnKV ? 2 * $1.numKVHeads * $1.headDim * 2 * cacheRows($1, context: context) : 0) }
+    }
 
     public init(directory: URL) throws {
         let url = directory.appendingPathComponent("config.json")