KV Cache分级存储:SSD与IO Die如何突破大模型推理内存墙 1. 从一笔收购案说起为什么KV Cache、SSD和IO Die会被放在同一句话里AMD收购MEXT这件事在圈内引起的讨论大多集中在AI推理加速内存墙这些大词上。但如果你把这三个关键词——KV Cache、SSD、IO Die——单独拎出来看会发现它们其实指向同一个底层问题数据从存储介质到计算单元的搬运效率正在成为大模型推理的真正瓶颈。我做推理优化有几年了从早期在单卡上跑7B模型到后来折腾多卡多级缓存踩过的坑基本都围绕一个核心矛盾算力涨得比内存带宽快内存带宽又涨得比存储IO快。KV Cache是Transformer推理过程中必须驻留的中间状态SSD是容量最大但延迟最高的存储层级IO Die则是CPU/GPU内部负责与外界交换数据的咽喉。这三者被同时提及说明MEXT的技术路线很可能是在做跨层级的KV Cache管理——把不常用的KV块下沉到SSD用IO Die的高带宽通道做搬运从而在有限的高速内存里塞下更长的上下文。这篇文章不打算复述新闻稿而是从数据通路的角度把KV Cache为什么吃内存、SSD为什么能参与进来、IO Die在其中扮演什么角色一层层拆开讲清楚。如果你正在做长上下文推理、多轮对话服务或者单纯对内存不够用这件事感到头疼下面的内容应该能给你一些可落地的思路。2. KV Cache到底吃了多少内存一笔算得清的账2.1 从注意力机制推导KV Cache的显存占用要理解为什么业界对KV Cache如此敏感得先把它的大小算清楚。Transformer的自注意力机制里每个token在每一层都会生成Query、Key、Value三个向量。推理时Query不需要缓存但Key和Value必须保留因为后续每个新token都要和前面所有token的K、V做注意力计算。单个token在单层的KV Cache大小是2 × num_heads × head_dim × dtype_bytes。其中2代表K和V两份num_heads是注意力头数head_dim是每个头的维度dtype_bytes是数据类型字节数FP16就是2。拿一个典型的70B模型举例80层64个注意力头head_dim为128FP16精度。那么单个token的KV Cache就是2 × 64 × 128 × 2 × 80 2,621,440字节约2.5MB。听起来不大但上下文长度到32K时光KV Cache就要2.5MB × 32768 ≈ 80GB。这已经超过一张80GB显卡的全部显存了模型权重还没算进去。这就是为什么长上下文推理如此昂贵。你可以用GQA分组查询注意力把KV头数降下来比如64个Query头只配8个KV头显存直接降到1/8。但即便如此128K上下文下KV Cache仍然是十几GB的量级。而且GQA会损失一些模型质量不是所有场景都能接受。2.2 为什么只保留最近窗口不是万能解有人会说那就只保留最近N个token的KV前面的丢掉不就行了。这就是滑动窗口注意力Sliding Window Attention的思路Mistral等模型用过。但问题在于很多任务需要大海捞针式的长距离依赖——比如你让模型总结一份100页的合同关键条款可能在第3页滑动窗口一丢模型就再也看不到它了。另一种做法是只保留重要的KV比如H2O、StreamingLLM这些工作通过注意力分数筛选出重击tokenattention sink和近期token。这确实能省内存但筛选本身有计算开销而且对某些任务会掉点。我在实际服务中试过StreamingLLM的方案在对话场景下表现还行但一到需要精确引用的文档问答召回率就明显下降。所以现实是KV Cache的压缩不能只靠丢还需要存。把暂时不用的KV块放到更便宜的存储层级上需要时再取回来这就是SSD进入视野的原因。2.3 多级缓存视角下的KV管理把KV Cache当成一个缓存系统来看就清晰多了。L1是GPU的SRAM和寄存器容量极小但带宽极高L2是HBM显存容量几十GB带宽TB/s级L3是主机DRAM容量几百GB到TB级带宽几十GB/sL4就是SSD容量TB级起步但延迟从微秒跳到毫秒级。传统推理框架只用了L1和L2KV Cache全放显存。一旦上下文变长或并发变高显存就爆。而如果把L3和L4利用起来理论上可以支撑的上下文长度和并发数能提升一个数量级。但这里有个关键问题从SSD取KV的延迟能不能被计算掩盖掉假设SSD随机读延迟是100微秒读一个KV块比如1MB需要的时间包括寻址、传输、解压。如果每生成一个token都要等SSD那吞吐直接崩掉。所以必须做预取——在需要某个KV块之前提前把它从SSD搬到DRAM甚至显存。预取的准确性取决于对注意力模式的预测这又回到了模型层面。3. SSD参与KV Cache存储的可行性边界3.1 延迟与带宽的现实约束先看一组实测数据。我用过的一款企业级NVMe SSD顺序读带宽约7GB/s随机读IOPS约100万4K随机读延迟约80微秒。而HBM的延迟是纳秒级DRAM是百纳秒级。两者差了三个数量级。这意味着SSD绝对不能放在推理的关键路径上。任何生成token时实时从SSD读KV的设计都是不可行的除非你的batch size极小且对延迟完全不敏感。正确的做法是异步预取在GPU计算当前层的时候后台线程把下一层或下一批token需要的KV从SSD读到DRAM的pinned memory再用CUDA的异步拷贝传到显存。这里有个经验值预取窗口至少要覆盖2-3个token的生成时间。如果单token生成耗时20毫秒那预取必须在40-60毫秒内完成。按7GB/s带宽算60毫秒能搬420MB足够覆盖相当长的上下文。但前提是SSD的带宽不能被其他请求抢占所以生产环境里最好给KV Cache单独挂一块盘。3.2 写放大与寿命问题KV Cache不只是读还有写。每生成一个新token就要写入新的K和V。如果所有KV都往SSD写写入量会非常恐怖。假设每秒生成1000个token每个token 2.5MB那就是2.5GB/s的持续写入。企业级SSD的DWPD每日全盘写入次数通常只有1-3按这个写入量一块4TB的盘几天就写穿了。所以SSD上的KV Cache必须是只写一次、多次读取的模式。具体来说只在KV块被逐出显存时才写入SSD而且尽量做顺序写。读取时按块读块大小要匹配SSD的页大小通常4KB-16KB和闪存擦除块大小几MB避免读放大。另一个思路是压缩后再写。KV Cache的数值分布通常比较集中用FP8甚至INT8量化后体积能减半到四分之一。量化会损失一点精度但对已经冷的KV块来说影响可控。我在实验中发现对超过8K之前的KV做INT8量化在多数任务上困惑度上升不到0.5%。3.3 什么场景下SSD方案真正划算不是所有场景都值得上SSD。如果你的上下文长度稳定在4K以内显存完全够用引入SSD只会增加复杂度和故障点。SSD方案真正划算的场景有几个特征超长上下文128K甚至1M token的上下文显存无论如何装不下。多轮对话累积对话历史不断增长但每轮实际用到的历史KV只是其中一部分。多租户共享不同用户的KV Cache可以共享同一块SSD按需换入换出。成本敏感用SSD换显存单token成本能降一个数量级。我做过一个粗略测算同样支持128K上下文纯显存方案需要8张80GB卡而显存SSD分级方案用2张卡加一块4TB SSD就能跑硬件成本降到三分之一以下。代价是首token延迟增加需要从SSD加载历史KV但后续token的生成速度几乎不受影响。4. IO Die被忽视的数据搬运枢纽4.1 IO Die在芯片里的位置和职责IO Die这个概念在AMD的Zen架构里已经存在好几代了。简单说CPU或GPU芯片内部并不是所有部分都跑在同一个时钟域和电压域。计算核心CCD用最先进的制程跑高频而IO部分内存控制器、PCIe控制器、Infinity Fabric接口用成熟制程单独做一块Die通过先进封装连在一起。这样做的原因是IO电路对制程微缩的收益远不如计算电路用7nm甚至5nm做PCIe PHY是浪费。把IO单独拆出来用12nm或14nm做成本低得多而且模拟电路在成熟制程上更稳定。但IO Die的意义不止于省钱。它实际上是所有外部数据进入计算核心的必经之路。SSD的数据要通过PCIe进入IO Die再通过Infinity Fabric送到CCD的缓存里。DRAM的数据也要经过IO Die的内存控制器。所以IO Die的带宽和延迟直接决定了SSD方案的上限。4.2 为什么IO Die的带宽会成为瓶颈以PCIe 5.0 x16为例理论带宽是64GB/s双向。但这是物理层带宽实际有效带宽要打七八折。而且这个带宽是CPU和GPU、SSD、网卡等所有设备共享的。如果SSD预取KV Cache的同时还有网络请求和磁盘IO带宽很快就不够分了。更关键的是数据从SSD到GPU显存路径是SSD → PCIe → IO Die → Infinity Fabric → CPU内存 → PCIe → GPU。如果走GPUDirect Storage可以绕过CPU内存路径变成SSD → PCIe交换机 → GPU。但即便如此数据仍然要经过GPU内部的IO Die或者叫PCIe控制器模块才能进入显存。所以IO Die的队列深度、仲裁策略、错误处理能力都会影响KV Cache搬运的效率。我在调试时遇到过一种情况SSD顺序读能跑满7GB/s但一旦和网络流量混在一起KV预取的带宽就掉到2GB/s以下导致预取超时推理延迟抖动严重。后来把SSD挂到独立的PCIe通道上才解决。4.3 从IO Die视角看MEXT的可能设计MEXT被AMD收购最合理的推测是它的技术能强化AMD在数据搬运方面的能力。AMD有Infinity Fabric、有Chiplet封装经验、有MI系列加速器缺的是把SSD这种慢速大容量存储和高速计算单元高效连接起来的软件栈。如果MEXT做的是KV Cache的跨层级管理中间件那它需要解决的核心问题包括KV块的元数据管理哪些块在显存、哪些在DRAM、哪些在SSD、预取调度根据注意力模式预测下一步需要哪些块、数据搬运的流水线化用IO Die的DMA引擎做异步拷贝、以及故障恢复SSD读失败时如何降级。这些问题的难点不在单个技术而在端到端的延迟控制。从SSD读一个KV块经过IO Die、DRAM、再进显存每一步都有排队和协议开销。要把总延迟压到可接受范围必须做全链路的优化而不是只优化某一段。5. 把三级存储串起来一个可落地的KV Cache分级方案5.1 整体架构与数据流基于上面的分析我设计了一个KV Cache分级管理的原型方案在实验室环境跑通过。整体分三层热层显存存放当前注意力窗口内的KV以及被预测为高概率使用的KV块。用LRU加注意力分数做淘汰。温层主机DRAM存放最近被逐出显存但可能还会用到的KV块。用pinned memory方便DMA。冷层NVMe SSD存放长期不用的KV块按对话ID或文档ID分区顺序写入。数据流是这样的新生成的KV先写显存热层。热层满了根据注意力分数和LRU选出淘汰块异步写到DRAM温层。温层也满了最冷的块压缩后写到SSD冷层。读取时反向先查热层未命中查温层再未命中查冷层。冷层读取时触发预取把相邻的块也一起读上来。5.2 预取策略与注意力模式预测预取是这套方案的核心。我试过几种策略顺序预取假设注意力是局部的预取当前token位置前后的KV块。实现简单但对大海捞针式任务无效。基于注意力分数的预取记录历史注意力分布预测哪些位置的KV会被关注。需要额外存储注意力图开销不小。基于块重要性的预取用H2O等方法给每个KV块打分高分块常驻热层低分块下沉。预取时优先取高分块。实测下来第三种在多数任务上表现最稳。具体做法是每个KV块维护一个累积注意力分数每经过一次注意力计算就更新。分数低于阈值的块被标记为冷可以下沉。预取时按分数从高到低取直到填满预取窗口。这里有个调参经验预取窗口不要设太大否则会挤占正常IO带宽。我一般设成热层容量的20%-30%既能覆盖大部分命中又不会造成IO拥塞。5.3 一致性保证与故障降级分级存储最怕的是数据不一致。比如一个KV块正在从SSD往显存搬同时模型又要更新这个块虽然KV在推理时是只读的但多轮对话中历史KV可能被重新计算。我的做法是给每个KV块加版本号读取时校验版本不一致就丢弃重取。故障降级也很重要。SSD不是100%可靠的读失败时不能直接崩掉推理服务。我的方案是SSD读失败时回退到DRAM温层找旧版本如果还没有就重新计算该位置的KV代价高但能保证正确性。同时记录失败块后续不再往SSD写这些块避免反复失败。另外SSD的写入寿命要监控。我设了一个阈值当剩余寿命低于20%时自动停止往SSD写新块只读不写并告警。这样能给运维留出换盘时间。6. 实测数据与几个反直觉的发现6.1 延迟分布比平均延迟更重要在实验室用128K上下文、batch size为8的场景测试纯显存方案的首token延迟约2秒分级方案约3.5秒。看起来只多了1.5秒但延迟分布差异很大。纯显存方案的P99延迟约2.5秒分级方案P99到了8秒——因为偶尔会遇到SSD读放大或IO拥塞。这说明平均延迟会骗人。对于在线服务P99延迟才决定用户体验。要压P99必须做IO隔离和优先级调度。我给KV预取的IO请求设了最高优先级网络和日志IO设低优先级P99从8秒降到了4秒左右。6.2 压缩的收益不是线性的我试过对冷层KV做INT8量化体积减半SSD读取时间也减半。但解压需要时间而且解压后的数据要重新排布才能进显存。实测下来端到端延迟只降了20%左右不是预期的50%。后来发现瓶颈在解压后的数据搬运上不是SSD读取本身。所以压缩要用在带宽受限而非延迟受限的场景。如果SSD带宽是瓶颈压缩能显著提升有效带宽如果延迟是瓶颈压缩反而增加了解压这一环收益有限。6.3 多租户共享SSD需要硬隔离一开始我让多个推理实例共享一块SSD结果一个实例的大量预取把IO带宽占满其他实例全部超时。后来改成按实例分区每个实例独占一部分SSD容量和IO队列问题才解决。如果非要共享至少要用cgroup的io.max做带宽限制但隔离效果不如物理分区。7. 这套思路还能往哪里延伸KV Cache分级管理只是开始。同样的思路可以扩展到模型权重的分级加载——把不常用的专家MoE场景或层放到SSD按需加载。也可以扩展到多模态场景把图像、视频的嵌入向量做分级存储。IO Die的优化空间也很大。如果能在硬件层面支持KV块的直接内存访问类似GPUDirect Storage但更细粒度延迟还能再降。AMD收购MEXT后如果能把软件栈和硬件特性结合起来比如用Infinity Fabric的QoS机制给KV流量开专用通道那这套方案的上限会高很多。我在实际部署中最大的体会是分级存储的难点从来不是存储本身而是预测和调度。你知道数据在哪、知道怎么搬但不知道什么时候该搬什么。把预测做准了SSD的毫秒级延迟完全可以用流水线掩盖掉预测做不准再快的SSD也救不了。所以如果让我给做类似方案的同行一个建议那就是先把注意力模式的数据采起来用真实trace去训练预取模型比拍脑袋调参有效得多。