HGCA: Hybrid GPU-CPU Attention for Long Context LLM Inference

文章主要内容总结

本文针对大型语言模型（LLMs）长上下文推理中GPU内存受限（尤其是KV缓存随上下文增长而膨胀）的问题，提出了一种混合CPU-GPU注意力机制HGCA。该机制通过以下方式实现高效推理：

1. 计算分工：在GPU内存中保留近期生成的KV条目，对其执行密集注意力计算；在CPU内存中存储选定的重要KV条目，对其执行并行稀疏注意力计算。
2. 结果融合：通过log-sum-exp融合策略高效合并GPU和CPU的注意力输出，最大限度减少PCIe传输开销。
3. 稀疏化策略：引入适合CPU执行的细粒度、按注意力头（per-head）的稀疏化策略，在减少计算量的同时保留上下文相关性。
4. 兼容性：无需重新训练模型，可无缝集成到现有LLM框架中。

实验表明，HGCA在商品GPU硬件上实现了更优的可扩展性（支持更长序列、更大批次），在性能和准确性上均优于现有稀疏注意力基线方法。

文章创新点

1. 混合注意力架构：首次将GPU密集注意力与CPU稀疏注意力结合，充分利用GPU的计算效率和CPU的内存容量及控制逻辑优势，解决了单一设备推理的内存或计算瓶颈。
2. log-sum-exp融合机制：通过融合GPU和CPU的部分注意力输出（而非传输原始KV缓存），显著降低PCIe传输开销，实现高效跨设备协作。