HGCA: Hybrid GPU-CPU Attention for Long Context LLM Inference
作者: Weishu Deng, Yujie Yang, Peiran Du, Lingfeng Xiang, Zhen Lin, Chen Zhong, Song Jiang, Hui Lu, Jia Rao
分类: cs.LG
发布日期: 2025-07-03
💡 一句话要点
HGCA:用于长上下文LLM推理的混合GPU-CPU注意力机制
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 长上下文LLM 混合CPU-GPU注意力 稀疏注意力 KV缓存 LLM推理
📋 核心要点
- 长上下文LLM推理受限于GPU内存,现有KV缓存卸载或稀疏注意力方法未能充分利用CPU资源且精度受损。
- HGCA采用混合CPU-GPU注意力,GPU处理近期KV缓存的密集注意力,CPU并行处理显著KV缓存的稀疏注意力。
- 实验表明,HGCA在商用GPU上实现了更优的可扩展性、更长的序列支持和更高的精度,优于现有稀疏注意力方法。
📝 摘要(中文)
大规模语言模型(LLM)的推理扩展日益受到有限GPU内存的制约,特别是由于长上下文生成所需的键值(KV)缓存不断增长。现有方法通常将KV缓存卸载到CPU内存或应用稀疏注意力以减少GPU负载,但这些方法往往未能充分利用CPU计算资源并降低了准确性。我们提出了HGCA,一种混合CPU-GPU注意力机制,它能够以接近完整注意力的质量实现可扩展、高吞吐量的LLM推理。HGCA对保留在GPU内存中的最近生成的KV条目执行密集注意力,并对CPU内存中选定的、显著的KV条目执行并行稀疏注意力。注意力输出使用log-sum-exp融合有效地合并,从而最大限度地减少PCIe传输开销。HGCA还引入了一种针对CPU执行优化的细粒度、逐头稀疏化策略,在减少计算量的同时保留了上下文相关性。我们的实现无缝集成到现有的LLM框架中,无需模型重新训练。在各种模型和工作负载上的实验表明,HGCA实现了卓越的可扩展性,支持更长的序列和更大的批处理大小,并且在性能和准确性方面均优于现有的稀疏注意力基线——所有这些都在商用GPU硬件上实现。
🔬 方法详解
问题定义:现有长上下文LLM推理方法面临GPU内存瓶颈,特别是KV缓存的增长。将KV缓存卸载到CPU或使用稀疏注意力虽然能缓解GPU压力,但前者引入了PCIe传输瓶颈,后者则可能损失精度,并且CPU计算资源利用率不高。
核心思路:HGCA的核心在于利用GPU和CPU各自的优势,将注意力计算任务分配到最合适的硬件上。GPU擅长处理小规模的密集计算,因此用于处理最近生成的、上下文关联性强的KV缓存;CPU则用于处理大规模的、但经过稀疏化处理的KV缓存。通过并行计算和高效的融合机制,实现整体性能的提升。
技术框架:HGCA包含以下主要模块:1) GPU密集注意力模块:处理GPU内存中的近期KV缓存,执行标准注意力计算。2) CPU稀疏注意力模块:处理CPU内存中的稀疏KV缓存,执行稀疏注意力计算。3) KV缓存管理模块:负责KV缓存的卸载、存储和检索。4) 注意力融合模块:使用log-sum-exp融合GPU和CPU的注意力输出,减少PCIe传输开销。
关键创新:HGCA的关键创新在于混合CPU-GPU注意力机制和细粒度的逐头稀疏化策略。混合注意力机制充分利用了GPU和CPU的计算资源,实现了更高的吞吐量和可扩展性。逐头稀疏化策略则针对CPU执行进行了优化,在减少计算量的同时保留了上下文相关性。
关键设计:HGCA的关键设计包括:1) 细粒度的逐头稀疏化策略,允许不同的注意力头具有不同的稀疏度,以适应不同的上下文信息。2) 使用log-sum-exp融合GPU和CPU的注意力输出,避免直接传输注意力权重,从而减少PCIe传输开销。3) 无需模型重新训练,可以无缝集成到现有的LLM框架中。
🖼️ 关键图片
📊 实验亮点
实验结果表明,HGCA在各种模型和工作负载上均优于现有的稀疏注意力基线。HGCA实现了卓越的可扩展性,支持更长的序列和更大的批处理大小,并且在性能和准确性方面均有提升。具体数据需要在论文中查找。
🎯 应用场景
HGCA适用于需要处理长上下文的大规模语言模型推理,例如文档总结、机器翻译、对话生成等。它能够提升LLM在资源受限环境下的推理性能,降低部署成本,并支持更大规模的模型和更长的上下文长度,从而拓展LLM的应用范围。
📄 摘要(原文)
Scaling inference for large language models (LLMs) is increasingly constrained by limited GPU memory, especially due to growing key-value (KV) caches required for long-context generation. While existing approaches offload KV caches to CPU memory or apply sparse attention to reduce GPU load, they often underutilize CPU compute resources and compromise accuracy. We present HGCA, a hybrid CPU-GPU attention mechanism that enables scalable, high-throughput LLM inference with near-full attention quality. HGCA performs dense attention on recently generated KV entries retained in GPU memory and parallel sparse attention on selected, salient KV entries in CPU memory. The attention outputs are efficiently merged using log-sum-exp fusion, minimizing PCIe transfer overhead. HGCA also introduces a finegrained, per-head sparsification strategy optimized for CPU execution, preserving contextual relevance while reducing computation. Our implementation seamlessly integrates into existing LLM frameworks without requiring model retraining. Experiments across diverse models and workloads show that HGCA achieves superior scalability, supports longer sequences and larger batch sizes, and outperforms existing sparse attention baselines in both performance and accuracy -- all on commodity GPU hardware.