From f428070074838f3f66a4a40a5b2f93783459996a Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Sun, 1 Nov 2026 18:16:57 +0100 Subject: [PATCH 1/3] Load weights into one aligned buffer; ring KV cache for sliding-window layers (1.2.1) MLX writes safetensors without the 16-byte alignment the kernels' vector loads need, so every tensor was copied into its own Metal buffer while the mapped file stayed referenced by the command buffer: the GPU held two copies of the weights, and a 9 GB Gemma 4 12B ran out of GPU memory on a 16 GB Mac. Shards are now copied once, in file order, into a page-aligned buffer and the mapping is released. Loading is faster (12B: 156 s -> 25 s from disk, 11 s from the page cache). Sliding-window layers keep a ring of window + 512 K/V rows (position p at row p mod rows) instead of the whole context: the 12B at 8K context needs 0.5 GB of cache instead of 2.7 GB and runs at 32K. Prefix reuse is refused when the previous run went on for more than a batch past the divergence point, since those rows are gone. A GPU out-of-memory failure is reported as such, with advice; the load-time budget includes the scratch buffers. README and the landing page: load-time and no-copy wording, Gemma rows in the performance tables. Existing models: greedy output token-identical; Qwen3 0.6B and Mistral 7B decode within 1% (3 alternating pairs); Gemma 3 1B and Gemma 4 E2B unchanged. Co-Authored-By: Claude Fable 5.1 --- CHANGELOG.md | 6 +++ README.md | 12 +++--- Sources/MLXEngine/Engine.swift | 10 ++++- Sources/MLXEngine/GPU.swift | 54 +++++++++++++++++++------- Sources/MLXEngine/Kernels.swift | 21 ++++++---- Sources/MLXEngine/ModelConfig.swift | 13 ++++++- Sources/MLXEngine/Transformer.swift | 43 ++++++++++++++++---- Sources/MLXEngine/Version.swift | 2 +- Tests/MLXEngineTests/GemmaTests.swift | 29 +++++++++++++- Tests/MLXEngineTests/KernelTests.swift | 47 ++++++++++++++++++++++ site/index.html | 8 +++- 11 files changed, 203 insertions(+), 42 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 532b40c..170940e 100644 --- a/CHANGELOG.md +++ b/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`): diff --git a/README.md b/README.md index 152a517..45c9ecd 100644 --- a/README.md +++ b/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 "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. diff --git a/Sources/MLXEngine/Engine.swift b/Sources/MLXEngine/Engine.swift index b973aca..386e25d 100644 --- a/Sources/MLXEngine/Engine.swift +++ b/Sources/MLXEngine/Engine.swift @@ -398,6 +398,7 @@ 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) stats.cachedPromptTokens = common let vocab = model.vocabSize @@ -591,8 +592,13 @@ 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) + // Scratch buffers for a 512-token batch (hidden, projections, MLP, logits), see Transformer.init. + let H = cfg.hiddenSize, I = cfg.intermediateSize * (cfg.isGemma4 ? 2 : 1) + let qkv = cfg.layers.map { (cfg.numHeads + 2 * $0.numKVHeads) * $0.headDim }.max() ?? 0 + let attn = cfg.layers.map { cfg.numHeads * $0.headDim }.max() ?? 0 + let scratch = 512 * 4 * (4 * H + 3 * I + qkv + attn) + 6 * cfg.vocabSize * 4 + let need = UInt64(weights + kv + scratch) let budget = GPU.shared.device.recommendedMaxWorkingSetSize guard need < budget else { let f = ByteCountFormatter() diff --git a/Sources/MLXEngine/GPU.swift b/Sources/MLXEngine/GPU.swift index e41308b..2b6139e 100644 --- a/Sources/MLXEngine/GPU.swift +++ b/Sources/MLXEngine/GPU.swift @@ -179,8 +179,15 @@ extension MTLComputeCommandEncoder { } } -// MARK: - Safetensors (zero-copy via mmap) +// MARK: - Safetensors +/// One safetensors shard as a Metal buffer. 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, into one 16-byte-padded arena and the mapping is released: the +/// GPU then holds exactly one copy of the weights. (Copying each tensor into its own buffer +/// while the mapping stayed referenced doubled the working set — a 9 GB model needed 18 GB.) final class SafeTensorFile { let buffer: MTLBuffer private(set) var infos: [String: (dtype: DType, shape: [Int], start: Int, end: Int)] = [:] @@ -201,13 +208,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,16 +252,40 @@ final class SafeTensorFile { } infos[name] = (dtype, shape.map { Int($0) }, dataStart + Int(offs[0]), dataStart + Int(offs[1])) } + + 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 + mappingOwned = false + return + } + // Arena: tensors in file order, each on a page boundary (what separate buffers gave; a + // 16-byte packing measured a few percent slower on the smallest model). + let align = page + let ordered = infos.sorted { $0.value.start < $1.value.start } + var total = 0 + var placed: [String: Int] = [:] + for (name, i) in ordered { + placed[name] = total + let (next, o) = total.addingReportingOverflow((i.end - i.start) + align - 1) + guard !o else { throw EngineError.invalidFile("\(file) is too large") } + total = next / align * align + } + guard let arena = device.makeBuffer(length: max(total, 16), options: .storageModeShared) else { + throw EngineError.outOfMemory("could not allocate \(ByteCountFormatter.string(fromByteCount: Int64(total), countStyle: .file)) for \(file)") + } + let dst = arena.contents() + for (name, i) in ordered { + memcpy(dst + placed[name]!, base + i.start, i.end - i.start) + infos[name] = (i.dtype, i.shape, placed[name]!, placed[name]! + (i.end - i.start)) + } + buffer = arena } 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 { - 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) - } return Tensor(buffer: buffer, offset: i.start, dtype: i.dtype, shape: i.shape) } } diff --git a/Sources/MLXEngine/Kernels.swift b/Sources/MLXEngine/Kernels.swift index 57d32eb..3ea2d23 100644 --- a/Sources/MLXEngine/Kernels.swift +++ b/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 @@ -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) @@ -1125,8 +1129,9 @@ kernel void flash_prefill(device const float* qkv [[buffer(0)]], for (int i = int(tid); i < BK * D; i += 128) { const int j = j0 + i / D; 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 % 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. diff --git a/Sources/MLXEngine/ModelConfig.swift b/Sources/MLXEngine/ModelConfig.swift index 02693b3..f40320c 100644 --- a/Sources/MLXEngine/ModelConfig.swift +++ b/Sources/MLXEngine/ModelConfig.swift @@ -97,8 +97,17 @@ 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) } } + /// 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 + 512) : 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") diff --git a/Sources/MLXEngine/Transformer.swift b/Sources/MLXEngine/Transformer.swift index 63c2de3..acda52f 100644 --- a/Sources/MLXEngine/Transformer.swift +++ b/Sources/MLXEngine/Transformer.swift @@ -300,6 +300,10 @@ final class TransformerLayer { /// nil for a KV-shared layer (it attends over `spec.kvSource`'s cache). let kCache: MTLBuffer? let vCache: MTLBuffer? + /// Rows in the K/V cache: the whole context, or for a sliding-window layer the window plus + /// one batch (position p lives at row p mod cacheLen; a batch's T new rows must not land + /// on rows its own queries still read). + let cacheLen: Int /// Gemma 4 per-layer input: gate (H → P), projection (P → H) and its norm. let pleGate: Linear? let pleProj: Linear? @@ -367,8 +371,9 @@ final class TransformerLayer { qkvFusable = qkv.count > 1 && Linear.fusable(qkv) gateUpFusable = gateUp.count > 1 && Linear.fusable(gateUp) qkvPostPSO = try GPU.shared.pipeline("qkv_post_\(qNorm?.weight.dtype.suffix ?? "f32")") + cacheLen = c.cacheRows(spec, context: maxSeq) if spec.hasOwnKV { - let cacheBytes = nKV * maxSeq * D * 2 + let cacheBytes = nKV * cacheLen * D * 2 kCache = try GPU.shared.buffer(cacheBytes) vCache = try GPU.shared.buffer(cacheBytes) } else { @@ -377,6 +382,18 @@ final class TransformerLayer { } } +extension EngineError { + /// The engine error for a failed command buffer: an out-of-memory failure gets an + /// actionable message, anything else is reported as a GPU error. + static func gpu(_ err: Error) -> EngineError { + let d = err.localizedDescription + if d.contains("OutOfMemory") || d.lowercased().contains("insufficient memory") { + return .outOfMemory("the GPU ran out of memory while running the model. Lower the context length (Settings) or use a smaller model.") + } + return .kernelCompile("GPU error: \(d)") + } +} + /// A decoder-only transformer running entirely on Metal. public final class Transformer { public let config: ModelConfig @@ -389,7 +406,8 @@ public final class Transformer { let quantDescription: String // Scratch buffers sized for `maxBatch` tokens. - let maxBatch = 512 + static let maxBatch = 512 + let maxBatch = Transformer.maxBatch /// Token ids by position (maxSeq + 1 entries): the CPU writes prompt tokens, the GPU sampler /// appends generated ones, and every step's embedding lookup reads its tokens from here. private let hist: MTLBuffer @@ -653,7 +671,7 @@ public final class Transformer { lastGPUTime = cb.gpuEndTime - cb.gpuStartTime if let err = cb.error { prepared = nil - throw EngineError.kernelCompile("GPU error: \(err.localizedDescription)") + throw EngineError.gpu(err) } } @@ -670,6 +688,17 @@ public final class Transformer { public let offset: Int } + /// How much of a previous run's KV cache (positions 0.. Int { + for l in layers where l.spec.hasOwnKV && l.cacheLen < maxSeq { + if previousLength - common > l.cacheLen - config.slidingWindow + 1 { return 0 } + } + return common + } + /// Seeds the GPU sampler's xorshift64 generator. Call before `prefill` (nothing may be in flight). public func seedSampler(_ seed: UInt64) { rngState.contents().storeBytes(of: seed == 0 ? 0x9E3779B97F4A7C15 : seed, as: UInt64.self) @@ -697,7 +726,7 @@ public final class Transformer { cb.commit() cb.waitUntilCompleted() lastGPUTime = cb.gpuEndTime - cb.gpuStartTime - if let err = cb.error { throw EngineError.kernelCompile("GPU error: \(err.localizedDescription)") } + if let err = cb.error { throw EngineError.gpu(err) } start += n } return try sampledToken(at: offset + tokens.count) @@ -719,7 +748,7 @@ public final class Transformer { public func finish(_ step: DecodeStep) throws -> Int { step.cb.waitUntilCompleted() lastGPUTime = step.cb.gpuEndTime - step.cb.gpuStartTime - if let err = step.cb.error { throw EngineError.kernelCompile("GPU error: \(err.localizedDescription)") } + if let err = step.cb.error { throw EngineError.gpu(err) } return try sampledToken(at: step.offset + 1) } @@ -790,7 +819,7 @@ public final class Transformer { enc.set(ropeFreqs[li], 0, 4); enc.set(ropeFreqs[li], 0, 5) } enc.params(QKVPostParams(D: Int32(D), nH: Int32(nH), nKV: Int32(nKV), rowStride: Int32(qkvCols), - offset: Int32(offset), maxSeq: Int32(maxSeq), traditional: c.ropeTraditional ? 1 : 0, + offset: Int32(offset), maxSeq: Int32(kv.cacheLen), traditional: c.ropeTraditional ? 1 : 0, rotDim: Int32(spec.rotDim), eps: c.rmsNormEps, hasNorm: layer.qNorm != nil ? 1 : 0, vNorm: c.valueNorm ? 1 : 0, kEqV: spec.kEqV ? 1 : 0), 6) // A KV-shared layer only prepares its queries. @@ -799,7 +828,7 @@ public final class Transformer { } let attnParams = AttnParams(D: Int32(D), nH: Int32(nH), nKV: Int32(nKV), qStride: Int32(qkvCols), - maxSeq: Int32(maxSeq), offset: Int32(offset), scale: c.ropeAttentionFactor * c.ropeAttentionFactor * c.queryScale, + maxSeq: Int32(kv.cacheLen), offset: Int32(offset), scale: c.ropeAttentionFactor * c.ropeAttentionFactor * c.queryScale, outStride: Int32(nH * D), T: Int32(T), window: Int32(spec.sliding ? c.slidingWindow : 0), softcap: c.attnLogitSoftcap) if skipAttn { diff --git a/Sources/MLXEngine/Version.swift b/Sources/MLXEngine/Version.swift index f7aa7ca..7b2d56a 100644 --- a/Sources/MLXEngine/Version.swift +++ b/Sources/MLXEngine/Version.swift @@ -1,7 +1,7 @@ /// The MLX Runner release version. The release workflow refuses a tag that doesn't match it, /// and scripts/bundle.sh writes it into the app's Info.plist. public enum MLXRunnerVersion { - public static let current = "1.2.0" + public static let current = "1.2.1" public static let repository = "BackendStack21/mlx-runner" public static let homepage = "https://mlx-runner.21no.de" } diff --git a/Tests/MLXEngineTests/GemmaTests.swift b/Tests/MLXEngineTests/GemmaTests.swift index 1e9cb6c..15215c1 100644 --- a/Tests/MLXEngineTests/GemmaTests.swift +++ b/Tests/MLXEngineTests/GemmaTests.swift @@ -99,7 +99,9 @@ struct GemmaConfigTests { #expect(c.layers.map(\.headDim) == [64, 64, 128, 64, 64, 128]) // Layers 4 and 5 reuse the caches of the last earlier layer of their type. #expect(c.layers.map(\.kvSource) == [-1, -1, -1, -1, 3, 2]) - #expect(c.kvBytesPerToken == 2 * 2 * (3 * 64 + 128)) + // Sliding layers keep window + one batch of rows, full layers the whole context. + #expect(c.kvCacheBytes(context: 4096) == 2 * 2 * (3 * 64 * (16 + 512) + 128 * 4096)) + #expect(c.kvCacheBytes(context: 256) == 2 * 2 * (3 * 64 + 128) * 256) // Proportional RoPE: the whole head rotates, a quarter of the frequencies are non-zero, // with exponents over the full head dim. let g = c.layers[2] @@ -153,6 +155,31 @@ struct GemmaModelTests { } } + /// With a context longer than window + batch the sliding layers keep a ring of rows: a prompt + /// past that point must still prefill and decode alike, and a run that went on for more than + /// a batch past the divergence point cannot reuse its cache. + @Test(.enabled(if: installedModel(gemma3) != nil)) + func ringCacheMatchesAndLimitsPrefixReuse() throws { + guard let dir = installedModel(Self.gemma3) else { return } + let model = try Transformer(directory: dir, contextLength: 2048) // window 512 → 1024 cache rows + let tok = try Tokenizer(directory: dir) + let text = (0..<120).map { "Line \($0) of a long document about rings, windows and caches." }.joined(separator: " ") + let ids = Array(tok.encode(text).map { Int32($0) }.prefix(1200)) + #expect(ids.count == 1200) + let V = model.vocabSize + let pre = Array(UnsafeBufferPointer(start: try model.forward(ids, offset: 0), count: V)) + _ = try model.forward(Array(ids[0..<1100]), offset: 0) + var inc: [Float] = [] + for i in 1100..= 8 ? 16 : 32 + let out = try GPU.shared.buffer(nH * D * 4) + try runGPU { enc in + enc.setComputePipelineState(try GPU.shared.pipeline("attention_\(e)_\(sg)")) + enc.set(makeBuffer(q), 0, 0); enc.set(makeBuffer(kc), 0, 1); enc.set(makeBuffer(vc), 0, 2); enc.set(out, 0, 3); enc.params(p, 4) + enc.dispatchThreadgroups(MTLSize(width: nH, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: sg * 32, height: 1, depth: 1)) + } + assertClose(readFloats(out, nH * D), ref, "ring decode D=\(D)") + // Split decode over the same ring. + let G = nH / nKV, gSub = e <= 4 ? G : min(G, 2) + let V = Transformer.attnSplitVirtualSG, NSG = Transformer.attnSplitSG + let part = try GPU.shared.buffer(nH * V * (D + 2) * 4), out2 = try GPU.shared.buffer(nH * D * 4) + try runGPU { enc in + enc.setComputePipelineState(try GPU.shared.pipeline("attention_split_\(e)_\(gSub)")) + enc.set(makeBuffer(q), 0, 0); enc.set(makeBuffer(kc), 0, 1); enc.set(makeBuffer(vc), 0, 2); enc.set(part, 0, 3); enc.params(p, 4) + enc.dispatchThreadgroups(MTLSize(width: nKV, height: V / NSG, depth: G / gSub), threadsPerThreadgroup: MTLSize(width: NSG * 32, height: 1, depth: 1)) + enc.setComputePipelineState(try GPU.shared.pipeline("attention_combine_\(e)")) + enc.set(part, 0, 0); enc.set(out2, 0, 1); enc.params(p, 2) + enc.dispatchThreadgroups(MTLSize(width: nH, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: 32, height: 1, depth: 1)) + } + assertClose(readFloats(out2, nH * D), ref, "ring split decode D=\(D)") + } + @Test(arguments: cases) func kernelsMatchReference(c: Case) throws { let D = c.D, nH = c.nH, nKV = c.nKV, maxSeq = 128, e = D / 32 diff --git a/site/index.html b/site/index.html index c179458..ae9314a 100644 --- a/site/index.html +++ b/site/index.html @@ -59,7 +59,7 @@

Free MIT · 100% Swift · zero dependencies

LLMs at the speed
of unified memory.

-

A native Mac app for MLX models from Hugging Face, with an inference engine written from scratch in Swift and Metal. No Python, no packages, no copy of your weights.

+

A native Mac app for MLX models from Hugging Face, with an inference engine written from scratch in Swift and Metal. No Python, no packages, nothing between unified memory and the GPU.

@@ -121,12 +121,14 @@

Same prompt. Five models. Real speeds.

  • Llama 3.2 1B 695 MB1.69 s
  • Phi-3.5 mini 2.15 GB5.26 s
  • Mistral 7B v0.3 4.08 GB9.33 s
    • +
  • Gemma 3 1B 732 MB1.67 s
    • +
  • Gemma 3 4B 2.56 GB5.30 s
    • - + @@ -134,6 +136,8 @@

    Same prompt. Five models. Real speeds.

    + +
    mlxr bench, MLX Runner 1.1.1, M2 Pro (16 GB), 4-bit models.mlxr bench, MLX Runner 1.2.1, M2 Pro (16 GB), 4-bit models.
    ModelWeightsPrefill (512 tok)Decode @16Decode @1K ctx
    Qwen2.5 0.5B Instruct278 MB2,960 tok/s358 tok/s295 tok/s
    Llama 3.2 1B Instruct695 MB1,180 tok/s203 tok/s188 tok/s
    Phi-3.5 mini2.15 GB320 tok/s70 tok/s61 tok/s
    Mistral 7B v0.34.08 GB166 tok/s41.0 tok/s39.5 tok/s
    Gemma 3 1B (QAT)732 MB1,450 tok/s195 tok/s178 tok/s
    Gemma 3 4B2.56 GB345 tok/s67 tok/s62 tok/s
    From b6507d223036ffc034c390375885e0c8274592bf Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Sun, 1 Nov 2026 19:01:56 +0100 Subject: [PATCH 2/3] Review fixes for the ring cache and loader The model tracks its own KV write high-water mark, so prefix reuse accounts for writes the engine's bookkeeping dropped (a cancelled prefill, rejected speculative drafts, a failed in-flight step); the engine resets it when a run starts from position 0. The ring size is derived from the one maxBatch constant and checked at load. GPU out-of-memory is recognised by the MTLCommandBufferError code, with the text match as fallback. The load-time scratch estimate is Transformer.scratchBytes(config:), and a test checks it covers what init allocates on every installed checkpoint. flash_prefill computes the ring slot per row. Prefill and decode A/B (3 alternating pairs): Mistral 7B and Gemma 3 1B within 1% of the previous build; prefill unchanged everywhere. Co-Authored-By: Claude Fable 5.1 --- Sources/MLXEngine/Engine.swift | 7 +-- Sources/MLXEngine/Kernels.swift | 6 +- Sources/MLXEngine/ModelConfig.swift | 6 +- Sources/MLXEngine/Transformer.swift | 81 +++++++++++++++++++++------ Tests/MLXEngineTests/GemmaTests.swift | 16 ++++++ 5 files changed, 90 insertions(+), 26 deletions(-) diff --git a/Sources/MLXEngine/Engine.swift b/Sources/MLXEngine/Engine.swift index 386e25d..f93aea2 100644 --- a/Sources/MLXEngine/Engine.swift +++ b/Sources/MLXEngine/Engine.swift @@ -399,6 +399,7 @@ public final class InferenceEngine: @unchecked Sendable { 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 @@ -593,11 +594,7 @@ extension InferenceEngine { let cfg = try ModelConfig(directory: directory) let ctx = min(contextLength, max(cfg.maxPositionEmbeddings, 256)) let kv = cfg.kvCacheBytes(context: ctx) - // Scratch buffers for a 512-token batch (hidden, projections, MLP, logits), see Transformer.init. - let H = cfg.hiddenSize, I = cfg.intermediateSize * (cfg.isGemma4 ? 2 : 1) - let qkv = cfg.layers.map { (cfg.numHeads + 2 * $0.numKVHeads) * $0.headDim }.max() ?? 0 - let attn = cfg.layers.map { cfg.numHeads * $0.headDim }.max() ?? 0 - let scratch = 512 * 4 * (4 * H + 3 * I + qkv + attn) + 6 * cfg.vocabSize * 4 + let scratch = Transformer.scratchBytes(config: cfg) let need = UInt64(weights + kv + scratch) let budget = GPU.shared.device.recommendedMaxWorkingSetSize guard need < budget else { diff --git a/Sources/MLXEngine/Kernels.swift b/Sources/MLXEngine/Kernels.swift index 3ea2d23..c767161 100644 --- a/Sources/MLXEngine/Kernels.swift +++ b/Sources/MLXEngine/Kernels.swift @@ -1126,10 +1126,12 @@ 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; - const ulong src = ulong(kv_slot(j, p.maxSeq)) * D + i % D; + 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; } diff --git a/Sources/MLXEngine/ModelConfig.swift b/Sources/MLXEngine/ModelConfig.swift index f40320c..30bfe62 100644 --- a/Sources/MLXEngine/ModelConfig.swift +++ b/Sources/MLXEngine/ModelConfig.swift @@ -97,10 +97,14 @@ public struct ModelConfig { public var perLayerInputSize = 0 public var valueNorm = false + /// 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 + 512) : context + spec.sliding && window > 0 ? min(context, window + maxBatch) : context } func cacheRows(_ spec: LayerSpec, context: Int) -> Int { Self.cacheRows(spec, context: context, window: slidingWindow) } diff --git a/Sources/MLXEngine/Transformer.swift b/Sources/MLXEngine/Transformer.swift index acda52f..eac569b 100644 --- a/Sources/MLXEngine/Transformer.swift +++ b/Sources/MLXEngine/Transformer.swift @@ -372,6 +372,10 @@ final class TransformerLayer { gateUpFusable = gateUp.count > 1 && Linear.fusable(gateUp) qkvPostPSO = try GPU.shared.pipeline("qkv_post_\(qNorm?.weight.dtype.suffix ?? "f32")") cacheLen = c.cacheRows(spec, context: maxSeq) + // A batch's new rows must never land on rows its own queries read (see cacheRows). + guard cacheLen == maxSeq || cacheLen >= c.slidingWindow + ModelConfig.maxBatch else { + throw EngineError.invalidFile("\(p): KV cache of \(cacheLen) rows for a window of \(c.slidingWindow)") + } if spec.hasOwnKV { let cacheBytes = nKV * cacheLen * D * 2 kCache = try GPU.shared.buffer(cacheBytes) @@ -382,12 +386,34 @@ final class TransformerLayer { } } +extension Transformer { + /// Upper bound on the scratch buffers `init` allocates for a config (the load-time memory + /// check adds it to the weights and the KV caches). A test keeps it above `scratchAllocated`. + public static func scratchBytes(config c: ModelConfig) -> Int { + let B = maxBatch, H = c.hiddenSize + let I = c.intermediateSize * (c.isGemma4 ? 2 : 1) // Gemma 4's KV-shared layers double the MLP + let qkv = c.layers.map { (c.numHeads + 2 * $0.numKVHeads) * $0.headDim }.max() ?? 0 + let attn = c.layers.map { c.numHeads * $0.headDim }.max() ?? 0 + let maxD = c.layers.map(\.headDim).max() ?? 0 + var bytes = B * 4 * (3 * H + qkv + attn + 3 * I) // h, xn, tmp, qkv, attn, gate|up, mlp + bytes += (1 + maxVerifyTokens) * max(c.vocabSize, 1) * 4 // logits, logitsAll + bytes += c.numHeads * attnSplitVirtualSG * (maxD + 2) * 4 // split-attention partials + if c.perLayerInputSize > 0 { + let LP = c.perLayerInputSize * c.numLayers + bytes += B * 4 * (2 * LP + c.perLayerInputSize) // per-layer inputs: raw, inputs, gate + } + return bytes + } +} + extension EngineError { /// The engine error for a failed command buffer: an out-of-memory failure gets an /// actionable message, anything else is reported as a GPU error. static func gpu(_ err: Error) -> EngineError { + let ns = err as NSError let d = err.localizedDescription - if d.contains("OutOfMemory") || d.lowercased().contains("insufficient memory") { + if (ns.domain == MTLCommandBufferErrorDomain && ns.code == MTLCommandBufferError.outOfMemory.rawValue) + || d.contains("OutOfMemory") || d.lowercased().contains("insufficient memory") { return .outOfMemory("the GPU ran out of memory while running the model. Lower the context length (Settings) or use a smaller model.") } return .kernelCompile("GPU error: \(d)") @@ -406,8 +432,14 @@ public final class Transformer { let quantDescription: String // Scratch buffers sized for `maxBatch` tokens. - static let maxBatch = 512 + static let maxBatch = ModelConfig.maxBatch let maxBatch = Transformer.maxBatch + /// One past the highest position any committed step has written to the KV caches. Ring + /// caches alias positions, so a resumed run must know about every write, including those + /// of a cancelled prefill or a rejected draft that the caller's own bookkeeping dropped. + public private(set) var writtenHigh = 0 + /// Scratch bytes `init` actually allocated (checked against `scratchBytes(config:)` by a test). + private(set) var scratchAllocated = 0 /// Token ids by position (maxSeq + 1 entries): the CPU writes prompt tokens, the GPU sampler /// appends generated ones, and every step's embedding lookup reads its tokens from here. private let hist: MTLBuffer @@ -535,15 +567,17 @@ public final class Transformer { sampleDiag = try gpu.buffer(16) samplePSO = try gpu.pipeline("sample_token") sampleThreads = min(1024, samplePSO.maxTotalThreadsPerThreadgroup) / 32 * 32 - h = try gpu.buffer(B * H * 4) - xn = try gpu.buffer(B * H * 4) - qkvBuf = try gpu.buffer(B * qkvMax * 4) - attnBuf = try gpu.buffer(B * attnMax * 4) - guBuf = try gpu.buffer(B * 2 * I * 4) - mlpBuf = try gpu.buffer(B * I * 4) - tmp = try gpu.buffer(B * H * 4) - logits = try gpu.buffer(lmHead.outDim * 4) - logitsAll = try gpu.buffer(Self.maxVerifyTokens * lmHead.outDim * 4) + var scratch = 0 + func buffer(_ n: Int) throws -> MTLBuffer { scratch += n; return try gpu.buffer(n) } + h = try buffer(B * H * 4) + xn = try buffer(B * H * 4) + qkvBuf = try buffer(B * qkvMax * 4) + attnBuf = try buffer(B * attnMax * 4) + guBuf = try buffer(B * 2 * I * 4) + mlpBuf = try buffer(B * I * 4) + tmp = try buffer(B * H * 4) + logits = try buffer(lmHead.outDim * 4) + logitsAll = try buffer(Self.maxVerifyTokens * lmHead.outDim * 4) var freqBuffers: [[Float]: MTLBuffer] = [:] ropeFreqs = c.layers.map { l in if let b = freqBuffers[l.invFreq] { return b } @@ -598,7 +632,7 @@ public final class Transformer { } kernels[D] = k } - attnPartials = partialBytes > 0 ? try gpu.buffer(partialBytes) : nil + attnPartials = partialBytes > 0 ? try buffer(partialBytes) : nil attnKernels = kernels if c.perLayerInputSize > 0 { @@ -612,11 +646,12 @@ public final class Transformer { modelProj: try Linear(store, "model.per_layer_model_projection", inDim: H, outDim: LP), // The reference normalises proj(h) / sqrt(H): RMSNorm(x / sqrt(H), eps) == RMSNorm(x, eps · H). norm: try RMSNorm(store, "model.per_layer_projection_norm.weight", eps: c.rmsNormEps * Float(H), dim: P), - raw: try gpu.buffer(B * LP * 4), inputs: try gpu.buffer(B * LP * 4), gateBuf: try gpu.buffer(B * P * 4), + raw: try buffer(B * LP * 4), inputs: try buffer(B * LP * 4), gateBuf: try buffer(B * P * 4), addPSO: try gpu.pipeline("add_scaled"), gatePSO: try gpu.pipeline("gelu_tanh_gate")) } else { ple = nil } + scratchAllocated = scratch switch c.activation { case .silu: siluPSO = try gpu.pipeline("silu_mul") case .geluTanh: siluPSO = try gpu.pipeline("gelu_tanh_mul") @@ -662,6 +697,7 @@ public final class Transformer { // Nothing is in flight (the previous step was waited for), so `hist` is free to write. writeHistory(tokens, at: offset) cb.commit() + writtenHigh = max(writtenHigh, offset + T) // Overlap CPU encoding of the likely next decode step (the following position) with this // step's GPU work. It is only committed after this buffer completes, on an exact match. if computeLogits, offset + T < maxSeq, let next = try? encode(T: 1, offset: offset + T, computeLogits: true, allLogits: false) { @@ -688,17 +724,24 @@ public final class Transformer { public let offset: Int } - /// How much of a previous run's KV cache (positions 0.. Int { + let written = max(previousLength, writtenHigh) for l in layers where l.spec.hasOwnKV && l.cacheLen < maxSeq { - if previousLength - common > l.cacheLen - config.slidingWindow + 1 { return 0 } + if written - common > l.cacheLen - config.slidingWindow + 1 { return 0 } } return common } + /// Forget the write history: the next run starts at position 0 and overwrites everything + /// it will read. + public func resetWriteTracking() { writtenHigh = 0 } + /// Seeds the GPU sampler's xorshift64 generator. Call before `prefill` (nothing may be in flight). public func seedSampler(_ seed: UInt64) { rngState.contents().storeBytes(of: seed == 0 ? 0x9E3779B97F4A7C15 : seed, as: UInt64.self) @@ -724,6 +767,7 @@ public final class Transformer { let cb = try encode(T: n, offset: at, computeLogits: last, allLogits: false, sample: last ? sample : nil) writeHistory(Array(tokens[start.. Date: Sun, 1 Nov 2026 19:09:30 +0100 Subject: [PATCH 3/3] One buffer per tensor instead of an arena A single arena for all of a shard's tensors measured 6% slower at 1K context on Qwen3 0.6B (3 alternating pairs, repeatedly); separate page-aligned buffers, the layout the engine always had, are within 1%. The mapping is still released after the copy, so the GPU holds one copy. Co-Authored-By: Claude Fable 5.1 --- Sources/MLXEngine/GPU.swift | 52 +++++++++++++++++-------------------- 1 file changed, 24 insertions(+), 28 deletions(-) diff --git a/Sources/MLXEngine/GPU.swift b/Sources/MLXEngine/GPU.swift index 2b6139e..e1b5788 100644 --- a/Sources/MLXEngine/GPU.swift +++ b/Sources/MLXEngine/GPU.swift @@ -181,17 +181,22 @@ extension MTLComputeCommandEncoder { // MARK: - Safetensors -/// One safetensors shard as a Metal buffer. The kernels use 16-byte vector loads, so every +/// 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, into one 16-byte-padded arena and the mapping is released: the -/// GPU then holds exactly one copy of the weights. (Copying each tensor into its own buffer -/// while the mapping stayed referenced doubled the working set — a 9 GB model needed 18 GB.) +/// 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 @@ -258,35 +263,26 @@ final class SafeTensorFile { munmap(p, len) }) else { throw EngineError.invalidFile("cannot wrap \(file) in a Metal buffer") } buffer = buf + byteCount = size mappingOwned = false return } - // Arena: tensors in file order, each on a page boundary (what separate buffers gave; a - // 16-byte packing measured a few percent slower on the smallest model). - let align = page - let ordered = infos.sorted { $0.value.start < $1.value.start } - var total = 0 - var placed: [String: Int] = [:] - for (name, i) in ordered { - placed[name] = total - let (next, o) = total.addingReportingOverflow((i.end - i.start) + align - 1) - guard !o else { throw EngineError.invalidFile("\(file) is too large") } - total = next / align * align - } - guard let arena = device.makeBuffer(length: max(total, 16), options: .storageModeShared) else { - throw EngineError.outOfMemory("could not allocate \(ByteCountFormatter.string(fromByteCount: Int64(total), countStyle: .file)) for \(file)") - } - let dst = arena.contents() - for (name, i) in ordered { - memcpy(dst + placed[name]!, base + i.start, i.end - i.start) - infos[name] = (i.dtype, i.shape, placed[name]!, placed[name]! + (i.end - i.start)) + // 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 + 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) } - buffer = arena } func tensor(_ name: String) -> Tensor? { - guard let i = infos[name] else { return nil } - return Tensor(buffer: buffer, offset: i.start, dtype: i.dtype, shape: i.shape) + 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) } } @@ -347,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 } } }