Apple Silicon 上的 LLM 推理：连续批处理与 SSD KV Cache 协同优化

在 Apple Silicon 上运行 LLM 推理曾是「高性能但低吞吐」的代名词 —— 单请求执行流畅，多并发场景却捉襟见肘。oMLX 项目从 vllm-mlx 分叉而来，通过连续批处理（Continuous Batching）与分层 KV Cache 体系，在 M 系列芯片的统一内存架构上实现了可扩展的高并发推理能力。本文聚焦其调度器设计与 SSD 缓存管理，深入剖析统一内存约束下的工程权衡。

统一内存架构的调度前提

Apple M 系列芯片采用统一内存架构，CPU、GPU（Neural Engine）与 DRAM 共享同一物理内存池。以 M4 Max 为例，可配置高达 128GB 统一内存，带宽达 546GB/s。这一架构消除了传统离散 GPU 的 PCIe 传输开销，但同时带来了内存分配策略的全新约束：LLM 推理所需的 KV Cache、模型权重与调度状态必须在一套内存池内协同管理，任何单一组件的过度消耗都会挤压其他组件的空间。

oMLX 的调度器在此约束下运行，其目标是最大化并发吞吐的同时，确保任意时刻的内存使用不超过系统阈值。关键参数为 --max-process-memory 80%（默认），即保留 8GB 供系统使用，剩余空间在模型权重、KV Cache 热区与调度元数据之间动态分配。这一保守阈值对于 192GB 配置的机器仍有较大余量，但在 16GB 入门级设备上却是防止 OOM 的生命线。

连续批处理的核心调度逻辑

oMLX 的连续批处理实现依赖 mlx-lm 的 BatchGenerator，其调度循环遵循三个阶段交替执行：请求入队、批量生成、请求退出。

请求入队阶段在每个 token 生成边界检查待处理队列与当前批次容量。若批次未满且队列非空，则从队列弹出一个新请求加入批次。这里的关键设计决策是「在 token 边界而非请求边界入队」—— 这确保了批次中的所有序列始终处于同一生成阶段，避免了 prefill 与 decode 的混合执行带来的内存碎片化问题。当 --max-concurrent-requests 设置为 16 时，调度器最多同时维护 16 个活跃序列，每个序列独立生成其下一个 token。

批量生成阶段对批次中的所有序列执行一次前向传播。MLX 的懒评估机制在此发挥关键作用：将多个序列的操作融合为单一 Metal kernel launch，减少内核启动开销。统一内存架构使得批次内所有序列的 KV Cache 无需跨设备传输即可直接访问，进一步降低了内存带宽压力。

请求退出阶段在每个 token 生成后立即检查序列是否完成。若 EOS token 生成或达到 max_tokens 限制，序列退出批次并返回完整响应，无需等待其他序列。这一设计使得短请求不会因长请求而被阻塞，从而显著改善了 P50 延迟。

实际调优中，--max-concurrent-requests 的取值需在延迟与吞吐之间权衡。默认值 8 适合单用户交互场景；16 适合多用户并发场景；超过 32 则可能导致内存争用反而降低有效吞吐。监控指标为 admin dashboard 中的批次填充率 —— 若持续低于 50%，说明并发请求不足，可适当调高参数。

分层 KV Cache：热区与冷区的换入换出

oMLX 的 KV Cache 管理采用热区（RAM）与冷区（SSD）两层架构，这一设计直接源于统一内存容量有限这一现实约束。以 Qwen3-8B 为例，上下文长度为 32K tokens 时，KV Cache 可达 20GB 以上；若同时加载多个模型，统一内存很快告急。热区缓存将最近访问的 KV blocks 保留在内存中，冷区缓存将不活跃的 blocks 溢出到 SSD。

PagedCacheManager 的块管理

oMLX 的块管理机制受 vLLM 的 PagedAttention 启发，但针对统一内存做了简化。KV Cache 按固定大小的物理块存储（默认 16 tokens/block），逻辑序列通过块链表映射到物理块。当新请求到达时，调度器首先查询文本前缀缓存（通过 SHA256 哈希匹配共享系统提示），若命中则复用已有 KV 状态，跳过 prefill 阶段；若未命中，则分配新块并执行完整的 prefill 计算。

关键配置参数为 --hot-cache-max-size，控制热区最大内存占用，默认 20% 即约 25GB（以 128GB 系统为例）。当热区填充率超过阈值时，最久未访问的块被换出至冷区。冷区存储格式为 safetensors，每个块一个文件，使用 mmap 映射以支持零拷贝读取。

冷区缓存的写入策略

冷区写入采用写回（write-back）策略：仅当块被驱逐出热区时才写入 SSD，而非每次更新后立即同步。这种延迟写入策略减少 SSD I/O 频率，但引入了崩溃恢复的风险 —— 若服务异常退出，未刷盘的块将丢失，下次请求需重新计算。oMLX 通过定期检查点机制缓解此问题，默认每 5 分钟将热区元数据（块映射表、哈希索引）写入磁盘，服务重启时从检查点恢复热区缓存结构。

对于需要高可靠性的场景，可通过 --paged-ssd-cache-dir 将 SSD 缓存路径指向具有电池保护的 SSD（如 Mac 内置 NVMe），以减少数据丢失概率。同时，--hot-cache-max-size 的合理设置可控制热区未刷盘数据量：热区越小，未刷新数据越少，恢复时重新计算的开销越低。

缓存命中的性能收益

当请求的文本前缀或图片内容命中缓存时，性能收益显著：

文本前缀缓存通过 SHA256 哈希匹配相同系统提示的场景，如客服机器人的固定指令模板或代码助手的工具描述。图片缓存则基于解码后像素值的哈希，确保同一图片的 URL、base64 或文件路径变体均命中同一缓存条目。

多模型场景下的内存压力管理

oMLX 支持在单一进程中加载多个模型（LLM、VLM、Embedding、Reranker），但统一内存容量限制了可同时驻留的模型数量。ProcessMemoryEnforcer 组件持续监控总内存使用，通过 LRU 策略自动卸载最久未访问的模型。

关键配置参数包括：

• --max-model-memory：单模型的内存上限，防止某个模型过度消耗
• --max-process-memory：进程级总内存上限，默认 auto 即 RAM 减 8GB
• Per-model TTL：闲置超时后自动卸载，如 --chat-template-kwargs ttl=300 表示 5 分钟无请求则卸载

模型固定（Pinning）功能允许将高频模型锁定在内存中，避免被 LRU 驱逐。固定操作通过 admin dashboard 的状态徽章或 CLI 的 omlx models pin <model-id> 执行，适用于始终保持加载状态的模型（如主力代码助手模型）。但需注意：固定模型会从 LRU 候选池中移除，当固定模型数量过多时，可能导致未固定模型频繁换入换出，反而降低整体效率。实践中建议将固定模型数量控制在总模型数量的 30% 以下。

调度参数调优的实用指南

基于统一内存约束与连续批处理机制，以下参数调优建议可作为初始配置参考：

低配置设备（24GB 以下 RAM）：建议 --max-concurrent-requests 4、--hot-cache-max-size 10%、--max-process-memory 70%。优先保障单请求延迟稳定，避免高并发导致内存压力触发 OOM。

中等配置设备（48GB–64GB RAM）：建议 --max-concurrent-requests 8、--hot-cache-max-size 20%。平衡并发吞吐与延迟，适合多用户或代理工作流场景。

高配置设备（128GB 以上 RAM）：建议 --max-concurrent-requests 16 或更高、--hot-cache-max-size 30%。充分利用内存资源支撑多模型并发，可运行 2–3 个 8B 模型同时服务。

此外，--hot-cache-max-size 与 --paged-ssd-cache-dir 的协同配置值得注意：SSD 缓存容量通常不受限制（取决于磁盘空间），但 SSD 读取延迟（约 100–500μs）远高于内存（约 10ns）。若冷区缓存命中率过高（如超过 30% 的请求需要从 SSD 恢复块），说明热区配置过小，应调高 --hot-cache-max-size 或减少并发模型数量。

工程权衡小结

oMLX 在 Apple Silicon 上实现了连续批处理与分层 KV Cache 的协同调度，其核心工程权衡体现在三个层面：

内存 vs 吞吐：更大热区提升缓存命中率但占用可用内存，可能限制并发模型数量；更小热区减少内存压力但增加冷区 I/O 开销。实践中 20% 热区是一个合理的折中点。

延迟 vs 吞吐：更高并发数（--max-concurrent-requests）提升总吞吐但增加队列等待时间；更低并发数保障低延迟但浪费 GPU 计算资源。BatchGenerator 的设计使得两者的权衡点比传统静态批处理更友好，但并非完全无代价。

可靠性 vs 性能：写回策略减少 SSD I/O 但引入数据丢失风险；同步写入增加 I/O 开销但保障崩溃恢复。默认配置优先性能，可通过定期检查点机制弥补可靠性缺口。

这些权衡没有绝对最优解，需根据具体硬件配置与业务场景（高吞吐 vs 低延迟、单模型 vs 多模型）进行参数迭代。oMLX 的 admin dashboard 提供了实时监控能力，建议在部署后观察批次填充率、冷区命中率与内存使用曲线，逐步调优至稳定工作点。

资料来源

• oMLX GitHub 仓库：https://github.com/jundot/omlx
• arXiv 2601.19139v1：Native LLM and MLLM Inference at Scale on Apple Silicon（vllm-mlx 原始论文）

ai-systems