Merge pull request #192 from SharpAI/fix/ssd-streaming-release

fix(ssd): Qwen3.6 --stream-experts works again (b769/b773 crash), first-run streaming, README re-measure

Author: solderzzc

README.md

@@ -65,14 +65,14 @@ The first SwiftLM numbers from a **32 GB** Mac. Every other table in this README
 | 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.6-35B-A3B-UD-MLX-4bit`** (MoE, ~3B active) | 21.6 GB | GPU | **47.0 tok/s** | 40.8K tokens | 21.5 GB |
+| `Qwen3.6-35B-A3B-UD-MLX-4bit` | 21.6 GB | `--stream-experts` | 13.2 tok/s | 40.8K tokens | 5.8 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).
+- **Qwen3.6-35B-A3B on a base M6 reaches 76%** of the M1 Ultra 64 GB decode speed below (47.0 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.
 
@@ -92,10 +92,12 @@ Every needle check passed. `--mtp` with the bf16 assistant (`gemma-4-26B-A4B-it-
 
 | 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 |
+| ~550 | 714 / 47.0 | 19.8 GB | 256 / 13.2 | 5.6 GB |
+| ~2.3K | 968 / 45.7 | 20.1 GB | 403 / 13.0 | 5.6 GB |
+| ~9.8K | 858 / 43.4 | 20.4 GB | 401 / 12.7 | 5.6 GB |
+| 40.8K | 615 / 36.1 | 21.5 GB | 336 / 12.0 | 5.8 GB |
+
+> ⚠️ **`--stream-experts` crashes on Qwen3.5/3.6 in releases b769 and b773** (`broadcast_shapes … (N,8,8,2048)` on the first request). The mlx-swift-lm upstream sync in #167 broke the SSD path. Earlier versions of this table were measured before that sync and were never re-checked afterwards. Fixed in SharpAI/mlx-swift-lm#69 and #71; the table above was re-measured with those fixes.
 
 ### Qwen3.8-27B-4bit (dense)
 
@@ -119,13 +121,13 @@ Every needle check passed. `--mtp` with the bf16 assistant (`gemma-4-26B-A4B-it-
 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.
+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` after that bump; the Qwen3.6 table was re-measured with SharpAI/mlx-swift-lm#69 and #71.
 
 > ℹ️ **`--turbo-kv` long-range recall is fixed** ([#175](https://github.com/SharpAI/SwiftLM/issues/175), SharpAI/mlx-swift-lm#65). Before the fix, once a prompt passed the 2,048-token compression threshold, attention only saw the recent hot window and positions restarted, so Qwen3.8-27B-4bit got exact lookups wrong. Attention now covers the compressed history too, which makes `--turbo-kv` slower than before (97 s vs 72 s on an 11.8K-token prompt on the M6). `--turbo-kv` still has no effect when `--ctx-size` is set (the attention layers use `RotatingKVCache`).
 >
 > ℹ️ QAT-quantized Gemma 4 MTP assistants (`…-qat-assistant-4bit`) now load and run (SharpAI/mlx-swift-lm#66).
 >
-> ℹ️ **`--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).
+> ℹ️ **`--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, 5.8 GB GPU).
 
 Reproduce:
 

Sources/SwiftLM/Server.swift

@@ -641,9 +641,35 @@ struct MLXServer: AsyncParsableCommand {
         // hand-copied or huggingface-cli model it returns nil — which skipped the MoE
         // guard *and* the ExpertStreamingConfig activation while still setting lazyLoad,
         // i.e. lazy weights with no streamer (the #112 memory shape) and no diagnostic.
-        let modelDirectory =
+        var modelDirectory =
             ModelStorage.validatedContentDirectory(for: modelId)
             ?? resolveModelDirectory(modelId: modelId)
+        if self.streamExperts, !self.info, modelDirectory == nil,
+            !FileManager.default.fileExists(atPath: modelId)
+        {
+            // Streaming must be activated for the directory the loader reads, so resolve it
+            // before loading. Same hub root as the loader below, which reuses these files.
+            let hub = HubApi(
+                downloadBase: URL.applicationSupportDirectory
+                    .appendingPathComponent("MLX", isDirectory: true)
+                    .appendingPathComponent("HuggingFace", isDirectory: true))
+            let localRepo = hub.localRepoLocation(Hub.Repo(id: modelId))
+            if FileManager.default.fileExists(
+                atPath: localRepo.appendingPathComponent("config.json").path)
+            {
+                modelDirectory = localRepo
+            } else {
+                // First run. A failed download is a model problem, not a binary one.
+                phase = .architectureProbe
+                print("[SwiftLM] --stream-experts: downloading \(modelId) before loading...")
+                let prefetchTracker = ProgressTracker(modelId: modelId)
+                modelDirectory = try await hub.snapshot(
+                    from: modelId, matching: ["*.safetensors", "*.json", "*.jinja"]
+                ) { progress in
+                    prefetchTracker.printProgress(progress)
+                }
+            }
+        }
         var mainModelProfile: ModelProfile? = nil
         if self.streamExperts, let dir = modelDirectory {
             mainModelProfile = ModelProfiler.profile(modelDirectory: dir, modelId: modelId)