Merge pull request #27 from SharpAI/feature/speculative-decoding-ci

test: add speculative decoding E2E test to CI pipeline

Author: solderzzc

.github/workflows/ci.yml

@@ -106,3 +106,139 @@ jobs:
           name: ci-test-logs
           path: /tmp/SwiftLM-test-*.log
           retention-days: 7
+
+  # ── Speculative Decoding E2E (dual-model: 0.8B draft + 4B main) ──
+  # Uses the standard macos-15 runner (7 GB RAM). 
+  # We test the 4B main model which safely fits within memory.
+  speculative-decoding:
+    runs-on: macos-15
+    timeout-minutes: 45
+    needs: ci  # Only run after core CI passes
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          submodules: recursive
+
+      - name: Install Metal Toolchain
+        run: xcodebuild -downloadComponent MetalToolchain || true
+
+      - name: Cache Swift packages
+        uses: actions/cache@v4
+        with:
+          path: .build
+          key: ${{ runner.os }}-spm-SwiftLM-v2-${{ hashFiles('Package.resolved') }}
+          restore-keys: |
+            ${{ runner.os }}-spm-SwiftLM-v2-
+
+      - name: Clear stale module cache
+        run: find .build -type d -name ModuleCache -exec rm -rf {} + 2>/dev/null || true
+
+      - name: Resolve dependencies
+        run: swift package resolve
+
+      - name: Build (Release)
+        run: swift build -c release
+
+      - name: Install MLX Metal library
+        run: |
+          python3 -m venv /tmp/mlx_venv
+          /tmp/mlx_venv/bin/pip install --quiet mlx
+          cp /tmp/mlx_venv/lib/python*/site-packages/mlx/lib/mlx.metallib .build/release/
+
+      - name: Cache MLX models (draft + main)
+        uses: actions/cache@v4
+        with:
+          path: ~/.cache/huggingface
+          key: mlx-speculative-qwen35-0.8b-9b
+
+      - name: Run speculative decoding E2E
+        env:
+          HF_HUB_DOWNLOAD_TIMEOUT: "900"
+        run: |
+          chmod +x tests/test-speculative.sh
+          for attempt in 1 2 3; do
+            echo "Attempt $attempt of 3..."
+            if tests/test-speculative.sh .build/release/SwiftLM 15414; then
+              exit 0
+            fi
+            if [ "$attempt" -lt 3 ]; then
+              echo "Test failed, retrying in 10s..."
+              sleep 10
+            fi
+          done
+          echo "All attempts failed"
+          exit 1
+
+      - name: Upload speculative test logs on failure
+        if: failure()
+        uses: actions/upload-artifact@v4
+        with:
+          name: speculative-test-logs
+          path: /tmp/SwiftLM-test-speculative.log
+          retention-days: 7
+
+  # ── Speculative Decoding Memory Evaluation ──
+  # Runs the 9B model with NUM_DRAFT_TOKENS=2 to check peak
+  # memory compression/efficiency. Allowed to OOM/fail.
+  speculative-decoding-eval:
+    runs-on: macos-15
+    timeout-minutes: 45
+    needs: ci
+    continue-on-error: true
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          submodules: recursive
+      
+      - name: Install Metal Toolchain
+        run: xcodebuild -downloadComponent MetalToolchain || true
+
+      - name: Cache Swift packages
+        uses: actions/cache@v4
+        with:
+          path: .build
+          key: ${{ runner.os }}-spm-SwiftLM-v2-${{ hashFiles('Package.resolved') }}
+          restore-keys: |
+            ${{ runner.os }}-spm-SwiftLM-v2-
+
+      - name: Clear stale module cache
+        run: find .build -type d -name ModuleCache -exec rm -rf {} + 2>/dev/null || true
+
+      - name: Resolve dependencies
+        run: swift package resolve
+
+      - name: Build (Release)
+        run: swift build -c release
+      
+      - name: Install MLX Metal library
+        run: |
+          python3 -m venv /tmp/mlx_venv
+          /tmp/mlx_venv/bin/pip install --quiet mlx
+          cp /tmp/mlx_venv/lib/python*/site-packages/mlx/lib/mlx.metallib .build/release/
+      
+      - name: Run speculative evaluation E2E
+        env:
+          HF_HUB_DOWNLOAD_TIMEOUT: "900"
+        run: |
+          chmod +x tests/test-speculative-eval.sh
+          for attempt in 1 2 3; do
+            echo "Attempt $attempt of 3..."
+            if tests/test-speculative-eval.sh .build/release/SwiftLM 15414; then
+              exit 0
+            fi
+            if [ "$attempt" -lt 3 ]; then
+              echo "Test failed, retrying in 10s..."
+              sleep 10
+            fi
+          done
+          echo "All attempts failed"
+          exit 1
+          
+      - name: Upload speculative eval logs on failure
+        if: failure()
+        uses: actions/upload-artifact@v4
+        with:
+          name: speculative-eval-logs
+          path: /tmp/SwiftLM-test-speculative-eval.log
+          retention-days: 7
+

Package.resolved

@@ -13,7 +13,7 @@ let package = Package(
         // Local Apple MLX Swift fork for C++ extensions
         .package(url: "https://github.com/SharpAI/mlx-swift.git", branch: "main"),
         // Apple's LLM library built on MLX Swift (SharpAI fork — with GPU/CPU layer partitioning)
-        .package(url: "https://github.com/ericjlake/mlx-swift-lm.git", branch: "feat/ssd-streaming-10x"),
+        .package(url: "https://github.com/SharpAI/mlx-swift-lm.git", branch: "main"),
         // HuggingFace tokenizers + model download
         .package(url: "https://github.com/huggingface/swift-transformers", .upToNextMinor(from: "1.2.0")),
         // Lightweight HTTP server (Apple-backed Swift server project)

Package.swift

@@ -81,7 +81,8 @@ Benchmark results for `gemma-4-26b-a4b-it-4bit` (26B MoE, 4-bit) on M5 Pro 64 GB
 - 🔌 **OpenAI-compatible**: Drop-in replacement for OpenAI SDKs (`/v1/chat/completions`, streaming, etc).
 - 🧠 **Smart Model Routing**: Loads HuggingFace format models directly, with native Safetensors parsing.
 - ⚡️ **TurboQuantization Integrated**: Custom low-level MLX Metal primitives that apply extremely fast quantization for KV caching out-of-the-box.
-- 💾 **SSD Expert Streaming**: *Experimental* zero-copy streaming that swaps Mixture of Experts (MoE) layers directly from the NVMe SSD to the GPU command buffer without trashing macOS Unified Memory (prevents Watchdog OS kernel panics on 122B+ models).
+- 💾 **SSD Expert Streaming (10x)**: High-performance NVMe streaming that loads Mixture of Experts (MoE) layers directly from SSD to GPU — engineered by [@ericjlake](https://github.com/ericjlake), achieving **10x speedup** (0.58 → 5.91 tok/s) on 122B+ models with only ~10 GB resident memory. Uses cross-projection batching, concurrent pread (QD=24), asyncEval pipeline, and runtime top-k expert selection.
+- 🔮 **Speculative Decoding**: Load a small draft model (e.g. 9B) alongside a large main model to generate candidate tokens and verify in bulk — accelerating in-RAM inference.
 - 🎛️ **Granular Memory Control**: Integrated Layer Partitioning (`--gpu-layers`) and Wisdom Auto-Calibration for squeezing massive models into RAM.
 
 ---
@@ -146,6 +147,64 @@ Reference implementations: [`turboquant-mlx`](https://github.com/sharpner/turboq
 
 ---
 
+## 💾 SSD Expert Streaming: 10x MoE Speedup
+
+SwiftLM implements a **rewritten SSD expert streaming pipeline** (engineered by [Eric Lake](https://github.com/ericjlake)) that achieves 10x generation speedup for massive Mixture of Experts (MoE) models running on memory-constrained Apple Silicon. This enables running models like **Qwen3.5-122B** (69.6 GB) and **Qwen3.5-397B** (209 GB) on a **64 GB Mac** by streaming expert weights from NVMe SSD.
+
+### Benchmark Results (M1 Ultra 64GB, Qwen3.5-122B-A10B-4bit)
+
+| Configuration | tok/s | vs. Original | Notes |
+|---|---|---|---|
+| Original `--stream-experts` | 0.58 | baseline | Sequential pread, 1 NVMe queue |
+| **This PR (top-k=8, full quality)** | **4.95** | **8.5×** | All 8 experts evaluated |
+| **This PR (top-k=6, default)** | **5.20** | **9.0×** | Recommended default |
+| **This PR (top-k=4, speed mode)** | **5.91** | **10.2×** | Best quality/speed tradeoff |
+| **This PR (top-k=2, turbo mode)** | **6.52** | **11.2×** | Still coherent output |
+
+> Memory stable at **~10.6 GB resident**, no swap activity. Tested over 200-token generation runs.
+
+### The Approach: Small Model Helps Large Model
+
+A novel aspect of this architecture is the **dual-model speculative decoding** pattern: a small draft model (e.g. Qwen3.5-9B at 73 tok/s) runs **entirely in RAM** while the large MoE model (e.g. 122B) streams experts from SSD. The draft model generates candidate tokens at high speed, and the main model verifies them in bulk — dramatically reducing the number of SSD-bound generation rounds needed.
+
+> **Important finding:** Speculative decoding is **counterproductive for SSD-streaming MoE** specifically. The verify pass sends N+1 tokens, each routing to *different* experts — SSD I/O scales with the *union* of all positions' expert selections. Speculative decoding is therefore routed exclusively to **in-RAM models**.
+
+### Optimization Techniques
+
+1. **Cross-Projection Batching**: Collapses ~1,400 per-expert `eval()` calls down to ~48 per token by orchestrating gate/up/down projections together in `SwitchGLU`.
+2. **Concurrent NVMe pread (QD=24)**: Replaces sequential pread with `DispatchQueue.concurrentPerform`, saturating the NVMe controller's queue depth (8 experts × 3 projections = 24 parallel reads).
+3. **AsyncEval Pipeline with Speculative Pread**: Overlaps GPU compute with SSD I/O — uses previous-token routing to speculatively pre-load experts for the next token during the GPU async window (~70% hit rate). Only missed experts (~30%) require on-demand pread after routing sync.
+4. **Persistent Metal Buffers**: Expert weight buffers are allocated once per `SwitchGLU` layer and reused across tokens, eliminating per-token allocation overhead.
+5. **Runtime Top-K Expert Selection**: The `SWIFTLM_TOP_K` environment variable reduces the number of active experts per token at runtime without model recompilation — trading marginal quality for significant speed gains.
+
+### Key Engineering Findings
+
+| Finding | Detail |
+|---|---|
+| **GPU compute is the bottleneck** | At steady state, GPU compute is ~190ms of ~200ms per-token time. The OS page cache serves ~90% of expert reads from RAM. |
+| **Don't cache experts in application memory** | An LRU expert cache *stole* from the OS page cache and regressed performance (4.84 → 4.01 tok/s). Let the kernel manage it. |
+| **MambaCache requires checkpoint rollback** | Unlike attention KV caches (trim = decrement offset), Mamba's recurrent state integrates all history and cannot be partially undone. We implemented `checkpoint()`/`restore()` for speculative decoding on hybrid Attention+Mamba architectures (Qwen3.5). |
+
+### Usage
+
+```bash
+# Standard SSD streaming (recommended, top-k=6):
+SWIFTLM_TOP_K=6 SwiftLM --port 8002 \
+  --model <path>/Qwen3.5-122B-A10B-4bit --stream-experts
+
+# Speed mode (top-k=4):
+SWIFTLM_TOP_K=4 SwiftLM --port 8002 \
+  --model <path>/Qwen3.5-122B-A10B-4bit --stream-experts
+
+# With speculative decoding (in-RAM models only):
+SwiftLM --port 8002 \
+  --model <path>/Qwen3.5-27B-4bit \
+  --draft-model <path>/Qwen3.5-9B-4bit \
+  --num-draft-tokens 4
+```
+
+---
+
 ## 💻 Benchmarks & Testing
 
 Run our automated benchmark suites via the interactive script:
@@ -226,8 +285,10 @@ curl http://localhost:5413/v1/chat/completions \
 | `--max-tokens` | `2048` | Max tokens limit per generation |
 | `--prefill-size`| `512`  | Prompt prefill chunk size (micro-batching for long contexts) |
 | `--gpu-layers` | `model_default`| Restrict the amount of layers allocated to GPU hardware |
-| `--stream-experts` | `false` | Enable experimental SSD streaming for MoE model expert matrices |
+| `--stream-experts` | `false` | Enable SSD expert streaming for MoE models (10x speedup) |
 | `--turbo-kv` | `false` | Enable TurboQuant 3-bit KV cache compression |
+| `--draft-model` | (none) | Draft model path/ID for speculative decoding (in-RAM models only) |
+| `--num-draft-tokens` | `4` | Number of draft tokens per speculation round |
 
 ## 📦 Requirements
 
@@ -247,7 +308,13 @@ The model instantly woke up from "whispering" whitespace and successfully respon
 
 ## 🙏 Acknowledgments & Credits
 
-`SwiftLM` leverages the powerful foundation of the Apple MLX community and relies heavily on the open-source ecosystem. While the custom C++ implementations, Metal optimizations, and high-performance pipeline architecture were engineered natively for this engine, we owe massive thanks to the following projects for their indispensable reference materials and underlying protocols:
+`SwiftLM` leverages the powerful foundation of the Apple MLX community and relies heavily on the open-source ecosystem. While the custom C++ implementations, Metal optimizations, and high-performance pipeline architecture were engineered natively for this engine, we owe massive thanks to the following projects and contributors for their indispensable reference materials and underlying protocols:
+
+### Contributors
+
+- **[Eric Lake](https://github.com/ericjlake)** — Engineered the **SSD Expert Streaming 10x rewrite** ([PR #26](https://github.com/SharpAI/SwiftLM/pull/26)), achieving 10× generation speedup on 122B+ MoE models via cross-projection batching, concurrent NVMe pread (QD=24), asyncEval pipeline with speculative pread, and runtime top-k expert selection. Also implemented the **speculative decoding infrastructure** with `DraftModelRef`, dual-model loading, and **MambaCache checkpoint/restore** for hybrid Attention+Mamba architectures.
+
+### Projects & References
 
 - **[mlx-swift](https://github.com/ml-explore/mlx-swift)** — The core Apple MLX wrapper bringing Metal-accelerated operations into the Swift ecosystem.
 - **[mlx-lm](https://github.com/ml-explore/mlx/tree/main/mlx_lm)** — The official Python language models implementation, serving as the core inspiration for our chunked-prefill architecture and attention manipulation logic.