oMLX 两级 KV 缓存深度解析vLLM 启发的分页缓存、CoW 与前缀共享【免费下载链接】omlxLLM inference server with continuous batching SSD caching for Apple Silicon — managed from the macOS menu bar项目地址: https://gitcode.com/GitHub_Trending/om/omlxoMLX是面向 Apple Silicon 的开源 LLM 推理服务器支持连续批处理并由 macOS 菜单栏统一管理。它的核心亮点之一就是两级 KV 缓存Hot/Cold Tiered KV Cache借鉴 vLLM 的分页Paged缓存架构用块Block 块表 链式哈希前缀查找 写时复制Copy-on-Write, CoW实现跨请求的前缀共享再把冷数据下沉到 SSD让重复对话秒级命中而不是从头重算。 为什么 KV 缓存是本地 LLM 的性能命门LLM 处理长提示词Prefill时每个 token 都要计算全部注意力层的 Key/Value 向量。如果缓存了这些结果下次遇到相同前缀比如同一份系统提示词、同一个长文档就能直接复用跳过重复计算。对本地 Mac 来说这往往意味着首字延迟从几十秒降到几秒。oMLX 的缓存栈分三层源码都集中在 omlx/cache/ 目录PagedCacheManager内存中的分页 KV 缓存块式分配、前缀共享、CoW实现在 omlx/cache/paged_cache.pyHot Cache高频块的内存热层write-backPagedSSDCacheManagerSSD 冷层按块存储为 safetensors 文件实现在 omlx/cache/paged_ssd_cache.py组装逻辑由 omlx/cache/factory.py 中的CacheFactory统一完成——三个缓存创建好后SSD 缓存会注册到分页缓存上实现内存未命中 → 自动查 SSD的两级回退。 一级热层vLLM 风格的分页 KV 缓存固定大小块 块表像虚拟内存一样管理 KV传统实现为每个请求分配一整块连续 KV 内存容易碎片化和浪费。oMLX 采用 vLLM 的思路把 KV 切成固定 token 数默认 64的块CacheBlock每个块带block_id、引用计数ref_count和内容哈希block_hash每个请求维护一张块表BlockTable把逻辑 token 位置映射到物理块 ID块可以物理不连续空闲块由一条双向链表 FreeKVCacheBlockQueue管理分配、释放、中间摘除都是 O(1)天然就是 LRU 淘汰顺序队首最旧、最先被回收块池还是弹性的初始只创建 256 块内存吃紧时按 256 步长自动扩容到max_blocks默认 1000既省内存又不浪费。链式哈希前缀共享的指纹前缀共享的关键是怎么判断两个请求的前缀完全一样。oMLX 在 omlx/cache/paged_cache.py#L78-L119 用链式哈希解决每块的哈希 SHA256(模型名 父块哈希 本块 token)第一块的父哈希是固定种子omlx-root。这样哈希天然编码了整条前缀前一个 token 不同后面所有块的哈希都不同。新请求到来时get_computed_blocks沿链逐块计算期望哈希并查表直到查不到为止——命中的部分直接复用未命中的才开始计算。O(1) 淘汰与弹性扩容引用计数为 0 的块回到空闲队尾最新位置下次分配时从队首最久未用取走淘汰与分配都是 O(1)。当空闲块不足时allocate_block会触发_grow_blocks动态扩容彻底回收则走evict_block_permanently同时清理哈希索引并通知前缀索引on_block_hash_dropped钩子保证两级索引永远一致。 CoW 与前缀共享多请求共用一份 KV多个请求共享同一段前缀例如多个会话共用同一系统提示词时块被多个请求引用ref_count 1即为共享块。这里有两个经典操作Fork分叉omlx/cache/paged_cache.py#L1232-L1253 中fork_block_table把源请求的块表浅拷贝给新请求并对每块引用计数 1——零拷贝共享全部前缀。写时复制CoW进入生成阶段时get_blocks_for_generation会检查每个块是否仍被共享omlx/cache/paged_cache.py#L1255-L1310是则复制出一个新块、替换块表指针并释放共享引用。只有在真正要写入不同内容时才发生复制这正是写时复制的精髓。一个容易被忽略的细节find_shared_prefix返回块时会同时返回查找瞬间的哈希取引用必须通过acquire_cached_block(id, hash)在同一把锁内校验哈希——这封死了并发淘汰 重新分配的 TOCTOU 窗口保证不会把别的块的内容拼进你的前缀链。 二级冷层Paged SSD 缓存内存终究有限。oMLX 的冷层 omlx/cache/paged_ssd_cache.py 把整个块池的数据都落盘到 SSDsafetensors 格式、按哈希分子目录组织并用 LRU 控制总大小默认上限 100GB。几个关键设计懒恢复Lazy Restore内存查不到某块、但 SSD 上存在时只先注册一个元数据块引用计数 0等真正有请求认领时才加载数据——查前缀几乎零成本见get_computed_blocks中的 lazy restore 分支异步写队列写盘在后台线程进行队列深度按块字节数 × 主机内存比例自动计算避免大模型场景撑爆内存也避免写盘卡住推理主循环跨重启持久化SSD 块带版本化的缓存格式当前 V3与兼容性签名模型名、层数、块大小、缓存类型重启服务后旧缓存仍可直接复用不匹配的旧块则安全地当作未命中此外还有一个 SSD-only 模式的前缀索引 omlx/cache/prefix_cache.pyBlockAwarePrefixCache在块粒度上维护更细的前缀索引与分页缓存通过生命周期钩子保持同步。如上图管理面板会实时展示已处理 token / 缓存 token / 缓存效率——官方示例场景下缓存效率达 81.2%意味着近 8 成 token 来自缓存复用而非重算。⚙️ 一键上手如何调整两级缓存大小oMLX 通过命令行或管理面板Resource Management 页即可调节热/冷层配比常用参数参数作用--hot-cache-max-size 20%热层内存 KV 块缓存占内存比例0 表示禁用SSD 缓存上限默认 100GB冷层 safetensors 块总量LRU 自动淘汰paged_cache_block_size默认 256每块 token 数影响粒度与开销的平衡 经验法则长会话、多 Agent 工具流如 Claude Code 场景前缀复用率极高可适当调大冷层上限短问答为主则把热层调大即可。缓存命中/未命中、CoW 复制次数等指标由 omlx/cache/stats.py 统一统计hits / misses / cow_copies / evictions / shared_blocks。️ 实际效果长对话越聊越快对普通用户最直观的感受是多轮对话、长文档问答时模型只计算新增部分。下图中模型连续处理多轮提问与图片历史上下文全部命中缓存响应明显更快✅ 小结分页化KV 按固定块分配 块表映射O(1) 分配/淘汰块池弹性扩容前缀共享链式哈希父哈希 token 模型名让任意相同前缀的块自动去重复用CoWfork 零拷贝共享写入时才复制共享块多会话前缀成本近乎一次两级缓存内存热层 SSD 冷层safetensors 块 LRU 异步写 懒恢复重启后缓存依然可用核心实现见 omlx/cache/paged_cache.py、omlx/cache/paged_ssd_cache.py端到端行为可用 tests/test_paged_cache.py 与 tests/test_paged_ssd_cache.py 验证【免费下载链接】omlxLLM inference server with continuous batching SSD caching for Apple Silicon — managed from the macOS menu bar项目地址: https://gitcode.com/GitHub_Trending/om/omlx创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 oMLX 两级 KV 缓存深度解析：vLLM 启发的分页缓存、CoW 与前缀共享
oMLX 两级 KV 缓存深度解析vLLM 启发的分页缓存、CoW 与前缀共享【免费下载链接】omlxLLM inference server with continuous batching SSD caching for Apple Silicon — managed from the macOS menu bar项目地址: https://gitcode.com/GitHub_Trending/om/omlxoMLX是面向 Apple Silicon 的开源 LLM 推理服务器支持连续批处理并由 macOS 菜单栏统一管理。它的核心亮点之一就是两级 KV 缓存Hot/Cold Tiered KV Cache借鉴 vLLM 的分页Paged缓存架构用块Block 块表 链式哈希前缀查找 写时复制Copy-on-Write, CoW实现跨请求的前缀共享再把冷数据下沉到 SSD让重复对话秒级命中而不是从头重算。 为什么 KV 缓存是本地 LLM 的性能命门LLM 处理长提示词Prefill时每个 token 都要计算全部注意力层的 Key/Value 向量。如果缓存了这些结果下次遇到相同前缀比如同一份系统提示词、同一个长文档就能直接复用跳过重复计算。对本地 Mac 来说这往往意味着首字延迟从几十秒降到几秒。oMLX 的缓存栈分三层源码都集中在 omlx/cache/ 目录PagedCacheManager内存中的分页 KV 缓存块式分配、前缀共享、CoW实现在 omlx/cache/paged_cache.pyHot Cache高频块的内存热层write-backPagedSSDCacheManagerSSD 冷层按块存储为 safetensors 文件实现在 omlx/cache/paged_ssd_cache.py组装逻辑由 omlx/cache/factory.py 中的CacheFactory统一完成——三个缓存创建好后SSD 缓存会注册到分页缓存上实现内存未命中 → 自动查 SSD的两级回退。 一级热层vLLM 风格的分页 KV 缓存固定大小块 块表像虚拟内存一样管理 KV传统实现为每个请求分配一整块连续 KV 内存容易碎片化和浪费。oMLX 采用 vLLM 的思路把 KV 切成固定 token 数默认 64的块CacheBlock每个块带block_id、引用计数ref_count和内容哈希block_hash每个请求维护一张块表BlockTable把逻辑 token 位置映射到物理块 ID块可以物理不连续空闲块由一条双向链表 FreeKVCacheBlockQueue管理分配、释放、中间摘除都是 O(1)天然就是 LRU 淘汰顺序队首最旧、最先被回收块池还是弹性的初始只创建 256 块内存吃紧时按 256 步长自动扩容到max_blocks默认 1000既省内存又不浪费。链式哈希前缀共享的指纹前缀共享的关键是怎么判断两个请求的前缀完全一样。oMLX 在 omlx/cache/paged_cache.py#L78-L119 用链式哈希解决每块的哈希 SHA256(模型名 父块哈希 本块 token)第一块的父哈希是固定种子omlx-root。这样哈希天然编码了整条前缀前一个 token 不同后面所有块的哈希都不同。新请求到来时get_computed_blocks沿链逐块计算期望哈希并查表直到查不到为止——命中的部分直接复用未命中的才开始计算。O(1) 淘汰与弹性扩容引用计数为 0 的块回到空闲队尾最新位置下次分配时从队首最久未用取走淘汰与分配都是 O(1)。当空闲块不足时allocate_block会触发_grow_blocks动态扩容彻底回收则走evict_block_permanently同时清理哈希索引并通知前缀索引on_block_hash_dropped钩子保证两级索引永远一致。 CoW 与前缀共享多请求共用一份 KV多个请求共享同一段前缀例如多个会话共用同一系统提示词时块被多个请求引用ref_count 1即为共享块。这里有两个经典操作Fork分叉omlx/cache/paged_cache.py#L1232-L1253 中fork_block_table把源请求的块表浅拷贝给新请求并对每块引用计数 1——零拷贝共享全部前缀。写时复制CoW进入生成阶段时get_blocks_for_generation会检查每个块是否仍被共享omlx/cache/paged_cache.py#L1255-L1310是则复制出一个新块、替换块表指针并释放共享引用。只有在真正要写入不同内容时才发生复制这正是写时复制的精髓。一个容易被忽略的细节find_shared_prefix返回块时会同时返回查找瞬间的哈希取引用必须通过acquire_cached_block(id, hash)在同一把锁内校验哈希——这封死了并发淘汰 重新分配的 TOCTOU 窗口保证不会把别的块的内容拼进你的前缀链。 二级冷层Paged SSD 缓存内存终究有限。oMLX 的冷层 omlx/cache/paged_ssd_cache.py 把整个块池的数据都落盘到 SSDsafetensors 格式、按哈希分子目录组织并用 LRU 控制总大小默认上限 100GB。几个关键设计懒恢复Lazy Restore内存查不到某块、但 SSD 上存在时只先注册一个元数据块引用计数 0等真正有请求认领时才加载数据——查前缀几乎零成本见get_computed_blocks中的 lazy restore 分支异步写队列写盘在后台线程进行队列深度按块字节数 × 主机内存比例自动计算避免大模型场景撑爆内存也避免写盘卡住推理主循环跨重启持久化SSD 块带版本化的缓存格式当前 V3与兼容性签名模型名、层数、块大小、缓存类型重启服务后旧缓存仍可直接复用不匹配的旧块则安全地当作未命中此外还有一个 SSD-only 模式的前缀索引 omlx/cache/prefix_cache.pyBlockAwarePrefixCache在块粒度上维护更细的前缀索引与分页缓存通过生命周期钩子保持同步。如上图管理面板会实时展示已处理 token / 缓存 token / 缓存效率——官方示例场景下缓存效率达 81.2%意味着近 8 成 token 来自缓存复用而非重算。⚙️ 一键上手如何调整两级缓存大小oMLX 通过命令行或管理面板Resource Management 页即可调节热/冷层配比常用参数参数作用--hot-cache-max-size 20%热层内存 KV 块缓存占内存比例0 表示禁用SSD 缓存上限默认 100GB冷层 safetensors 块总量LRU 自动淘汰paged_cache_block_size默认 256每块 token 数影响粒度与开销的平衡 经验法则长会话、多 Agent 工具流如 Claude Code 场景前缀复用率极高可适当调大冷层上限短问答为主则把热层调大即可。缓存命中/未命中、CoW 复制次数等指标由 omlx/cache/stats.py 统一统计hits / misses / cow_copies / evictions / shared_blocks。️ 实际效果长对话越聊越快对普通用户最直观的感受是多轮对话、长文档问答时模型只计算新增部分。下图中模型连续处理多轮提问与图片历史上下文全部命中缓存响应明显更快✅ 小结分页化KV 按固定块分配 块表映射O(1) 分配/淘汰块池弹性扩容前缀共享链式哈希父哈希 token 模型名让任意相同前缀的块自动去重复用CoWfork 零拷贝共享写入时才复制共享块多会话前缀成本近乎一次两级缓存内存热层 SSD 冷层safetensors 块 LRU 异步写 懒恢复重启后缓存依然可用核心实现见 omlx/cache/paged_cache.py、omlx/cache/paged_ssd_cache.py端到端行为可用 tests/test_paged_cache.py 与 tests/test_paged_ssd_cache.py 验证【免费下载链接】omlxLLM inference server with continuous batching SSD caching for Apple Silicon — managed from the macOS menu bar项目地址: https://gitcode.com/GitHub_Trending/om/omlx创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考