HiCache 系统设计与优化#

本文档提供了 SGLang HiCache 的全面概述,涵盖了其系统架构、工作流程和关键组件。同时详细介绍了配置参数、优化技术以及与各种 L3 存储后端的集成,为用户和开发者提供了完整的参考,以便理解和调优 HiCache 以实现高效的 LLM 推理。

为什么需要 HiCache 及其定义是什么?#

在大语言模型推理中,预填充阶段通常耗时较长:输入序列需要首先转换为键值缓存(KV cache)以供后续解码使用。当多个请求共享相同前缀时,该前缀的 KV cache 是完全相同的。通过缓存和重用这些共享的 KV cache,可以避免冗余计算。为此,SGLang 引入了 RadixAttention,它利用空闲的 GPU 内存来缓存和重用前缀 KV cache,以及 HiCache,它将这一思想扩展到主机内存和分布式存储。

受现代 CPU 经典三级缓存设计的启发,HiCache 将 GPU 内存组织为 L1,主机内存组织为 L2,分布式存储组织为 L3。这种层次结构使 HiCache 能够充分利用 GPU 和 CPU 的"空闲"存储空间,同时整合了 Mooncake、3FS、NIXL 和 AIBrix KVCache 等分布式缓存系统,用于全局 KV cache 的存储和调度。因此,HiCache 在保持强读取性能的同时显著扩展了 KV cache 容量——特别是在多问答(multi-QA)和长上下文推理等工作负载中,KV cache 的重用频率很高。有关详细的基准测试结果,请参阅这篇博客

系统设计#

总体架构#

在许多现代 CPU 架构中,小巧快速的 L1 和 L2 缓存是每个核心专用的,能够快速访问最热的数据,而较大的 L3 缓存则是所有核心共享的,显著减少了缓存中的冗余。类似地,在 HiCache 中,L1 和 L2 KV cache 是每个推理实例专用的,而 L3 KV cache 则在集群内的所有推理实例之间共享。

HiRadixTree:HiCache 中的元数据组织#

对于 KV cache 数据组织,HiCache 基于 RadixAttention 中引入的 RadixTree 结构,并提出了 HiRadixTree。在 RadixAttention 中,RadixTree 的每个节点对应 GPU 内存中连续令牌范围的 KV cache。从根节点到叶节点的路径代表一个请求的前缀,多个请求之间的共享前缀可以重用相同的节点,从而避免冗余存储。

HiRadixTree 扩展了这一理念:每个节点对应连续令牌范围的 KV cache,并记录该 KV cache 的存储位置——无论是在本地 GPU 内存、CPU 内存、L3 存储中,还是这些层级的组合。如果存储在本地,HiRadixTree 会维护精确的元数据,包括确切的存储地址。然而,为了减少开销,HiRadixTree 不会存储或持续同步 L3 KV cache 的元数据。相反,当访问 L3 数据时,它会实时查询后端以获取必要的元数据,例如数据是否存在以及它位于哪台服务器和位置。

整体工作流程#

HiCache 的工作流程主要涉及三个关键操作:本地匹配预取写回。当系统接收到新请求时,它首先搜索本地的 L1 和 L2 缓存以查找匹配的 KV cache。对于在本地未找到的部分,尝试从 L3 预取。预取完成后,所有必需的 KV cache 都被加载到 GPU 进行计算。预填充计算完成后,系统考虑将新生成的数据存储到 L2 或 L3。

HiCache 工作流程

本地匹配#

本地匹配是 HiCache 工作流程的第一步,在此步骤中,传入的请求令牌与 HiRadixTree 进行匹配,以定位本地内存层级(L1 GPU 内存和 L2 主机内存)中的缓存 KV 数据。

匹配算法从根节点遍历 HiRadixTree,跟随与令牌序列前缀匹配的子节点。在每个节点,传入的令牌序列与节点存储的令牌序列进行比较。当 page_size > 1 时,以页面粒度执行匹配以优化内存访问模式。如果匹配在节点存储的序列中终止,则该节点会被自动分割以创建精确边界,提高未来匹配的效率。

该算法返回请求的连续前缀,其中第一部分位于 L1,后一部分位于 L2。

由于该过程仅需要遍历本地 HiRadixTree 且不涉及任何实际数据复制,因此本地匹配极其快速。

从 L3 预取#

数据预取是 HiCache 的核心优化技术之一,旨在主动将 KV cache 从 L3 存储加载到本地 L2 内存,从而减少后续操作中的访问延迟。

预取触发条件: 本地匹配后,对于在 L1 或 L2 中未找到的部分,系统查询 L3 以获取下一个连续匹配的 KV cache 的元数据。如果 L3 中命中的缓存长度超过阈值(默认:256 个令牌,可配置),则触发预取操作。

预取策略:HiCache 提供三种不同的预取终止策略以满足不同场景需求:

  • best_effort:当 GPU 可以执行预填充计算时立即终止,无需等待时间,适合对延迟极其敏感的场景。

  • wait_complete:必须等待所有预取操作完成后才继续,适合需要高缓存命中率的场景。

  • timeout:在指定时间或完成后终止,平衡延迟和缓存命中率需求。

预取停止后,已获取的数据与本地数据一起用于预填充计算。

对于 timeout 策略,HiCache 引入了两个配置参数以支持对预取超时条件的细粒度控制:

  • prefetch_timeout_base:基础超时时间,代表与令牌数量无关的开销(例如调度和同步)。

  • prefetch_timeout_per_ki_token:每千令牌的增量超时时间。

超时时间计算如下:

timeout = prefetch_timeout_base + prefetch_timeout_per_ki_token * num_token_to_fetch / 1024

数据写回#

写回机制负责将频繁访问的 KV cache 从 L1 移动到 L2 和 L3,实现更大更长期的存储以及跨实例的缓存共享。

可配置的写回策略:HiCache 支持三种写回策略:

  • write_through:每次访问都立即写回到下一层级。当带宽充足时,此策略提供最强的缓存优势。

  • write_through_selective:仅在访问频率超过阈值后才写回数据。此策略仅备份热数据,减少 I/O 开销。

  • write_back:仅当数据从上层被驱逐时才写回到下一层级。此策略减轻存储压力,适合存储容量有限但必须最大化内存利用率的场景。

跨实例共享:当数据从 L2 写回到 L3 时,仅传输尚未存在于 L3 中的数据。存储在 L3 中的 KV cache 可以在集群中的所有 SGLang 实例之间共享(取决于 L3 后端实现),显著提高了相同内存预算内的缓存命中率。

多级同步#

在多 GPU 并行计算(如张量并行(TP))期间,HiCache 必须确保不同级(rank)之间的一致状态。因此,关键计算步骤需要使用 all_reduce 进行状态同步。

例如,在预取期间,使用 all_reduce(op=min) 确保所有级获得相同数量的 L3 命中,防止对是否达到预取阈值产生不一致的判断。类似地,预取完成或终止后,再次需要 all_reduce(op=min) 以确保所有级就成功检索的 KV cache 的前缀长度达成共识。

数据传输优化#

零拷贝数据传输:预取和写回都涉及大量数据移动。最小化数据复制数量可以显著提高系统性能。HiCache 在从 L2 内存传输数据到 L3 后端时支持直接传递内存地址和大小。

"面向批次"的数据组织:数据读写的粒度对性能有重大影响。为此,HiCache L3 以页面(pages)的粒度存储和传输 KV cache 数据,并支持现有 layer first 方案之外的不同数据布局,包括 page firstpage_first_direct。在 page firstpage first direct 布局下,属于同一页的所有 KV cache 数据被放置在连续内存中,允许作为单个对象使用零拷贝传输传递给 L3。

HiCache L2 内存布局

然而,因为 GPU KV 计算是天然按层执行的,GPU 本身采用 layer first 布局。当从 L2 传输 page first 数据到 GPU 时,必须以每层一个令牌的粒度传输数据。page first direct 布局通过将给定页面内的所有层令牌分组在一起,缓解了这一问题,允许从 L2 到 GPU 的传输在页面-层级级别进行聚合。

CPU 到 GPU 传输优化:在 HiCache 中,将数据从 CPU 内存移动到 GPU 与从 L3 预取数据到 L2 一样关键。HiCache 对此过程采用了几种优化:

  • 计算-传输重叠:在预填充阶段,当从 CPU 向 GPU 传输数据时,HiCache 通过在计算第 N 层的同时并发加载第 N+1 层的 KV cache 来实现层重叠。这有效地隐藏了数据传输延迟。

  • GPU 辅助 I/O 内核:在 cudaMemcpyAsync 基础上,HiCache 实现了一组专门为 CPU 和 GPU 之间 KV cache 传输优化的 GPU 辅助 I/O 内核。与基线方法相比,这些内核实现了高达 3 倍的传输速度。

MLA 写回优化:对于多 TP 下的 MHA(多头注意力)模型,每个级持有令牌 KV 数据的 1/tp_size。相比之下,对于 MLA(多层注意力)模型,所有级持有每个令牌完整且相同的 KV 数据。HiCache 包含专门针对 MLA 的优化:仅一个级发起写回操作,确保数据不会在多个级中冗余存储。

与 PD-解耦部署模式的集成#

SGLang 通过 Mooncake TransferEngine 支持 PD(预填充-解码)解耦部署模式(详见此文档)。在 PD 解耦部署模式下,可以在预填充节点和解码节点上都启用 HiCache 以优化预填充性能。如果在解码节点上启用,解码输出也会写回到 L3。

统一接口和丰富的 L3 存储后端#

HiCache 将在 L3 后端上的所有读取、写入和查询操作封装在 class HiCacheStorage(ABC) 中,提供一套简单一致的接口。这种设计支持多种 L3 存储后端,允许用户根据特定用例选择最适合的方案。

  • Mooncake:Mooncake 是一个高性能的 LLM 推理缓存系统,利用 RDMA 和多 NIC 资源实现零拷贝、超高速数据传输。在此处尝试 Mooncake。

  • DeepSeek 3FS (HF3FS):HF3FS 是一种 Kubernetes 原生的分布式存储解决方案,基于 operator 部署。在此处尝试 HF3FS。

  • NIXL:NIXL 提供统一的 API 访问各种存储插件,包括但不限于 DeepSeek 的 3FS、GPU Direct Storage (GDS) 和 Amazon S3 兼容的对象存储。在此处尝试 NIXL。

  • AIBrix KVCache:AIBrix KVCache 是一个生产就绪的 KVCache 卸载框架,实现高效的内存分层和低开销的跨引擎重用。在此处尝试 AIBrix KVCache。

  • HiCacheFile:一个用于演示目的的简单基于文件的存储后端。

具体而言,LMCache 是一个针对企业规模 LLM 推理的高效 KV cache 层,提供了 HiCache 的替代方案。在此处尝试 LMCache。

相关参数#

  • --enable-hierarchical-cache:启用分层缓存功能。使用 HiCache 需要此参数。

  • --hicache-ratio HICACHE_RATIO:主机 KV cache 内存池大小与设备池大小的比率。例如,值为 2 表示主机内存池是设备内存池的两倍。此参数的值必须大于 1,因为当前实现要求分配给 KV cache 的主机内存大于分配给 KV cache 的设备内存。

  • --hicache-size HICACHE_SIZE:主机 KV cache 内存池大小(以 GB 为单位)。如果设置了此参数,它将覆盖 hicache-ratio。例如,--hicache-size 30 为每个级分配 30GB(1GB = 1e9 字节)的主机内存池。如果有 8 个级,则总内存大小为 240GB。与 hicache-ratio 类似,此参数的值必须大于分配给 KV cache 的设备内存大小。

注意--hicache-ratio--hicache-size 是两个关键参数。一般来说,更大的 HiCache 大小会导致更高的缓存命中率,从而提高预填充性能。然而,缓存大小与命中率之间的关系不是线性的。一旦大多数可重用的 KV 数据——特别是热令牌——已经被缓存,进一步增加大小可能只会带来边际性能提升。用户可以根据其工作负载特征和性能要求设置这些参数。

  • --page-size PAGE_SIZE:每页的令牌数量。此参数决定 KV cache 存储和检索的粒度。较大的页面大小会减少元数据开销并提高存储后端的 I/O 效率,但当只有页面的一部分与存储的 KV cache 匹配时,可能会降低缓存命中率。对于具有长公共前缀的工作负载,较大的页面可以提高性能,而对于前缀更多样化的工作负载,较小的页面可能更有益。有关页面粒度如何影响 I/O 性能的详细信息,请参阅数据传输优化。

  • --hicache-storage-prefetch-policy {best_effort,wait_complete,timeout}:控制何时停止从存储预取。详情请参阅从 L3 预取。

    • best_effort:在不阻塞的情况下尽可能多地预取

    • wait_complete:等待预取完成后再继续

    • timeout:在指定时间或完成后终止(生产环境推荐,设置适当的超时有助于系统满足所需的 SLOs)

  • --hicache-write-policy {write_back,write_through,write_through_selective}:控制数据如何从较快的内存层级写入较慢的层级。详情请参阅数据写回。

    • write_through:立即将数据写入所有层级(最强的缓存优势)

    • write_through_selective:使用命中计数跟踪仅备份频繁访问的数据

    • write_back:仅在需要驱逐时将数据写回到较慢的层级(减少 I/O 负载)

  • --hicache-io-backend {direct,kernel}:选择 CPU 和 GPU 之间 KV cache 传输的 I/O 后端。详情请参阅数据传输优化。

    • direct:标准 CUDA 内存复制操作

    • kernel:GPU 辅助 I/O 内核(推荐用于更好的性能)

  • --hicache-mem-layout {layer_first,page_first,page_first_direct}:主机内存池的内存布局。详情请参阅数据传输优化。

    • layer_first:兼容 GPU 计算内核(GPU 内存默认)

    • page_first:针对 I/O 效率优化

    • page_first_direct:将给定页面的所有层令牌分组在一起,允许从 L2 到 GPU 的传输在页面-层级级别进行聚合

  • --hicache-storage-backend {file,mooncake,hf3fs,nixl,aibrix,dynamic}:选择 L3 层的存储后端。内置后端:file、mooncake、hf3fs、nixl、aibrix。对于动态后端,使用 --hicache-storage-backend-extra-config 指定:backend_name(自定义名称)、module_path(Python 模块路径)、class_name(后端类名)。有关可用后端,请参阅统一接口和丰富的 L3 存储后端。

  • --enable-lmcache:使用 LMCache 作为替代的分层缓存解决方案。

  • --hicache-storage-backend-extra-config HICACHE_STORAGE_BACKEND_EXTRA_CONFIG:包含存储后端额外配置的 JSON 字符串,例如 --hicache-storage-backend-extra-config '{"prefetch_threshold":512, "prefetch_timeout_base": 0.5, "prefetch_timeout_per_ki_token": 0.25}'