NEE's Blog

Flash-MoE：在笔记本上运行 397B 参数大模型

March 22, 2026

本文翻译自 Flash-MoE: Running a 397B Parameter Model on a Laptop，原载于 Hacker News。

惊人的壮举

想象一下：在一台普通的笔记本电脑上运行一个 3970 亿参数 的大语言模型，而且能以每秒 4.4 个 token 的速度生成文本。这不是科幻小说，这是 Flash-MoE 项目的真实成果。

Flash-MoE 是一个纯 C/Metal 推理引擎，它可以在配备 48GB 内存的 MacBook Pro 上运行 Qwen3.5-397B-A17B（一个 397B 参数的 Mixture-of-Experts 模型），达到 4.4+ tokens/秒 的推理速度，并且支持完整的工具调用（tool calling）功能。

整个 209GB 的模型通过自定义的 Metal 计算管道从 SSD 流式加载。没有 Python，没有框架，只有 C、Objective-C 和手工优化的 Metal 着色器。

性能表现

配置 tok/s 质量 备注
4-bit experts, FMA kernel 4.36 优秀 当前最佳配置。完整工具调用支持。磁盘占用 209GB。
4-bit experts, baseline 3.90 优秀 FMA kernel 优化前的版本。
2-bit experts, trust OS 5.74 良好* 磁盘占用 120GB。*会破坏 JSON/工具调用。
2-bit peak single token 7.05 良好* 热缓存峰值。*不适合工具使用。

2-bit 量化会在 JSON 输出中产生 \name\ 而不是 "name"，导致工具调用不可靠。4-bit 是生产环境推荐配置。

硬件配置

这项成就的硬件基础是一台 MacBook Pro：

  • 机器：MacBook Pro, Apple M3 Max
  • 芯片：16 核 CPU（12P + 4E），40 核 GPU，16 核 ANE
  • 内存：48 GB 统一内存（约 400 GB/s 带宽）
  • SSD：1TB Apple Fabric，实测 17.5 GB/s 顺序读取
  • macOS：26.2 (Darwin 25.2.0)

模型架构

该模型包含 60 个 transformer 层：45 个 GatedDeltaNet（线性注意力）+ 15 个标准全注意力层。每层有 512 个专家（experts），每个 token 激活 K=4 个专家（外加一个共享专家）。隐藏维度为 4096。

核心技术

1. SSD 专家流式加载（SSD Expert Streaming）

专家权重（4-bit 下 209GB）通过并行 pread() 和 GCD dispatch groups 从 NVMe SSD 按需读取。每层只有 K=4 个活跃专家被加载（每个约 6.75MB）。操作系统页面缓存管理所有缓存——不需要自定义缓存（”Trust the OS” 原则）。这个设计灵感来自 Apple 的 “LLM in a Flash” 论文。

2. FMA 优化的反量化内核（FMA-Optimized Dequant Kernel）

4-bit 反量化矩阵-向量乘法的内循环将数学运算从 (nibble * scale + bias) * x 重排为 fma(nibble, scale*x, bias*x)。预先计算 scale*x 和 bias*x 让 GPU 的融合乘加单元（FMA）在一条指令中完成反量化和乘法。比原始实现快 12%。

3. Metal 计算着色器

手工编写的 Metal 内核包括：

  • 4-bit 和 2-bit 反量化矩阵-向量乘法（分块、SIMD 约简、共享输入缓存、FMA 优化）
  • 融合 SwiGLU 激活
  • RMS 归一化（两遍：平方和约简 + 应用）
  • 批量 GPU 注意力（Q@K^T, softmax, scores@V）用于全注意力层
  • GPU RoPE（与 Q 解交织和 K 归一化融合）
  • MoE combine + residual + sigmoid gate（融合内核）

4. 延迟 GPU 专家计算

CMD3（专家前向传播）被提交但不等待。GPU 执行它的同时 CPU 准备下一层。combine + residual + norm 也在 GPU 上，直接馈入下一层的注意力投影。

5. Accelerate BLAS 用于线性注意力

GatedDeltaNet 递归使用 cblas_sscal、cblas_sgemv 和 cblas_sger 进行 64 头 × 128×128 状态矩阵更新。比标量代码快 64%。

6. 信任操作系统（Trust the OS）

没有自定义专家缓存。操作系统页面缓存（约 35GB）通过标准 LRU 管理专家数据缓存。我们测试的每种自定义缓存方法（Metal LRU、malloc 缓存、LZ4 压缩缓存）都因为 GPU 内存压力或开销而更慢。页面缓存自然达到约 71% 的命中率。

每层流水线（4-bit 平均 4.28ms）

CMD3(prev) → CMD1: attention projections + delta-net  [1.22ms GPU]
           → CPU: flush results                       [0.01ms CPU]
           → CMD2: o_proj + norm + routing + shared    [0.55ms GPU]
           → CPU: softmax + topK routing               [0.003ms]
           → I/O: parallel pread K=4 experts           [2.41ms SSD]
           → CMD3: expert forward + combine + norm     [0.04ms encode, DEFERRED]

统一内存约束

在 Apple Silicon 上，SSD DMA 和 GPU 计算共享同一内存控制器，无法有效重叠。GPU 的反量化内核在约 418 GiB/s 时达到带宽饱和。即使很小的后台 SSD DMA 也会通过内存控制器仲裁导致不成比例的 GPU 延迟峰值。串行流水线（GPU → SSD → GPU）是硬件最优的。

快速开始

cd metal_infer
make
# 4-bit 推理（需要 packed_experts/ 目录）
./infer --prompt "Explain quantum computing" --tokens 100

# 2-bit 推理（更快但会破坏工具调用）
./infer --prompt "Explain quantum computing" --tokens 100 --2bit

# 带工具调用的交互式聊天
./chat

# 每层时间分解
./infer --prompt "Hello" --tokens 20 --timing

项目结构

metal_infer/
  infer.m              # 完整推理引擎（约 7000 行）
  shaders.metal        # Metal 计算内核（约 1200 行）
  chat.m               # 带工具调用的交互式聊天 TUI
  tokenizer.h          # C BPE tokenizer（单头文件，449 行）
  main.m               # MoE 基准测试
  Makefile             # 构建系统
  extract_weights.py   # 从 safetensors 创建 model_weights.bin
  repack_experts_2bit.py  # 4-bit → 2-bit 专家重量化

我们尝试过的方法（及其效果）

保留的优化

方法 结果 影响
FMA dequant kernel GPU 计算 -12% +12% tok/s
信任 OS 页面缓存 删除 Metal LRU → +38% 基础性改进
GPU combine+norm in CMD3 消除 CPU 往返 流水线优化
BLAS delta-net (Accelerate) cpu_attn 0.78→0.28ms +64% attention
F_NOCACHE for 2-bit +3%（避免页面抖动） 仅限 2-bit
GPU fused attention (RoPE) 全注意力层 +2% 小幅提升
C BPE tokenizer 180ms vs 3500ms 启动 20x 启动速度
Deferred CMD3 execution GPU/CPU 重叠 流水线优化

舍弃的方法（58 个实验中的亮点）

方法 结果 原因
LZ4 专家压缩 -13% 解压开销 > 热缓存节省
F_RDADVISE 预取 净 0% 统一内存：SSD DMA 使 GPU 变慢 -73%
时序专家预测 -18% 25% 命中率，SSD 带宽浪费
MLP 路由预测器 31% 准确率 比时序基准更差
GPU LUT dequant kernel -2% 间接寄存器访问序列化
dispatch_io -70% dispatch_data 管理开销
mmap 专家文件 -5x 冷数据上的每页故障开销
MTP 推测解码 持平 MoE I/O 按 token 扩展（与密集模型不同）

内存安全

这是一台主要开发机器。引擎明确控制内存：

  • 非专家权重：5.5GB（mmap’d，只读）
  • Metal 临时缓冲区：约 200MB
  • 总计：约 6GB，为 OS + 页面缓存留出 42GB
  • 无 OOM 风险。专家数据按需从 SSD 流式加载。
  • 没有自定义缓存。信任操作系统。

我的思考

这个项目展示了几个非常有趣的设计理念：

  1. “Trust the OS” - 作者尝试了多种自定义缓存策略，最终发现操作系统的页面缓存效果最好。这是一个很好的提醒：有时候最简单的方案就是最优解。

  2. 纯 C/Metal 实现 - 在当今 Python 和各种框架盛行的时代，作者选择从头用 C 和 Metal 实现，最终获得了极致的性能。这让我想起了早期程序员的精神。

  3. SSD 作为扩展内存 - 通过流式加载，209GB 的模型可以在 48GB 内存的机器上运行。这种思路对于资源受限的开发者来说非常有价值。

  4. 统一内存的双刃剑 - Apple Silicon 的统一内存架构既是优势（高带宽），也是限制（SSD DMA 和 GPU 计算不能并行）。作者深入理解硬件特性后做出了正确的串行流水线设计。

对于想要在本地运行大模型的开发者来说，这个项目提供了宝贵的参考。虽然 48GB 内存的 M3 Max MacBook Pro 仍然是高端设备，但相比于专业 GPU 服务器，这已经是触手可及的消费级硬件了。

关键要点

  • 🚀 397B 模型，48GB 内存 - 通过 SSD 流式加载实现
  • ⚡ 4.4+ tok/s - 生产质量的推理速度
  • 🔧 完整工具调用支持 - 4-bit 量化配置
  • 💡 Trust the OS - 操作系统页面缓存比自定义缓存更高效
  • 🛠️ 纯 C/Metal - 无 Python、无框架的极致优化
comments powered by Disqus