Skip to content
Search

Flash-MoE

lastmod 2026-03-24created 2026-03-241 edits

Running a 397-billion-parameter model on a laptop

About

  • A C/Metal-based inference engine that runs a 397B-parameter Mixture-of-Experts model on a MacBook Pro (48GB RAM) at over 4.4 tokens per second
  • Streams the entire 209GB model from SSD, implemented with only C and Metal shaders, without Python or frameworks
  • Maximizes GPU·SSD·CPU parallel efficiency with SSD Expert Streaming, FMA-optimized kernels, Deferred GPU Compute, and more
  • The 4-bit quantization configuration balances quality and speed, producing production-grade output including tool calling
  • A case of lightweighting and optimization that makes real-time inference of a very large MoE model possible even on a laptop

Performance results

  • 4.36 tok/s in the 4-bit experts (FMA kernel) configuration, excellent quality, full 209GB on disk
  • The 4-bit baseline configuration runs at 3.90 tok/s, the stage before FMA optimization
  • The 2-bit experts (trust OS) configuration is faster at 5.74 tok/s, but tool calling fails due to JSON output errors
  • 2-bit peak single token reaches up to 7.05 tok/s, but is unsuitable for real use
  • 4-bit quantization is the configuration suited to actual operation

Hardware environment

  • MacBook Pro (Apple M3 Max), 16-core CPU (12P+4E), 40-core GPU, 16-core ANE
  • 48GB unified memory, about 400GB/s bandwidth
  • 1TB Apple Fabric SSD, 17.5GB/s sequential read speed
  • macOS 26.2 (Darwin 25.2.0) environment

Model architecture

  • 60 transformer layers in total: 45 GatedDeltaNet (linear attention) + 15 full attention
  • Each layer has 512 experts, with K=4 activated per token (including 1 shared expert)
  • Hidden dimension 4096
  • Core techniques
    • SSD Expert Streaming
      • Expert weights (209GB at 4-bit) are loaded on demand from NVMe SSD with parallel pread()
      • Only the 4 experts activated in each layer are loaded (about 6.75MB each)
      • The OS page cache manages caching automatically, so no separate cache is needed
      • A design inspired by Apple's “LLM in a Flash” paper
    • FMA-optimized dequant kernel
      • Rearranges the (nibble * scale + bias) * x operation into the form fma(nibble, scalex, biasx)
      • Precomputes scalex and biasx so the GPU FMA unit performs it in a single instruction
      • 12% speedup over a naive implementation
    • Metal Compute Shaders
      • 4-bit/2-bit dequant matrix-vector multiplication, SwiGLU activation, RMS normalization, GPU attention (Q@Kᵀ, softmax, scores@V), RoPE, MoE combine+residual+gate, etc. implemented as hand-coded Metal kernels
    • Deferred GPU Expert Compute
      • Submits CMD3 (expert forward pass) commands asynchronously so the CPU prepares the next layer while the GPU is running
      • Combine+normalization+residual operations are also performed on the GPU and passed directly to the next layer
    • Using Accelerate BLAS
      • Uses cblas_sscal, cblas_sgemv, cblas_sger for GatedDeltaNet's recurrent computation
      • 64% faster than scalar code
    • Trust the OS
      • Custom cache removed; the OS page cache (LRU-based, about 35GB) handles caching of expert data
      • Custom Metal LRU, malloc cache, and LZ4 compressed cache were all slower
      • Achieves a natural 71% cache hit rate
  • Unified memory constraints
    • On Apple Silicon, SSD DMA and GPU computation share the same memory controller
    • Running in parallel saturates GPU bandwidth, causing latency to spike
    • A sequential GPU → SSD → GPU pipeline is the hardware-optimal form

See also

Favorite site

author Zer0 sitemap