Merge pull request #182 from SharpAI/feat/gpu-layers-moe-warning

Warn when --gpu-layers puts MoE layers on the CPU; mark #176 fixed in README

Author: solderzzc

README.md

@@ -123,7 +123,9 @@ Every needle check passed. `--mtp` with the bf16 assistant (`gemma-4-26B-A4B-it-
 
 > ⚠️ **`--turbo-kv` loses exact long-range recall on every chip:** once a prompt passes the 2,048-token compression threshold, Qwen3.8-27B-4bit with `--turbo-kv` gets exact lookups wrong. Asked how many numbered lines a prompt has, it answers "1,000", "1,314" or "14" instead of 315 / 500 / 700. Reproduced on both M5 and M6; without `--turbo-kv` it answers correctly. The cause is a cache-eviction bookkeeping regression: after compression, attention only sees the recent hot window, and positions restart. A fix is in progress; tracked in [#175](https://github.com/SharpAI/SwiftLM/issues/175). Until then, avoid `--turbo-kv` when exact recall matters. Also note that `--turbo-kv` currently has no effect when `--ctx-size` is set (the attention layers use `RotatingKVCache`).
 >
-> ⚠️ **Known issues:** `--gpu-layers N` (CPU/GPU layer partitioning) hits a Metal GPU timeout on the first request, on both M5 and M6 (repro: `--model mlx-community/gemma-4-26b-a4b-it-4bit --gpu-layers 23`). Tracked in [#176](https://github.com/SharpAI/SwiftLM/issues/176); the fix is [SharpAI/mlx-swift#17](https://github.com/SharpAI/mlx-swift/pull/17). Even once it's fixed, CPU-resident MoE layers are very slow (~0.4 tok/s prefill), so on a 32 GB Mac try `--stream-experts` first (Qwen3.6-35B-A3B: 13.2 tok/s decode, 7.7 GB GPU). QAT-quantized Gemma 4 MTP assistants (`…-qat-assistant-4bit`) fail with `unhandledKeys pre_projection/post_projection`; use `gemma-4-26B-A4B-it-assistant-bf16`.
+> ⚠️ **Known issues:** QAT-quantized Gemma 4 MTP assistants (`…-qat-assistant-4bit`) fail with `unhandledKeys pre_projection/post_projection`; use `gemma-4-26B-A4B-it-assistant-bf16`.
+>
+> ℹ️ **`--gpu-layers N` with MoE models:** the Metal GPU timeout on the first request ([#176](https://github.com/SharpAI/SwiftLM/issues/176)) is fixed in #177. CPU-resident MoE layers still run on a single core and are very slow (Gemma 4 26B-A4B with `--gpu-layers 23`: ~0.4 tok/s prefill, ~3 s per decoded token on both M5 and M6), and SwiftLM now warns about this at startup. Treat `--gpu-layers` as a way to avoid running out of memory, not as a speed trade-off. On a 32 GB Mac try `--stream-experts` first (Qwen3.6-35B-A3B: 13.2 tok/s decode, 7.7 GB GPU).
 
 Reproduce:
 

Sources/SwiftLM/Server.swift

@@ -781,9 +781,11 @@ struct MLXServer: AsyncParsableCommand {
         }
 
         var partitionPlan: PartitionPlan?
+        var modelIsMoE = false
         if let modelDir = modelDirectory {
            let profile = mainModelProfile ?? ModelProfiler.profile(modelDirectory: modelDir, modelId: modelId)
            if let profile = profile {
+            modelIsMoE = profile.isMoE
             let system = ModelProfiler.systemProfile()
             let contextSize = self.ctxSize ?? 4096
             let plan = ModelProfiler.plan(model: profile, system: system, contextSize: contextSize, draftWeightBytes: draftFootprintBytes)
@@ -1132,6 +1134,12 @@ struct MLXServer: AsyncParsableCommand {
                 let total = partitionPlan?.totalLayers ?? actual
                 let cpuCount = total - actual
                 print("[SwiftLM] 🔀 Layer split active: \(actual) GPU / \(cpuCount) CPU")
+                if modelIsMoE && cpuCount > 0 {
+                    // CPU-resident MoE layers run the quantized expert matmuls on a
+                    // single core (~0.4 tok/s prefill on Gemma 4 26B-A4B, see #176).
+                    print("[SwiftLM] ⚠️  \(cpuCount) MoE layers will run on the CPU, which is very slow (expect well under 1 tok/s).")
+                    print("[SwiftLM]    For MoE models that don't fit in GPU memory, prefer --stream-experts over --gpu-layers.")
+                }
                 // Update the partition plan to reflect actual split
                 partitionPlan?.gpuLayers = actual
             } else {