Merge pull request #173 from SharpAI/fix/fullgpu-cache-limit

fix(memory): run cleanly on 32 GB Macs + Mac mini M6 benchmarks

Author: solderzzc

README.md

@@ -51,6 +51,90 @@ Then start the server (models download automatically if not cached):
 
 *(Add `--stream-experts` when running oversized MoE models to bypass macOS virtual memory swapping and stream expert layers directly from NVMe SSD.)*
 
+## 📊 Performance: Mac mini M6 (base, 32 GB)
+
+The first SwiftLM numbers from a **32 GB** Mac. Every other table in this README comes from 64 GB hardware. Running the same context lengths on this M6 exposed memory bugs that 64 GB machines had been hiding, and they are fixed in this release (see [What 32 GB exposed](#what-32-gb-exposed)).
+
+![Gemma-4-26B-A4B streaming at 55 tok/s on a base Mac mini M6 32 GB](docs/profiling/m6/media/m6_gemma4_26b_a4b_stream.gif)
+
+> *Hardware:* Mac mini (Mac18,5), Apple M6, 12-core GPU, 32 GB unified memory (170 GB/s), macOS 27.0. Metal working set 26.8 GB.
+> *Method:* [`scripts/profiling/m6_bench.py`](scripts/profiling/m6_bench.py). One warm-up, then the median of 3 runs (1 run at 32K and above), temperature 0. Every prompt starts with a unique nonce, so the prompt cache can't hit, and hides a code word that the answer must return. A memory guard aborts any case whose swap grows by more than 2 GB. Raw results: [`docs/profiling/m6/`](docs/profiling/m6/).
+
+### What runs well on a 32 GB M6
+
+| Model (4-bit unless noted) | Weights | Mode | Decode, short prompt | Longest prompt that passed | Peak GPU |
+|---|---|---|---|---|---|
+| **`gemma-4-26b-a4b-it-4bit`** (MoE, ~4B active) | 15.3 GB | GPU | **52.2 tok/s** | 80.7K tokens | 19.5 GB |
+| **`Qwen3.6-35B-A3B-UD-MLX-4bit`** (MoE, ~3B active) | 21.6 GB | GPU | **46.7 tok/s** | 40.8K tokens | 22.4 GB |
+| `Qwen3.6-35B-A3B-UD-MLX-4bit` | 21.6 GB | `--stream-experts` | 13.2 tok/s | 40.8K tokens | 7.7 GB |
+| `Qwen3.8-27B-4bit` (dense) | 11.3 GB | GPU | 9.3 tok/s | 40.8K tokens | 18.4 GB |
+| `gemma-4-26b-a4b-it-8bit` | ~26 GB | GPU | swaps (+3.1 GB on the first prompt) | — | — |
+| `gemma-4-26b-a4b-it-8bit` | ~26 GB | `--stream-experts` | 8.8 tok/s | 9.5K tokens (32K swapped) | 7.6 GB |
+
+- **MoE models are the sweet spot at 32 GB.** Only the active experts are read for each token, so they decode 5–6× faster than a dense 27B. A 4-bit MoE with up to about 22 GB of weights runs entirely on the GPU.
+- **Qwen3.6-35B-A3B on a base M6 reaches 76%** of the M1 Ultra 64 GB decode speed below (46.7 vs 61.7 tok/s).
+- **Dense 27B decode is bandwidth-bound.** 9.3 tok/s × 11.3 GB is about 105 GB/s, roughly 60% of the M6's rated 170 GB/s.
+- **An 8-bit 26 GB model needs SSD streaming** and tops out at about 10K tokens of context.
+
+### Gemma-4-26B-A4B 4-bit — by prompt length
+
+| Prompt tokens | Vanilla prefill / decode (tok/s) | TTFT | `--turbo-kv` prefill / decode (tok/s) | Peak GPU · swap growth |
+|---|---|---|---|---|
+| ~530 | 733 / **52.2** | 0.8 s | 785 / 53.2 | 14.5 GB · 0 |
+| ~2.3K | **963** / 50.2 | 2.5 s | 969 / 50.5 | 15.0 GB · 0 |
+| ~9.5K | 959 / 45.1 | 10.1 s | 971 / 45.5 | 15.8 GB · 0 |
+| ~39.7K | 757 / 31.0 | 53.3 s | 787 / 31.5 | 17.9 GB · 0 |
+| ~80.7K | 622 / 24.3 | 131.3 s | 630 / 24.5 | 19.5 GB · 0 |
+
+Every needle check passed in both modes. `--mtp` with the bf16 assistant (`gemma-4-26B-A4B-it-assistant-bf16`) works but is slower on the M6: 45.2 / 35.6 / 30.4 tok/s decode at ~530 / 2.3K / 9.5K tokens, against 53.0 / 50.8 / 46.0 without it. A 4-bit MoE is compute-bound, so verifying the drafted tokens costs more than it saves (the same finding as the M5 Pro tables below).
+
+### Qwen3.6-35B-A3B 4-bit — GPU vs SSD streaming
+
+| Prompt tokens | GPU prefill / decode (tok/s) | GPU peak | `--stream-experts` prefill / decode (tok/s) | SSD peak |
+|---|---|---|---|---|
+| ~550 | 808 / 46.7 | 20.4 GB | 321 / 13.2 | 6.0 GB |
+| ~2.3K | 969 / 45.5 | 21.0 GB | 402 / 13.1 | 6.2 GB |
+| ~9.8K | 849 / 43.4 | 21.1 GB | 403 / 12.9 | 6.6 GB |
+| 40.8K | 635 / 35.1 | 22.4 GB | 340 / 11.9 | 7.7 GB |
+
+### Qwen3.8-27B-4bit (dense) — Vanilla vs TurboKV
+
+| Prompt tokens | Vanilla prefill / decode (tok/s) | TurboKV prefill / decode (tok/s) | Peak GPU · swap growth |
+|---|---|---|---|
+| ~550 | 233 / 9.3 | 231 / 9.3 | 15.2 GB · 0 |
+| ~2.3K | 242 / 9.1 | 261 / 9.2 | 16.1 GB · 0 |
+| ~9.8K | 238 / 8.9 | 250 / 8.9 | 16.9 GB · 0 |
+| ~40.8K | 200 / 7.9 | 202 / 7.9 | 18.4 GB · 0 (TurboKV 16.9 GB) |
+
+TurboKV barely changes speed on this model. Only 16 of its 64 layers use full attention (the other 48 are GatedDeltaNet), so the KV cache is already small. It saves about 1.5 GB at 40K tokens.
+
+### What 32 GB exposed
+
+| 8.5K-token prompt, Qwen3.8-27B-4bit | Before (old pin) | After (this release) |
+|---|---|---|
+| Prefill | 33.4 tok/s | **~240 tok/s** (≈7×; 238 tok/s measured at 9.8K) |
+| Peak memory | 38 GB process footprint | **≤19 GB** process (16.9 GB GPU peak at 9.8K) |
+| Swap growth | +15 GB | **0** |
+
+1. **The MLX buffer cache was unbounded on full-GPU loads.** It could grow to the whole 26.8 GB working set. It is now sized from the RAM left after weights and KV.
+2. **The KV-cache estimate counted every layer as full attention.** Gemma 4 (25 of 30 layers use a 1,024-token sliding window) was overestimated 10×, and Qwen3.5/3.8 (48 of 64 layers are linear attention) 4×. On 32 GB that pushed Gemma into CPU/GPU layer partitioning, which crashed with a Metal GPU timeout.
+3. **An auto-detected VLM that failed to load exited the server.** `Qwen3.6-35B-A3B-UD-MLX-4bit` ships a `preprocessor_config.json` without `image_mean`. SwiftLM now falls back to text-only unless you pass `--vision`.
+4. **Vision-capable models skipped chunked prefill.** On the older mlx-swift-lm pin, a text-only prompt on the VLM path ran through the model in a single pass. It's fixed by the mlx-swift-lm bump in #167. Every number in this section was measured on `main` with that bump.
+
+> ⚠️ **Known issues:** `--gpu-layers N` (CPU/GPU layer partitioning) hits a Metal GPU timeout on the first request (repro: `--model mlx-community/gemma-4-26b-a4b-it-4bit --gpu-layers 23`). QAT-quantized Gemma 4 MTP assistants (`…-qat-assistant-4bit`) fail with `unhandledKeys pre_projection/post_projection`; use `gemma-4-26B-A4B-it-assistant-bf16`.
+
+Reproduce:
+
+```bash
+./build.sh
+.build/release/SwiftLM --model mlx-community/gemma-4-26b-a4b-it-4bit --port 5431 --ctx-size 48000 &
+python3 scripts/demo/stream_client.py short
+python3 scripts/profiling/m6_bench.py --model mlx-community/gemma-4-26b-a4b-it-4bit \
+  --config "Vanilla=" --contexts 512,2048,8192,32768,65536 --out docs/profiling/m6/gemma4_26b_a4b_4bit
+```
+
+More recordings: [Gemma, 41K-token prompt (4×)](docs/profiling/m6/media/m6_gemma4_26b_a4b_41k_prompt_4x.gif) · [Qwen3.8-27B streaming](docs/profiling/m6/media/m6_qwen38_27b_stream.gif) · [Qwen3.8-27B, 8.6K-token prompt (4×)](docs/profiling/m6/media/m6_qwen38_27b_8k_prompt_4x.gif)
+
 ## 📊 Performance: MTP Speculative Decoding — Gemma 4-26B (MacBook Pro M5 Pro 64 GB)
 
 Benchmarked with `gemma-4-26b-a4b-it-4bit` running three configurations across 512 / 40K / 100K token contexts.

Sources/SwiftLM/ModelProfiler.swift

@@ -31,6 +31,14 @@ struct ModelProfile: Sendable {
     let numActiveExperts: Int?
     let weightFileSizeBytes: Int
     let modelId: String
+    /// Attention layout from `layer_types`. When a config has no `layer_types`,
+    /// every layer counts as full attention, which is the old estimate.
+    var fullAttentionLayers: Int? = nil
+    var slidingAttentionLayers: Int = 0
+    var slidingWindow: Int? = nil
+    /// Gemma 4 gives its full-attention layers their own KV head count and head dim.
+    var globalKVHeads: Int? = nil
+    var globalHeadDim: Int? = nil
 
     /// Estimated total parameters in billions (rough)
     var estimatedParamsB: Double {
@@ -55,10 +63,19 @@ struct ModelProfile: Sendable {
 
     /// KV cache memory in GB for a given context length
     func kvCacheMemoryGB(contextLength: Int) -> Double {
-        // KV cache = 2 (K + V) × layers × kv_heads × head_dim × context × 2 bytes (FP16)
+        // KV cache = 2 (K + V) × layers × kv_heads × head_dim × tokens × 2 bytes (FP16).
+        // Only attention layers hold a KV cache. Linear-attention layers (GatedDeltaNet)
+        // keep a fixed-size state instead, and sliding-window layers stop growing at the
+        // window size. Counting every layer as full attention overestimated Gemma 4 by
+        // ~10× and Qwen3.5/3.8 by ~4×, which pushed 32 GB machines into CPU
+        // partitioning they didn't need.
         let bytesPerElement = 2 // FP16
-        let kvBytes = 2 * numLayers * numKVHeads * headDim * contextLength * bytesPerElement
-        return Double(kvBytes) / 1e9
+        let fullLayers = fullAttentionLayers ?? numLayers
+        let fullBytes = 2 * fullLayers * (globalKVHeads ?? numKVHeads) * (globalHeadDim ?? headDim)
+            * contextLength * bytesPerElement
+        let slidingTokens = min(contextLength, slidingWindow ?? contextLength)
+        let slidingBytes = 2 * slidingAttentionLayers * numKVHeads * headDim * slidingTokens * bytesPerElement
+        return Double(fullBytes + slidingBytes) / 1e9
     }
 
     /// Total memory required in GB (weights + KV cache + overhead)
@@ -177,6 +194,10 @@ enum ModelProfiler {
         let vocabSize: Int?
         let quantizationConfig: QuantConfig?
         let textConfig: TextConfig?
+        let layerTypes: [String]?
+        let slidingWindow: Int?
+        let numGlobalKeyValueHeads: Int?
+        let globalHeadDim: Int?
 
         enum CodingKeys: String, CodingKey {
             case modelType = "model_type"
@@ -189,6 +210,10 @@ enum ModelProfiler {
             case vocabSize = "vocab_size"
             case quantizationConfig = "quantization_config"
             case textConfig = "text_config"
+            case layerTypes = "layer_types"
+            case slidingWindow = "sliding_window"
+            case numGlobalKeyValueHeads = "num_global_key_value_heads"
+            case globalHeadDim = "global_head_dim"
         }
     }
 
@@ -200,6 +225,10 @@ enum ModelProfiler {
         let headDim: Int?
         let intermediateSize: Int?
         let vocabSize: Int?
+        let layerTypes: [String]?
+        let slidingWindow: Int?
+        let numGlobalKeyValueHeads: Int?
+        let globalHeadDim: Int?
 
         enum CodingKeys: String, CodingKey {
             case numHiddenLayers = "num_hidden_layers"
@@ -209,6 +238,10 @@ enum ModelProfiler {
             case headDim = "head_dim"
             case intermediateSize = "intermediate_size"
             case vocabSize = "vocab_size"
+            case layerTypes = "layer_types"
+            case slidingWindow = "sliding_window"
+            case numGlobalKeyValueHeads = "num_global_key_value_heads"
+            case globalHeadDim = "global_head_dim"
         }
     }
 
@@ -275,7 +308,8 @@ enum ModelProfiler {
         // Measure weight file sizes on disk (only for MoE to avoid slow walks on dense models)
         let weightSize = isMoE ? measureWeightFiles(directory: modelDirectory) : 0
 
-        return ModelProfile(
+        let layerTypes = config.layerTypes ?? config.textConfig?.layerTypes
+        var profile = ModelProfile(
             modelType: modelType,
             numLayers: numLayers,
             hiddenSize: hiddenSize,
@@ -291,6 +325,22 @@ enum ModelProfiler {
             weightFileSizeBytes: weightSize,
             modelId: modelId
         )
+        if let layerTypes, layerTypes.count == numLayers {
+            // Anything that isn't sliding or linear attention is treated as full attention.
+            let sliding = layerTypes.filter { $0 == "sliding_attention" }.count
+            let linear = layerTypes.filter { $0 == "linear_attention" }.count
+            profile.slidingAttentionLayers = sliding
+            profile.fullAttentionLayers = layerTypes.count - sliding - linear
+            profile.slidingWindow = config.slidingWindow ?? config.textConfig?.slidingWindow
+            if sliding > 0 && profile.slidingWindow == nil {
+                // Can't cap a window we don't know, so count those layers as full.
+                profile.fullAttentionLayers = layerTypes.count - linear
+                profile.slidingAttentionLayers = 0
+            }
+        }
+        profile.globalKVHeads = config.numGlobalKeyValueHeads ?? config.textConfig?.numGlobalKeyValueHeads
+        profile.globalHeadDim = config.globalHeadDim ?? config.textConfig?.globalHeadDim
+        return profile
     }
 
     /// Routed-expert count keys, in precedence order within a container.
@@ -505,7 +555,15 @@ enum ModelProfiler {
         switch strategy {
         case .fullGPU:
             memoryLimit = Int(Double(system.recommendedWorkingSetBytes) * 1.5)
-            cacheLimit = system.recommendedWorkingSetBytes // default
+            // MLX's default cache limit is the whole working set. Prefill chunk
+            // buffers change shape as the KV offset grows, so they are rarely
+            // reused and the cache just keeps growing. On a 32 GB machine that
+            // pushed macOS into swap by ~8K tokens. Give the cache half of what
+            // is left after weights, KV and an 8 GB OS reserve, and never more
+            // than the old default.
+            let headroomGB = system.totalRAMGB - (weightGB + draftGB) - kvGB - 8.0
+            let budget = Int(max(1.0, headroomGB / 2) * 1e9)
+            cacheLimit = min(system.recommendedWorkingSetBytes, budget)
         case .swapAssisted:
             memoryLimit = 200 * 1024 * 1024 * 1024 // 200 GB sentinel to bypass MLX eval_impl spin loop (let macOS swap handle it)
             cacheLimit = 2 * 1024 * 1024 // 2MB — let OS manage caching

Sources/SwiftLM/Server.swift

@@ -205,6 +205,28 @@ private struct TransformersTokenizerBridge: MLXLMCommon.Tokenizer, Sendable {
 
 /// Returns `nil` when the value must be dropped (JSON `null` / NSNull), otherwise a
 /// structure with every nested null removed. See `TransformersTokenizerBridge.applyChatTemplate`.
+/// True when a VLM load failed because the checkpoint doesn't match the VLM code:
+/// its config doesn't decode, its weights don't line up with the module tree, or the
+/// factory doesn't know the model/processor type. Only these justify retrying an
+/// auto-detected VLM as a text-only LLM. Anything else (cancellation, download, I/O)
+/// would fail the same way again and just hide the real error.
+func isVLMCheckpointMismatch(_ error: any Error) -> Bool {
+    switch error {
+    case is DecodingError, is UpdateError:
+        return true
+    case let factoryError as ModelFactoryError:
+        switch factoryError {
+        case .unsupportedModelType, .unsupportedProcessorType, .configurationDecodingError,
+            .invalidConfiguration:
+            return true
+        default:
+            return false
+        }
+    default:
+        return false
+    }
+}
+
 func sanitizeForJinja(_ value: any Sendable) -> (any Sendable)? {
     if value is NSNull { return nil }
     let mirror = Mirror(reflecting: value)
@@ -771,7 +793,8 @@ struct MLXServer: AsyncParsableCommand {
             // Apply memory strategy
             switch plan.strategy {
             case .fullGPU:
-                print("[SwiftLM] \(plan.strategy.emoji) Memory strategy: FULL GPU (\(String(format: "%.1f", plan.weightMemoryGB))GB model, \(String(format: "%.1f", system.availableRAMGB))GB available)")
+                Memory.cacheLimit = plan.recommendedCacheLimit
+                print("[SwiftLM] \(plan.strategy.emoji) Memory strategy: FULL GPU (\(String(format: "%.1f", plan.weightMemoryGB))GB model, \(String(format: "%.1f", system.availableRAMGB))GB available, cache limited to \(plan.recommendedCacheLimit / (1024*1024))MB)")
             case .swapAssisted:
                 if self.streamExperts {
                     // SSD Streaming: expert weights are mmap'd from SSD via the OS page cache.
@@ -863,7 +886,7 @@ struct MLXServer: AsyncParsableCommand {
         let speculativeDecodingRequested = self.draftModel != nil || self.dflash || self.mtp
         let autoDetectedVision = !self.audio && architecture.supportsVision
             && !speculativeDecodingRequested
-        let isVision = self.vision || autoDetectedVision
+        var isVision = self.vision || autoDetectedVision
         if architecture.supportsVision, !self.vision, !self.audio, speculativeDecodingRequested {
             print(
                 "[SwiftLM] Note: \(architecture.modelType ?? "unknown") reports vision support, but speculative/MTP decoding was requested; loading as a text-only LLM."
@@ -895,12 +918,29 @@ struct MLXServer: AsyncParsableCommand {
             }
         } else if isVision {
             print("[SwiftLM] Loading VLM (vision-language model)...")
-            container = try await VLMModelFactory.shared.loadContainer(
-                from: downloader,
-                using: TransformersTokenizerLoader(modelId: resolvedModelId),
-                configuration: modelConfig
-            ) { progress in
-                tracker.printProgress(progress)
+            do {
+                container = try await VLMModelFactory.shared.loadContainer(
+                    from: downloader,
+                    using: TransformersTokenizerLoader(modelId: resolvedModelId),
+                    configuration: modelConfig
+                ) { progress in
+                    tracker.printProgress(progress)
+                }
+            } catch where !self.vision && isVLMCheckpointMismatch(error) {
+                // Vision was only auto-detected, and the vision side of the checkpoint
+                // doesn't match the VLM code (e.g. a preprocessor_config.json without
+                // image_mean). The text model can still serve, so fall back rather than
+                // exit. Cancellation, network and I/O errors still propagate, as does
+                // any error under an explicit --vision.
+                print("[SwiftLM] ⚠️  Auto-detected VLM failed to load (\(error)); loading as a text-only LLM. Pass --vision to make this fatal.")
+                isVision = false
+                container = try await LLMModelFactory.shared.loadContainer(
+                    from: downloader,
+                    using: TransformersTokenizerLoader(modelId: resolvedModelId),
+                    configuration: modelConfig
+                ) { progress in
+                    tracker.printProgress(progress)
+                }
             }
         } else if isAudio {
             print("[SwiftLM] Loading ALM (audio-language model)...")