论文
用于长上下文推理的具有 CPU-GPU 并行性的高效混合稀疏注意力
An Efficient Hybrid Sparse Attention with CPU-GPU Parallelism for Long-Context Inference
摘要
长上下文推理越来越多地在 CPU 驻留的 KV 缓存上运行,这要么是因为解码时 KV 状态超出了 GPU 内存容量,要么是因为分解的预填充解码系统将 KV 数据放置在主机内存中。尽管块稀疏注意力在这种情况下降低了注意力成本,但仅稀疏性不足以提高端到端效率。仅 GPU 设计仍然受到 PCIe 带宽和元数据内存开销的限制,而 CPU-GPU 混合设计仍然面临大量 GPU 空闲时间以及 CPU 端 top-k 选择和稀疏注意力计算方面的瓶颈。 Fluxion 基于三个关键见解:输出感知的 KV 预算、头部特定和粒度感知的稀疏配置,以及跨设备协调执行,以对 CPU 驻留的 KV 缓存进行稀疏关注。在这些见解的指导下,Fluxion 结合了轻量级头部属性预测器、粒度预算选择器和基于优先级的调度器,共同优化预算分配、稀疏配置和 CPU-GPU 执行重叠。这种协同设计使混合稀疏注意力能够在长上下文推理中实现准确性和系统效率。在 2 个模型、3 个基准测试和 40 个任务中,Fluxion 很好地保持了质量 - 相对于 FULL,最差的平均降级仅为 -0.26,同时比最强的固定稀疏混合基线(其 KV 预算仅为 0.05)提供 1.5$\times$-3.7$\times$ 加速。