AI资讯 / 工程

工程 / 工程

在 Amazon SageMaker HyperPod 上为大型 LLM 构建分层 KV 缓存

AWS Machine Learning

在大规模部署大型语言模型推理时,KV 缓存始终面临两难困境:要么为不断增长的缓存需求购置高规格 GPU 实例,要么接受因重复计算相同提示词而导致的高首字延迟。AWS 机器学习团队在 Amazon SageMaker HyperPod 上构建了一套三层 KV 缓存架构,将缓存层级从 GPU 显存(L0)、本地 CPU 内存(L1)延伸至基于 Curvine 的共享分布式 NVMe 存储池(L2)。Curvine 是一款轻量级分布式缓存文件系统,能够将 G6e 或 P5 实例本地 NVMe 硬盘汇聚为统一命名空间,并以 FUSE 客户端方式挂载至每个推理 Pod。配合 HyperPod 内置的缓存感知路由策略,请求可被精准调度至最可能命中缓存的副本。实测结果显示,跨 Pod 缓存命中率最高可达 100%,首字延迟最多改善 2.7 倍,约 1900 个 token 的提示词跨节点 L2 读取延迟约为 56 毫秒,同时原本需要 P5 实例的工作负载可迁移至成本更低的 G6e 实例运行。

大型语言模型在推理阶段会将已处理 token 的注意力键值对存入 KV 缓存,避免每步重复计算。前缀缓存(Prefix Caching)进一步将这一机制扩展至共享相同前导 token 的跨请求复用场景,例如多个请求共用同一系统提示词。然而在 ml.g6e.4xlarge 这类成本效益较高的实例上,48 GB 显存在扣除模型权重和运行时分配后,留给前缀缓存的空间十分有限。随着模型规模增大或并发量提升,缓存命中率会显著下降,水平扩展的多个 vLLM 副本各自维护独立缓存,路由至不同副本实际上等同于冷启动。

为解决上述问题,AWS 团队设计了三层缓存架构。L0 层为 GPU 高带宽内存中的原生 vLLM 分页注意力缓存,访问延迟最低但容量受限于模型权重占用后的剩余显存。以 48 GB GPU 为例,7B 参数模型约占用 14 GB 权重,剩余空间充裕;而 32B 模型权重超过 64 GB,即便分片后可用于 KV 缓存的空间也极为有限,在高并发场景下缓存极易被驱逐。L1 层为本地 CPU 内存卸载层,由 LMCache 在 GPU 块被驱逐时将其捕获至主机 DRAM,通过 SageMaker HyperPod 推理算子自动管理,建议初始分配比例设为实例内存的 20%。

L2 层是实现跨副本缓存复用的关键所在。Curvine 将 G6e 或 P5 实例本地 NVMe 硬盘汇聚为单一命名空间,通过 FUSE 客户端以 ReadWriteMany PVC 形式挂载至每个推理 Pod。LMCache 通过 fs:// 连接器进行读写,分布式存储池对上层应用呈现为普通本地目录。由于所有 Pod 挂载同一命名空间,某个副本写入的 KV 块可立即被其他副本读取。Curvine 的主节点负责元数据管理与日志持久化,数据持久化依托 Amazon EBS;各 GPU 节点上的工作节点将数据存储于本地 NVMe。即便某个工作节点宕机,丢失的 KV 块也可通过重新计算恢复,不存在数据丢失风险。

仅有分层缓存还不够,请求必须被路由至已持有相关 KV 块的副本,缓存才能真正发挥价值。HyperPod 推理算子内置路由器支持三种策略:前缀感知路由(prefix-aware)适用于多轮对话和共享系统提示词场景;KV 感知路由(kv-aware)适合长文档处理和扩展会话;轮询路由(round-robin)则适用于无状态批量推理和负载测试。前缀感知路由通过维护前缀树选择最可能命中缓存的副本,KV 感知路由则查询每个工作节点的缓存状态进行决策,整个过程对客户端完全透明,无需任何客户端改动。

在具体实现层面,推理算子作为 Amazon EKS 插件安装,负责管理完整生命周期,包括启动带有 LMCache sidecar 的 vLLM Pod、配置 L1 和 L2 后端、部署路由器并暴露单一负载均衡端点。用户通过 InferenceEndpointConfig CRD 声明缓存拓扑,算子自动渲染对应的环境变量、卷挂载和路由规则。目前 CRD 的 l2CacheBackend 字段原生仅支持 redis 或 tieredstorage,若需将 L2 指向 Curvine FUSE 挂载点,需在 vLLM 容器规格中手动修补 LMCACHE_REMOTE_URL 环境变量,将其设置为 fs://localhost:0/mnt/curvine/l2cache/。

实测基准测试结果验证了该架构的实际效果:跨 Pod 缓存命中率最高达到 100%,首字延迟最多改善 2.7 倍,针对约 1900 个 token 的提示词,跨节点 L2 读取延迟约为 56 毫秒。从成本角度看,原本需要 P5 实例才能承载的工作负载,现在可以迁移至成本更低的 G6e 实例运行,实际节省幅度取决于模型规模和流量特征。对于部署 Qwen、Llama、DeepSeek 等主流开源基础模型的团队,以及运行 RAG 流水线或多轮对话应用的场景,这套分层缓存方案提供了一条在不牺牲用户体验的前提下控制基础设施成本的可行路径。

要点

  • 三层 KV 缓存架构(GPU 显存 L0、CPU 内存 L1、共享 NVMe L2)可将缓存复用从单副本扩展至跨节点范围,解决水平扩展场景下各副本缓存孤岛问题
  • Curvine 通过 FUSE 客户端将多节点本地 NVMe 汇聚为统一命名空间,使所有推理 Pod 共享同一 L2 缓存层,实测跨节点读取延迟约 56 毫秒
  • HyperPod 内置的缓存感知路由策略(前缀感知、KV 感知、轮询)可将请求精准调度至最可能命中缓存的副本,无需客户端改动
  • 实测首字延迟最多改善 2.7 倍,跨 Pod 缓存命中率最高达 100%,部分原需 P5 实例的工作负载可迁移至成本更低的 G6e 实例
  • 当前实现需手动修补 LMCACHE_REMOTE_URL 环境变量以将 L2 指向 Curvine 挂载点,该限制源于 CRD 原生后端支持范围
查看原始来源

原始标题:Tiered KV cache for large LLMs on Amazon SageMaker HyperPod with Curvine

本文由 DataHub 基于公开来源整理,用于信息发现与摘要阅读;具体事实、数据和后续更新以原始来源为准。