arXiv:2605.18856cs.LGcs.CL2026-05被引 1

通过角度域注意力与率失真保留,显著降低长文本推理的内存占用和显存压力。

SPHERICAL KV: Angle-Domain Attention and Rate-Distortion Retention for Efficient Long-Context Inference

论文配图:SPHERICAL KV: Angle-Domain Attention and Rate-Distortion Retention for Efficient Long-Context Inference
图 1 · 摘自论文原文
  • 用球面参数化存储键值,直接在解码中计算注意力,避免重建密集向量。
  • 在固定预算下动态分配保留策略与精度层级,减少显存读写开销。
  • 适合高并发长序列生成场景,尤其优化了真实服务中的显存瓶颈。

长上下文推理受限于键值缓存(KV cache):缓存内存随上下文长度增长,解码受制于重复的高带宽内存(HBM)数据流而非算术运算。现有方法如淘汰、滑动窗口、量化和卸载虽能减少内存占用,但未彻底解决关键路径瓶颈,尤其当压缩状态仍需在解码时重构为密集向量。本文提出 Spherical KV,将 KV 分配建模为基于注意力几何的率失真问题,实现高效解码。其核心包含两点:(i) 在解码热循环中廉价表示方向信息;(ii) 根据未来效用估计分配保留与精度。第一部分,角度域注意力(ADA)采用球面参数化(标量半径 + 紧凑角度编码),直接从编码计算注意力得分,无需重构密集键。第二部分,率失真保留(RDR)在固定预算下联合决定每令牌/头的保留/丢弃与精度层级,生成同质页、轻量元数据与聚合读取。二者协同提供面向部署的机制,在降低 KV 居住的同时保持解码效率。

原文摘要 · Abstract (English)

Long-context inference is increasingly constrained by the KV cache: resident memory grows with context length, and decoding becomes limited by repeated High Bandwidth Memory (HBM) streaming rather than arithmetic. Existing methods such as eviction, windowing, quantization, and offloading reduce footprint, but often leave the critical-path bottleneck only partially addressed, especially when compressed states must still be reconstructed into dense vectors during decoding. We present Spherical KV, a long-context inference method that treats KV allocation as a rate-distortion problem grounded in attention geometry for efficient decoding. The method is built on two ideas: (i) represent directional information cheaply in the decode hot loop, and (ii) allocate retention and precision according to estimated future utility. Its first component, Angle-Domain Attention (ADA), stores keys in a spherical parameterization consisting of a scalar radius and compact angle codes, and computes attention logits directly from these codes without reconstructing dense keys. This preserves a paged, block-local, fusion-friendly decode path and directly targets HBM traffic in realistic serving settings. Its second component, Rate-Distortion Retention (RDR), jointly chooses keep/drop decisions and precision tiers per token and head under a fixed budget, producing tier-homogeneous pages with lightweight metadata and coalesced reads. Together, ADA and RDR provide a deployment-oriented mechanism for reducing KV residency while preserving decode efficiency.

长上下文注意力机制内存优化KV缓存

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。