当大语言模型(LLM)走向千行百业,推理效率与显存成本的矛盾日益尖锐。
KV Cache (Key-Value Cache)作为提升生成速度的核心技术,却像一个 " 吞存巨兽 " —— 每增加一个 token,就需要更多显存存储键(Key)和值(Value)向量,最终成为制约模型规模扩张、并发能力提升的 " 紧箍咒 " 如何高效、经济地扩展 KV Cache 存储空间,已成为全球 AI 产研界亟待攻克的高地。
业界探索
Nvidia 开源的 Dynamo 项目,实现存储系统多级缓存算法,热数据在显存、温数据在主机内存、冷数据在 SSD 或远端对象存储,并通过一套统一的索引 + 异步流水线实现自动迁移与透明访问,但是多级存储之间的数据迁移流程复杂,延迟开销难以压缩。
微软推出的 LMCahce 存储系统,高度兼容 vLLM 等推理框架,但是对分布式存储支持较低,空间上限低。
阿里巴巴提出一种将 KV Cache 空间扩展到 Tair 数据库的远端存储方案,存储空间易扩展,但是读写性能难以满足 LLM 推理业务的低延迟需求。
业界关于 CXL 存储加速 LLM 推理的研究仍然较少,探索如何利用 CXL 等新型介质扩展 KV Cache 空间,进而将成熟的软件栈迁移到 CXL 硬件场景,是一项非常有意义的工作。
中兴通讯 Mariana 探索
中兴通讯公司和华东师范大学的研究团队联合提出了一种名为 Mariana(马里亚纳)的分布式共享 KV 存储技术,该工作《Mariana: Exploring Native SkipList Index Design for Disaggregated Memory》论文发表在 IEEE Transactions on Parallel and Distributed Systems ( TPDS ) 。
Mariana 作为一个面向计算 - 内存分离架构的高性能分布式 KV 索引,可以完美适配大模型推理场景的 GPU、KV Cache 存储,通过把节点锁细化到条目级、用自适应分裂 / 合并抑制热点写放大,并以 SIMD 友好的数据布局和轻量级热点缓存,实现比现有方案高 1.7 倍吞吐、尾延迟降低 23% 的纯有序索引。
Mariana 解决了 " 如何在高并发、热点倾斜、延迟敏感的高性能计算场景,让 KV 索引同时做到高吞吐、低延迟、可扩展?" 的问题,为 KV Cache 存储提供新的优化方案。
1. 细粒度并发控制方案(MSCS - Multi-Slot lock-based Concurrency Scheme)
将并发控制的粒度从节点级下降到条目级。在每个叶子节点内预留多个槽位 Slots 和对应的闩锁 Latch,写入操作时,客户端通过 RDMA_CAS 竞争一个空闲槽位的闩锁,而不是竞争整个节点的锁,实现同一个节点的不同槽位并发写入,极大地减少了竞争。叶子节点支持动态调整叶子数据范围、自动扩缩,将热点区域隔离以减少后续该节点上的竞争,显著降低了写密集型和高偏斜(skewed)工作负载下的争用,提高吞吐量同时降低了尾延迟。
采用分离式存储,Key 在内存中连续存放,value 则与校验和在另一内存块连续存储;连续键数组可一次性加载至 SIMD 寄存器,大幅提升查找速度;优化读写操作序列,避免多次 RDMA 读写指令。
提出一种能够快速感知并缓存变化热点的轻量级缓存机制。使用 Count-Min Sketch 算法及时识别出热点数据,维护一个按热度排序的链表。计算节点缓存热点 L1 节点的最小键及存储地址,大幅降低热点数据加载延迟,提升存储系统的读性能。
MARIANA 的创新设计与大模型 KV Cache 需求高度契合:
大容量存储:Mariana 的解耦内存架构,支持将数据分布在远端 CPU DRAM 甚至 PMem/SSD 组成的共享内存池内存节点池,理论上存储空间没有上限。
高吞吐:Mariana 通过硬件加速(SIMD) 和智能缓存(元数据缓存) 两种方式,显著减少了完成一次 KV Cache 查找所需的计算和网络开销,极大地提升了读吞吐量。
低延迟读:Mariana 的整个数据路径(本地缓存 ->RDMA->SIMD 搜索)都是为低延迟而优化的。其延迟远低于需要经过远端 CPU 协议栈的解决方案,能够满足推理流水线的严格延迟要求。
水平扩展:Mariana 通过去中心化的细粒度并发控制实现了良好的水平扩展能力,非常适合作为大规模分布式推理集群的存储基础。
基于 vLLM 框架搭建大模型推理应用,验证关闭 KV Cache、开启显存 KV Cache、开启 Mariana 多级存储几个场景的推理性能,配置足够多的并发测试请求,GPU 显存空间有限,只能存放 50% 的 KV 数据。
拥抱 CXL 新生态
Mariana 的强大之处在于其设计理念与底层硬件解耦。其核心算法无需重新设计,即可从 RDMA 网络平滑迁移至未来的 CXL 硬件生态。仅需将远程访问 API 替换为 CXL.mem 操作,便能充分利用 CXL 的低延迟和一致性优势,成为构建下一代大模型推理基础设施的坚实基石。
从优化细粒度并发控制到适配 CXL 新生态,Mariana 的突破不仅是一次技术创新,更重新定义了大模型推理的存储逻辑——当显存不再是不可逾越的壁垒,当分布式存储能在高吞吐与低延迟间找到完美平衡点,大模型的规模化应用将迎来真正的普惠时代。
或许在不久的将来,随着 CXL 技术的成熟与 Mariana 这样的技术方案的落地,我们会看到百亿 / 千亿参数模型在普通硬件上高效运行,让 AI 的算力红利真正渗透到每一个需要它的场景中。
* 本文系量子位获授权刊载,观点仅为原作者所有。
一键三连「点赞」「转发」「小心心」
欢迎在评论区留下你的想法!
— 完 —
点亮星标
科技前沿进展每日见