01. AttentionStore — KV 三级缓存 + 跨会话复用

元数据

论文标题:AttentionStore: Cost-effective Attention Reuse across Multi-turn Conversations in Serving Large Language Models
出处:USENIX ATC 2024
作者方向:KAUST + 阿里(早期版本叫 CachedAttention)
关键词:KV cache、multi-turn dialogue、hierarchical storage、prefix reuse

一句话精髓

把 KV cache 当成”会话级冷热数据”做三级存储(HBM-DRAM-SSD),让多轮对话不再每次重算 prefill。

解决的问题

聊天场景下,同一个用户的多轮对话前缀越来越长,每一轮新请求都要把整个历史 prefill 一遍——计算量随轮数 O(N²) 增长。已有 prefix cache 只能驻留 HBM,容量极小,踢出概率高——相当于做了等于没做。

关键 idea

维度	设计
存储层次	HBM(active session) → DRAM(near-active) → SSD(cold) 三级
入卡时机	用户发起新轮次时,根据 session ID 把对应 KV 块预取到 HBM
粒度	按 layer + token block 拆分,layer-wise 流水预取
替换策略	跨会话 LRU + 用户最近活跃度感知
对外语义	框架透明——上层 vLLM/TGI 看到的还是”prefix cache 命中”

关键数据(公开报告)

多轮对话场景下,首 token 延迟降低数倍(对比每轮重算 prefill)
KV 数据 从 HBM 溢出到 DRAM/SSD 后,命中率提升 10× 以上
三级架构在长会话下的 token 单位成本明显下降(具体倍数依工作负载)

关键架构图(我重画)

   用户请求(session_id, new_turn)
            │
            ▼
   ┌─────────────────────┐
   │ Session Manager     │ ── 查 session 在哪一级
   └──────────┬──────────┘
              │
   ┌──────────┼──────────┐
   ▼          ▼          ▼
  HBM       DRAM        SSD
 (1-2GB)   (100GB)    (TB级)
   │          │          │
   └──── 预取 / 流水 ─────┘
              │
              ▼
       Attention 计算(只 decode 新 token)

局限

只针对对话场景——长记忆 Agent / RAG / 多模态没覆盖
假设 session 边界清晰,对 cross-user prefix 共享不友好
SSD 用起来还是依赖 OS page cache + buffered I/O,没用 GPUDirect Storage 直 DMA
KV 在三级间搬运是整块,没做语义压缩(对比 CacheGen)

对本项目的启示

🌟 直接竞品——第一模块”跨层级 KV 管理”的最强基线之一,论文必须做精确对比。

我们的差异化空间:

多类型数据:AttentionStore 只管 KV,我们要把向量索引、多模态 blob 一起纳入同一套 hierarchy
跨用户 prefix 共享:我们要对 RAG 文档前缀、系统 prompt 这类”准全局”做层级化共享(参考 LMCache)
GPUDirect Storage 直通:把 SSD 升级成”直接 DMA 进 HBM”,绕过 CPU bounce buffer
token 单位成本建模:AttentionStore 用启发式 LRU,我们要做模型驱动的最优放置

横向对比

系统	KV 多级	跨实例共享	多类型数据	主要场景
vLLM PagedAttention	单卡 HBM 内	❌	❌	单卡内分页
AttentionStore	HBM-DRAM-SSD	弱	❌	多轮会话
LMCache	HBM-DRAM	✅ 跨实例	❌	跨实例复用
Mooncake	HBM-DRAM-SSD	✅ 池化	❌	大规模 PD 解耦
本项目目标	三级 + 多类型 + 跨实例 + 模型驱动放置	✅	✅	长记忆系统

待精读 / 待补问题

论文里给出的 token cost 计算模型是什么形式?是否可作为我们建模的起点?
它的预取策略对 SSD 的随机 IO 放大了多少?
实测 SSD 用什么 IO 模式(O_DIRECT? io_uring? GDS?)

01. AttentionStore — KV 三级缓存 + 跨会话复用

搜索