oMLX SSD Cache 死锁修复

2026-02-15 ~ 16 · PR #16 · 从发现到彻底修复的完整故事

问题

oMLX 启用 SSD Paged Cache 后，推理 2-15 分钟后静默死锁：进程活着但推理永远不返回，无任何错误日志。

禁用 SSD cache (--no-cache) 后完全稳定。

根因

save_block() 和 load_block() 在推理线程中执行同步磁盘 I/O：

• save_block: 写入 ~10MB x 80 layers 到 SSD
• load_block: 从 SSD 读取 prefix cache

两者都在 self._lock (RLock) 保护下执行，阻塞了推理调度器，推理完全停止。

方案演进：4 个版本

v1: 后台线程直接写入 ❌

后台线程调用 MLX 保存函数导致 Metal command buffer 断言失败。

原因: Metal command buffers 不是线程安全的。

v2: numpy 转换后台保存 ❌

numpy 不支持 bfloat16 dtype，转换直接失败。

v3: 主线程写临时文件 + 后台原子重命名 ✅

• 主线程: 物化 lazy arrays + 写入临时文件 — Metal-safe
• 后台 daemon thread: os.rename(temp, final) — O(1) 原子操作
• _pending_writes 内存暂存避免竞态

100 次短请求压测通过，28 分钟零失败。

v4: 移除 load executor ✅

v3 通过短请求压测后，发现长请求仍然挂：

ThreadPoolExecutor(max_workers=1)  # 只有 1 个 worker

1. 请求 A 的 mx.load() 在 worker 线程卡住（Metal 争用）
2. 5s timeout 触发，返回 None，但 worker 仍然卡着
3. 请求 B 排队等待，推理挂起

v4 方案: 移除 executor，mx.load() 改为主线程直接同步调用。10MB @ 5GB/s SSD = 2ms，延迟可忽略。

最终架构

写入路径

主线程                          后台 daemon thread
  │                                  │
  ├─ 物化 lazy arrays (Metal-safe)   │
  ├─ 写入临时文件 (Metal-safe)       │
  ├─ 更新 _pending_writes            │
  └─ 入队 ────────────────────────→ os.rename(tmp, final)
      ↑                              │
      └──── 立即返回 ────────────────┘

读取路径

1. 查 _pending_writes → 命中 → 从内存直接返回（零 I/O）
2. 未命中 → 主线程同步 mx.load()（~2ms）

关键设计决策

为什么压测没发现 v3 的问题？

教训: 短请求压测无法发现 prefix cache 死锁，必须用长 prompt + prefix 匹配场景测试。

测试结果

• 56 个单元测试全部通过
• 100 次短请求压测：28 分钟零失败
• 20 次混合长/短 prompt：零错误、零超时、零死锁
• SSD cache 不影响推理速度：稳定 ~25 tok/s

开源贡献

• Issue: jundot/omlx#15↗
• PR: jundot/omlx#16↗
• 变更: +785/-329 行，9 个新测试

教训

1. 同步 I/O 绝不能在推理热路径执行 — 哪怕有锁保护，锁本身就是问题
2. MLX/Metal 操作必须留在主线程 — 后台线程调用 MLX 函数都可能导致 Metal crash
3. pending_writes 内存暂存是关键 — 避免 save 后立即 load 的竞态
4. 压测要覆盖真实场景 — 短请求 ≠ 长请求，prefix cache 行为完全不同
5. GPU Hang 和代码死锁是两个独立问题 — 不要混淆

修复这个问题花了两天，但这是 AIOS 五天里最有技术含量的工作。从发现问题到根因分析到四个版本的方案迭代，最终贡献回开源社区。