针对大模型推理中的kvcache，存储系统可以有那些工作？

KV Cache在大模型推理里确实是个大头，占内存、吃带宽，怎么优化它基本就决定了推理系统能跑多快、能服务多少用户。

从存储系统的角度看，能做的事情挺多的，比如：

最直接的——分层存储

GPU显存就那么点，动不动几十GB的KV Cache塞不下怎么办?最自然的想法就是把不常用的丢到CPU内存里。这个思路其实FlexGen在2023年就做过了,他们把GPU、CPU内存、甚至SSD都用上,搞了个三层的存储架构。但是这里有个问题就是要预测哪些KV Cache马上会用到,需要提前预取回GPU,这样才能把传输延迟藏起来。

更激进一点的做法是直接上NVMe SSD。对于那些超长上下文的场景,比如处理整本书或者一个完整的代码仓库,几百万个token的KV Cache真的没办法全放内存。这时候SSD就成了目前唯一的选择,配合GPU Direct Storage这种技术,跳过CPU直接在GPU和SSD之间传数据,速度其实也还凑活。

压缩则是另一个大方向

KV Cache本质上就是一堆浮点数,FP16占16位,但实际上很多时候用不着这么高精度。INT8甚至INT4量化在很多任务上精度损失很小,但存储能省一半甚至四分之三。Google的那篇KIVI就专门研究了这个问题,他们发现对Key做per-channel量化、对Value做per-token量化效果最好。

还有个更狠的思路是结构化压缩。注意力机制里其实有大量的稀疏性,很多token之间的attention score接近零,那些对应的KV Cache理论上可以直接扔掉。H2O好像有个这个工作就是基于这个发现,只保留"Heavy Hitters"那些重要的KV,其他的直接prune掉。但是这个跟剪枝一样,砍多了模型效果就崩了。

共享复用

同样的system prompt，每个请求都单独存一份KV Cache完全是浪费。vLLM提出的PagedAttention就做了这个事情,把KV Cache分成固定大小的块,不同请求可以共享相同的prefix块。这个设计借鉴了操作系统的虚拟内存思想。

SGLang更进一步,搞了个RadixAttention,用前缀树来管理所有的KV Cache。当新请求来的时候,先在树里找有没有匹配的前缀,有的话直接复用,没有的话再生成新的。这对多轮对话场景特别有用,因为历史对话都是可以共享的。

长上下文是最大的问题

现在的模型动不动就支持100K、1M token的上下文,这对存储系统是个巨大考验。传统的连续内存分配根本hold不住,必须搞分块管理。

可以把KV Cache按照固定窗口切分,比如每4096个token一个块,然后像数据库一样建索引。需要访问某个位置的时候,通过索引快速定位到对应的块,加载进来就行。这个思路有点像Mamba那种状态空间模型的做法,虽然它不是transformer,但分块管理状态的思想是相通的。

流式处理也很重要。超长文本不可能一次性全塞进去,得边生成边写入,边prune掉不重要的部分。这需要存储系统支持高效的append和删除操作,传统的数组结构可能就不够用了,链表、B树这些数据结构可能更合适。

实际工程的细节

上面说的都是理论，但是真正做起来还有很多坑。比如内存碎片问题,KV Cache的size不固定,分配释放频繁,很容易产生碎片导致可用内存不连续。需要像Buddy System那样做内存管理,或者定期做compaction。

异步I/O也是必须的。不能因为等数据传输就把GPU晾在那里,得让计算和传输overlap起来。CUDA Stream这些机制用好了,能把传输开销隐藏掉大半。

监控也不能少。得知道每个请求的KV Cache占了多少空间,访问模式是什么样的,热点在哪里。有了这些数据才能做针对性优化,不然就是瞎调参。

总的来说,KV Cache的存储优化是个系统工程,需要算法、系统、硬件各层面的配合。现在这个领域还在快速发展,每隔几个月就有新的paper出来。