论文略读：OmniKV: Dynamic Context Selection for Efficient Long-Context LLMs

2025 ICLR

• 核心思想：在单个生成步骤内，被模型高度关注的（注意力得分高的）Token 集合，在不同的 Transformer 层之间表现出高度的相似性
  • 
• 论文提出OmniKV
  • 仅选择少数几个层（称为 “Filter 层”）来计算完整的注意力分数并识别重要的 Token 子集
  • 其他大多数层则直接复用（共享）来自最近 Filter 层识别出的 Token 索引。
  • 在计算注意力时，仅加载并计算这个稀疏子集 KV Cache，从而大幅减少计算量和数据传输量
  • 
    • Prefill 阶段
      • 对输入的 Prompt 进行编码，生成完整的 KV Cache
      • 将大部分非 Filter 层的 KV Cache 卸载（offload）到 CPU 内存
      • 仅保留少量 Filter 层的 KV Cache 在 GPU 上
    • Decode 阶段（生成阶段）
      • Filter 层：计算完整注意力，并使用 Context Selector 动态选择当前步骤最重要的 Top-K 个 Token 索引
      • 非 Filter 层：直接从 CPU 加载（load）由前一个 Filter 层选择出的 Token 索引对应的 KV Cache 子集，并在该子集上执行稀疏注意力计算