ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

稀疏注意力落地新拼图:HiSparse 分层 KV Cache 管理实践

稀疏注意力落地新拼图:HiSparse 分层 KV Cache 管理实践 做 LLM 推理优化的人最近半年聊得最多的已经不只是算子融合而是 KV cache。我自己的体感非常直接同样是 7B 模型输出长度从 2K 拉到 32K单请求的显存占用能差出一个数量级吞吐也跟着肉眼可见地往下掉。HiSparse 这个工作标题写得很清楚Sparse-Attention Decoding 配上 Hierarchical KV Cache Management。换句话说它尝试把“注意力算得少”和“缓存存得对”这两件事放到同一个系统里一起解决。我读完的第一反应是这才是稀疏注意力落地时真正该补的那块拼图。如果你只把注意力稀疏化当成“少算几个点”大概率会在工程上踩坑因为 GPU 解码阶段的制约瓶颈根本不在计算而在数据搬运。HiSparse 的思路是先把 KV cache 当成一个有层级、有索引、有热度统计的资源系统来管理再让稀疏注意力产生的 block 级访问模式反过来指导缓存迁移。这篇文章我就从工程实现的角度把问题、设计逻辑和落地要点一次讲透。1. 先看清楚瓶颈解码阶段到底被什么卡住1.1 一次只算一个 token但要把全上下文搬来搬去自回归解码的本质是一个串行循环每个 step 只生成一个 token然后把新 token 追加到序列后面继续做下一轮。这个流程的 FLOPs 其实不高但它每一步都要把当前序列所有 KV cache 从显存读到计算单元里和当前 query 做 attention。用一组典型参数算个账模型 32 层、8 个 KV head、head dim 128上下文 4096 tokenKV cache 用 FP16 存。单个解码 step 至少要读取的 KV 数据量是每层 KV 4096 * 8 * 128 * 2 * 2 (K 和 V) 16.8 MB 全部层 16.8 MB * 32 536 MB但这一步实际要做的矩阵乘只有一次 query维度 128和所有 KV 相乘计算量撑死几十到几百 MFLOPs。以 A100 的 2TB/s 左右 HBM 带宽来算光把 536MB 搬进片上就要大约 0.27ms而执行这些操作的算术强度极低GPU 的算力基本处于“饿肚子”状态。这就是经典的 memory-bound 场景解码性能上限写死在带宽上。这也是为什么很多人一开始用稀疏注意力目标并不是省 FLOPs而是减少“需要搬运的数据量”。只要能从 4096 个 token 里挑出 512 个相关 token 来做 attention理论上读取量就降一个数量级。1.2 KV cache 不只是“大”它的访问方式才是关键KV cache 占显存这个大家都能理解。更麻烦的是它的访问方式在标准实现里KV cache 是一块连续增长的缓冲区每个 token 的 K、V 按序排布。服务端推理时不同请求长度不一样做不到简单复用要么提前预留 max_len 的空间要么像 vLLM 那样用 paged attention 做分块。做了分块之后情况变成GPU 上有一个全局 block table记录每个逻辑 block 对应的物理地址。解码时 attention kernel 需要根据 block table 找到对应 KV 块读进共享内存计算。但这里有个新问题如果是标准 dense attention所有 block 都要读按顺序遍历就行一旦换成 sparse attention每个解码 step 只访问少数几个 block而且这些 block 在序列上的位置可能非常分散。于是系统从“顺序扫一遍数组”变成了“按一堆随机地址去查表、去搬数据”这个开销如果不管控稀疏带来的收益会被随机访存吃掉一大半。1.3 稀疏注意力减少计算却引入了新的缓存问题我们试过几种稀疏注意力方案滑动窗口、stride 模式、按注意力分数 top-k 选择、检索式稀疏等。每一种都在“少算”这件事上有效但工程落地时都会面对同一个尴尬注意力内核在 GPU 上决定“我要哪些 block”可 KV cache 的物理布局不一定配合。比如按 top-k 选择 block如果选中的 block 在显存里分布在不同的页、甚至部分已经被换到 CPU 内存那内核要么只能处理“命中部分”要么得把缺失 block 的加载过程暴露在关键路径上。HiSparse 把这个问题总结成“缓存管理必须理解稀疏注意力的访问模式”而不是继续把 KV cache 当作一个单纯的大数组来分配和释放。2. HiSparse 的设计主线让缓存管理跟着注意力走2.1 从“连续数组”变成“带索引的块设备”传统 KV cache 管理更像一个大 vectorappend 一个新 token 就把数据写在末尾释放时整块回收。而 HiSparse 的做法更像在维护一个微型数据库或者一个分级存储系统的 page cache。KV cache 被切割成固定大小的 block每个 block 有虚拟地址、物理地址、所在层级、访问热度、迁移状态等元数据。这听起来和 vLLM 的 paged KV cache 很像但有个本质差异vLLM 的分页主要解决“显存碎片”和“按请求分配”的问题注意力仍是 dense 的所以每个 block 都要访问只是物理位置灵活一点。HiSparse 的假设前提是稀疏注意力每个 step 只访问一部分 block那么缓存管理的目标就从“保证所有 block 都能被快速访问”变成“保证被访问的这部分 block 尽量在最快层级里”。这个转变很关键。一旦允许部分 block“不在最快层级”你就有空间做很多优化冷 block 可以挪到 CPU 内存更冷的甚至可以落盘。而这一切能不能成立取决于两个问题第一稀疏注意力选择的 block 有没有时间局部性第二缓存迁移的开销能不能被调度和预取覆盖。2.2 三层缓存与一张“热度地图”HiSparse 这类方案一般会把 KV block 按访问热度放到不同层级。我在实际工程里常分成三层L1GPU HBM也就是默认显存放当前正在被频繁访问的热 block要求低延迟、随机访问友好L2同节点 CPU 内存或者通过 NVLink/PCIe 能访问的扩展显存放温 block访问代价在几十到几百微秒级别L3主机 SSD 或者远端内存放冷 block基本只在大跨度回溯时才会用到。真正支撑这套体系的是一个“热度地图”它的度量单位是 block。每个 KV block 维护一个访问计数每次解码 step 结束后稀疏注意力内核会返回本次实际访问了哪些 block。这些命中信息会以带衰减的方式累加到 block 上频繁访问的 block 热度上升长期不访问的热度下降。热度地图的作用有两个。一是做“位置决策”热度高的 block 提升到 L1热度低的降级到 L2/L3。二是做“预取决策”热度在上升期的 block 提前往上搬避免真正要用的时候才发现它还在冷层。2.3 为什么分层能省显存而不是换一种碎片有人可能觉得把 KV cache 搬到 CPU 内存只是换了个地方存GPU 显存压力确实小了但访问延迟不是更高了吗答案在于“解码过程是批量并发的”。服务端推理很少只有一个请求在跑。通常一个 batch 里有几十个请求每个请求都有自己的 KV cache。如果每个请求的完整 KV cache 都必须待在 GPU 显存里显存很快就不够用最终只能限制 batch size吞吐自然上不去。而分层缓存允许你只把每个请求“当前正在高频访问”的那部分 KV block 放 GPU 上其余放远端。这样同一个显存容量可以服务更多请求虽然少数 block 需要偶尔从 L2/L3 搬回但这个搬移动作可以并行做并且摊到大量请求后有很好的平均收益。本质上这是用“更聪明的放置”换“更大的 batch”再拿“异步迁移”换“显存空间”。我觉得这才是分层管理的核心价值它优化的不是单请求延迟而是整个推理服务在长上下文场景下的吞吐密度。3. 分层管理的核心机制拆解3.1 块粒度选多大16、32 还是 64KV block 粒度是整个缓存系统的地基选错后面全难受。粒度太小比如 8 个 token 一块稀疏选择更灵活但 block table 会非常大而且每次迁移的单位也小传输效率低。粒度太大比如 128 个 token 一块block 数量少、元数据省了但一个 block 里只有个别 token 被 attention 选中时你也得整块加载稀疏性收益被稀释。工程上比较常见的选择是 16 或 32 个 token。假设 head dim 128、KV head 8、FP1632 个 token 的一个 KV block 体积大约是32 * 8 * 128 * 2 * 2 128 KB加上元信息后一个 block 的迁移量在几百 KB 级别用现代 PCIe 或 NVLink 传输延迟可控。同时32 token 的粒度在 top-k 稀疏注意力里基本能保留足够的选取精度不会误伤太多邻接 token 的局部性。我自己的经验是先默认 32如果观察到的稀疏选择非常离散、命中率偏低再降到 16 试试。3.2 两级索引虚拟 Block 到物理 Page既然 KV cache 分块并且分布在多个层级就必须有一套索引结构把逻辑上的 block 编号映射到“它在哪一层、物理地址是多少”。HiSparse 这类系统通常用两级索引全局 Block Table记录每个虚拟 block 的当前状态包括 level、物理 page id、迁移状态、热度计数。每层 Page Table对应层内部的分配器管理该层级里物理页的分配和释放。我把它类比成 OS 的页表虚拟地址连续不代表物理页连续每次访问前先查页表得到物理地址。在 GPU kernel 里直接查全局表往往太慢所以一般会先把当前 active 请求的 block 映射表同步到 GPU 上的一个小表里。这个同步不需要每个 step 都做只需要在 block 发生迁移、映射变化时更新控制好频率开销很小。为了防止索引本身占太多显存block table 可以用紧凑的固定大小结构体比如 16 字节一个条目8 字节物理指针、4 字节 level 和状态标志、4 字节热度计数的整数近似值。这样 100 万个 block 的索引大约是 16MB完全可接受。3.3 命中、预取和迁移的配合分层缓存真正难的部分不是数据结构而是迁移调度。默认做法是每个解码 step 结束后内核把本次缺失的 block 列表交给 CPU 侧的调度器调度器为这些缺失项发起点对点异步拷贝。注意顺序很重要先做“缺失即迁移”再做“热度驱动的预取”。缺失迁移是硬需求step N 需要的 block 不在 GPU 上就得等它搬上来否则这轮 attention 做不了。这种情况下能优化的地方在于“提前”如果预测到 step N1、N2 可能也会用到同一批邻居 block就在 step N 时顺手把它们也拉上来这就是预取。预取不能拍脑袋做。最可靠的特征是解码序列的“时序局部性”一个 block 被访问后紧挨着它的若干 block 在接下来几个 step 被访问的概率会升高。所以比较稳的预取策略是每次访问某个 block 时把它后续 2 到 4 个相邻 block 也标记为“预取候选”。另一个特征是请求级共享。比如多请求共享同一个系统 prompt或者同一批 beam 之间共享前缀那么一个请求访问到某个前缀 block 时其他请求大概率也会访问。这类复用信息可以从 block table 的 entry 上加一个“共享引用计数”来识别优先级最高。3.4 淘汰策略不能只看 LRU还要看注意力很多人一上来就套 LRU按最近访问时间往后淘汰。但我在实验里发现在稀疏注意力场景下纯 LRU 表现很差。原因是有些 KV block 虽然最近被访问过但它在注意力分数里的权重极低属于“顺手读了一下但基本没什么用”的情况而有些 block 虽然最近几轮没碰但只要一碰就会提供关键历史信息。更实用的淘汰计分是综合两个维度访问频次和注意力权重。具体做法是给每个 block 维护一个“累计注意力热度”每次被选中时把注意力分数加进去再做时间衰减。淘汰时优先淘汰“累计热度低且最近未访问”的 block而不是简单地看 LRU 时间戳。注意淘汰前要确认这个 block 是否还有别的请求正在共用。如果 block 的共享引用计数大于一就只能做“降级”而不能做“释放”否则会造成 use-after-free 类错误排查起来非常难受。4. 落地实操把 HiSparse 思路接到推理服务里4.1 给解码内核加一个“访问报告”要让缓存管理跟着注意力走第一步是让注意力内核告诉你“我这次到底碰了哪些 block”。这一步不需要把整个 cache 管理逻辑写进 kernel只需要在 kernel 输出里增加一个 block bitmap 或 block id list。伪代码大概是这样的// 假设 block_size 32 tokens // n_selected 稀疏注意力实际选中的 block 数 __global__ void sparse_attn_decode( const float* query, // [num_heads, head_dim] const void** kv_block_ptrs, const int* block_ids, int n_selected, float* output, int* accessed_block_ids, // 输出实际访问的 block id int* accessed_count)这里的accessed_block_ids是每个 step 结束后回传给 host 的“访存痕迹”。它的量很小比如一个请求选 32 个 block4 个 head 总共也就几百个 int几乎可以忽略。关键是 host 侧拿到这份信息后要更新 block 的热度统计并生成下一轮的预取列表和缺失列表。有了“访问报告”你还能做更细致的观察比如稀疏选择到底稳不稳定、命中率是高还是低、有没有明显的顺序访问模式。这些数据都是调缓存参数的直接依据。4.2 异步迁移队列怎么写才不会阻塞解码迁移 KV block 是纯数据拷贝完全可以用 CUDA stream 异步执行问题在于不能让迁移等待阻塞住 decode stream。我在项目里的做法是维护两条 stream一条是 decode 主计算流另一条是 cache 迁移流。迁移流只做 D2H、H2D 拷贝和 pinned memory 的 staging 操作。具体迁移流程1. 解码 kernel 完成 step N返回缺失 block 列表 2. host 调度器计算哪些缺失 block 在 L2/L3 可被拷贝 3. 分配 pinned staging buffer 4. 在迁移流上发起 cudaMemcpyAsync 5. 迁移流加一个 eventdecode 流在下一轮 attention 前 cudaStreamWaitEvent 这个 event 6. 只有真正被依赖的 block 拷贝完成后decode 才继续这里有一个大多数人容易忽略的点staging buffer 要提前分配好并固定住pinned memory不要每次迁移都现 malloc不然 DRAM 到 CPU 的拷贝带宽会被急出来的分配开销拖垮。4.3 预取要顺着解码的时序做预取策略的落地其实不复杂。在拿到当前步的 access list 之后对每个被访问的 block找出它在逻辑序列上的后续几个 block如果它们不在 GPU 上且热度预测会升高就加入预取队列。我常用一个“lookahead 窗口”参数默认取 4。也就是访问 block B 后预取 B1 到 B4。窗口太大容易浪费带宽因为后续 block 不一定真会被选中窗口太小又很可能出现下一个 step 又要等迁移。这个参数在我的测试里和 block_size 高度相关block_size32 时窗口取 4 比较稳block_size16 时窗口取 6 到 8 会更平滑。预取最好做成带优先级的队列缺失迁移 共享前缀预取 顺序邻居预取 热度上升预取。低优先级的预取如果在带宽高峰期可以直接放弃不要挤占正常计算资源。4.4 配套算子接口怎么设计如果你想把这套缓存管理接到自己推理框架里我建议不要用现成的 dense attention kernel 硬改而是让稀疏注意力算子和缓存管理器约定一个清晰的接口。接口需要支持两部分一是输入侧传入“当前可用的 KV block 指针表”二是输出侧返回“访问过的 block id 列表”。比较实用的接口形式// 每次 decode 前调用 struct SparseAttnParams { const void** resident_kv_blocks; // 只包含 GPU 上 resident 的 block const int* resident_block_ids; int num_resident; int* missing_block_ids; // 输出本轮需要但从 GPU 缺失的 block int* accessed_block_ids; // 输出本轮实际访问的 block };这里 resident 列表是 host 根据 block table 生成的每次生成都是一个新的前缀较短的数组。因为 attention 内核只用 resident 列表它不会去访问一个已经不存在的 block安全性也更好。有个细节需要提醒resident 列表的更新要小心并发。如果在迁移流还在拷贝某个 block 时decode 流已经把它的新指针传给了 kernel那 kernel 可能出现读到半块数据的情况。解决方法是给每个 block 加一个migration_generation字段指针表里带上 generationkernel 启动前校验 generation 一致才允许读。4.5 评估指标怎么设做这类优化指标不能用单点延迟一竿子打死。我会同时看四类数据Block 命中率当前解码 step 需要的 block 有多少已经在 GPU 上。这是最直观的缓存有效指标。TPOTtime per output token体现迁移是否是关键路径上的瓶颈。吞吐requests/s 或 tokens/s aggregate体现 batch 扩大后整体收益。长尾延迟 P99迁移抖动往往会体现在尾延迟上必须盯住。如果你发现命中率很高但 TPOT 没改善多半是索引查询或者迁移锁竞争成了瓶颈如果命中率低但吞吐还行那说明预取窗口给多了带宽浪费不少。用这几个指标交叉定位比单纯看显存占用有用得多。5. 常见问题与排查技巧实录5.1 命中率很高但 TPOT 反而更差了这是我第一次实现时踩的坑。当时 block 命中率已经到 95% 以上但 TPOT 比不用分层缓存还慢。查了半天发现问题出在“每次 decode step 都要重新生成 resident 指针表”而生成这个表的过程涉及 GPU 到 CPU 的同步一旦每次 step 都有一次同步流水线就被打断了。解决办法是把指针表生成放到单独的 CPU 线程里异步做并且只在映射变化时更新映射完全没变时直接复用上一轮的表。实际操作下来把同步次数从每个 step 一次降到每几十个 step 一次TPOT 立刻恢复到正常水平。5.2 稀疏策略“切换太吵”导致热度统计失效有些稀疏注意力策略会根据当前 token 的分布动态更换选择方式比如 window 模式和 top-k 模式混着切。这时 block 的热度信号会非常不稳定同一个 block 前几步每次都被访问后几步又完全不碰。我一度以为是缓存迁移太慢后来把访问日志拉出来看才发现是“假热点”热度计数的衰减周期太短了。我把热度窗口从几十个 step 拉长到几百个 step并且改用指数衰减而不是滑动窗口统计问题就缓解了很多。核心经验是热度的更新要匹配稀疏策略的时间尺度而不是匹配解码 step 的时间尺度。5.3 显存碎片和层级索引膨胀分层缓存跑久了之后可能会出现一种诡异现象显存总量看起来还有剩但分配新 block 老是失败。原因是每个层级内部如果都用简单 freelist长时间高频迁移后会产生大量碎片。我后来在每个物理层级里改成了按 block_size 对齐的 slab 分配器并加了一个“block 迁移计数”的指标。当某个层级的 block 迁移频率特别高时可能是这个层级的容量和访问模式不匹配这时候最简单的方法是调大/调小该层容量而不是继续优化分配器。5.4 Host 带宽顶不住当 batch 大到一定程度瓶颈会从 GPU 显存变成 CPU 内存到 GPU 之间的传输带宽。这时候你会发现 block 命中率一直上不去因为迁移数据量超过了互连带宽。我自己用过两个有效的缓解手段一是对 KV block 做低损压缩FP16 转 INT8 需要处理精度问题但很多场景下影响不大二是做“批量迁移”把多个 block 的拷贝合并成一次较大的 memcpy而不是每个 block 单独发一次这样 PCIe/NVLink 的有效带宽能拉高不少。5.5 长尾和抢占分层缓存天然会引入一次性的迁移延迟如果某个请求刚切到一段很久没访问的历史上下文它得把一整套冷 block 从 L3 拉回来这次 TPOT 会突然飙高。为了不让这一发长尾打爆 P99我给这类“冷区启动”事件做了异步限速允许它慢但不允许它阻塞其他请求。具体做法是给迁移队列按请求分优先级冷区预热请求的迁移任务挂在低优先级队列同时通过预测“复用概率”提前几轮开始预热。实验下来P99 最长尾可以压到普通 TPOT 的 3 倍以内。6. 性能边界与调参经验6.1 什么场景收益最大从原理就能看出来HiSparse 这套分层缓存思路最适合“长上下文 低显存 高并发”的组合。比如 128K 上下文的外部文档问答、长时间对话、代码仓库级自动补全这些场景下每请求的 KV cache 可能到几十 GB任何单卡的显存都装不下全部请求分层带来的价值非常大。另外如果一个系统里存在大量请求共享公共前缀比如系统提示词很长、知识库前缀一致收益也会被放大。因为公共前缀的 KV block 会被不同请求反复命中热度极高几乎总是稳定驻留在 GPU 上迁移开销很小。6.2 什么场景别硬上短上下文场景例如 1K 以内我不建议上这套方案。因为 block 数量少分层索引和迁移调度的开销可能比省下的显存带宽还多。还有一种是稀疏注意力本身命中极不稳定的场景比如选块策略完全随机那热度地图基本是噪声分层缓存失去了预测价值。另外如果单请求超长上下文而且 batch 只有个位数迁移延迟会直接暴露在关键路径上收益有限。这时候不如老老实实把 KV cache 做压缩或者用更激进的稀疏定义。6.3 一套保守的初始参数如果你要在自己的服务里试这套思路可以先用下面这组参数起步参数推荐值说明block_size32 tokens平衡粒度和迁移效率L1 容量可容纳 batch 内 85% 活跃 block过高浪费显存过低命中率差L2 容量尽量大用 CPU 内存温层是带宽性价比最高的地方预取窗口4 个相邻 block小于等于 8 都合理热度衰减半衰期256 个 step匹配长上下文的一般访问节奏迁移批大小32 个 block 起步优先合并成一次大拷贝这组参数不是最优解但它能让你快速验证分层缓存的方向是否有效。之后再用命中率、TPOT、P99 三个指标去反推调整。我在实际项目中体会最深的一点是稀疏注意力本身并不能带来多少解码加速真正让它跑起来的是外界配合它的访存策略。HiSparse 给我的启发是不要把 KV cache 当成一个被动的存储区而要把缓存迁移、预取、淘汰都变成解码过程的一部分。最后再分享一个小技巧就算你暂时不想引入完整的分层缓存也可以先在现有 paged KV cache 上加上“block 访问报告”这个观测层光是把每个 step 在访问哪些 block 的数据拉出来看往往就能找到一批此前完全没意识到的显存浪费点。
返回列表