跳到主要内容
← 返回研究笔记
论文笔记 USENIX ATC 2024 2024

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

USENIX ATC 2024 — KAUST + 阿里。把 KV cache 做成 HBM/DRAM/SSD 三级层级 + 跨多轮对话的 prefix 复用

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 管理”的最强基线之一,论文必须做精确对比。

我们的差异化空间:

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

横向对比

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

待精读 / 待补问题

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