通过统一稀疏注意力与分层内存,显著提升长文本大模型推理速度。
Unifying Sparse Attention with Hierarchical Memory for Scalable Long-Context LLM Serving
- 设计统一页面抽象,兼容多种稀疏粒度的注意力机制。
- 动态分配显存预算,减少跨GPU-CPU传输次数,提速7-9倍。
- 适合需要高效处理超长上下文的大模型部署场景。
长上下文大模型服务受限于不断增长的键值(KV)缓存访问开销。动态稀疏注意力通过每步仅访问少量、与查询相关的KV状态并扩展存储至CPU内存来缓解问题。然而,实际中算法节省难以转化为系统级性能提升,因不同稀疏方法粒度不一,需各自定制实现。同时,分层KV存储引入新瓶颈:在GPU-CPU边界频繁获取细粒度、不规则的KV子集会抵消稀疏优势。本文提出SPIN,一种面向稀疏注意力的推理框架,通过三项技术协同优化执行流程与分层存储:(1) 统一分区抽象,将不同稀疏粒度映射到共享的基于页的KV基础结构;(2) 局部性感知的缓存管理器,动态分配每请求显存预算,并采用GPU友好的分桶LRU策略减少PCIe往返;(3) 两级元数据布局,按活跃工作集大小而非最坏地址空间配置。基于vLLM集成三种代表性稀疏注意力算法,SPIN实现端到端吞吐率提升1.66-5.66倍,首次响应时间(TTFT)降低7-9倍,端到端延迟(TPOT)相比原始稀疏实现最高降低58%。
原文摘要 · Abstract (English)
Long-context LLM serving is bottlenecked by the cost of attending over ever-growing KV caches. Dynamic sparse attention promises relief by accessing only a small, query-dependent subset of the KV state per decoding step and extending the KV storage to CPU memory. In practice, however, these algorithmic savings rarely translate into end-to-end system-level gains because sparse methods typically operate at different granularities and thus rely on ad hoc, per-algorithm implementations. At the same time, hierarchical KV storage introduces a new systems bottleneck: retrieving fine-grained, irregular KV subsets across the GPU-CPU boundary can easily erase the benefits of sparsity. We present SPIN, a sparse-attention-aware inference framework that co-designs the execution pipeline with hierarchical KV storage through three techniques: (1) a unified partition abstraction that maps different sparsity granularities onto a shared page-based KV substrate; (2) a locality-aware KV cache manager that dynamically sizes per-request HBM budgets and uses a GPU-friendly bucketed LRU policy to cut PCIe round-trips; and (3) a two-level hierarchical metadata layout sized to the active working set rather than the worst-case address space. Built on vLLM with three representative sparse attention algorithms, SPIN delivers 1.66-5.66x higher end-to-end throughput and 7-9x lower TTFT than vLLM, and reduces TPOT by up to 58% over the original sparse-attention implementations.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。