通过哈希编码实现语义稀疏,显著加速注意力计算。
HashAttention: Semantic Sparsity for Faster Inference
- 将关键标记识别转化为推荐问题,用哈希空间捕获语义相似性。
- 在GPT-FAST上实现32倍稀疏,延迟降低4.3倍,吞吐提升3.12倍。
- 仅需每标记32比特辅助内存,适合部署在资源受限场景。
长上下文对先进AI系统至关重要,但注意力计算存在可扩展性挑战。尽管缩放点积注意力(SDPA)具有标记稀疏性,即仅有少数关键标记显著影响输出,但利用这种稀疏性仍具挑战。现有方法或导致质量下降,或需大量额外资源。本文表明,关键标记识别是最大内积搜索(MIPS)问题。然而,现有MIPS方案不适用于SDPA,因缺乏GPU友好性且因查询与键分布分离而表现不佳。本文提出HashAttention,将关键标记识别建模为推荐问题。给定查询,HashAttention使用学习映射函数将键和查询编码至汉明空间,捕捉所需语义相似性。通过位运算高效识别关键标记,并仅用这些标记计算注意力,提升整体效率。在通用数据上训练,HashAttention最多减少16倍计算标记数,质量损失极小,每标记仅需32比特辅助内存。通过任务特定微调,稀疏度可进一步提升至32倍。在A100 GPU上,32倍稀疏下,HashAttention使GPT-FAST的注意力延迟降低4.3倍,FlashDecode降低2.54倍,GPT-FAST吞吐量最高提升3.12倍。
原文摘要 · Abstract (English)
Leveraging long contexts is crucial for advanced AI systems, but attention computation poses a scalability challenge. While scaled dot-product attention (SDPA) exhibits token sparsity, i.e. only a few pivotal tokens significantly contribute to output, exploiting this sparsity remains challenging. Existing methods either suffer from quality degradation or require substantial additional resources. We show that identifying pivotal tokens is a Maximum Inner Product Search (MIPS) problem. However, existing MIPS solutions are not well-suited for SDPA, as they are not GPU-friendly and often underperform due to the separated query and key distributions. This paper introduces HashAttention, framing pivotal token identification as a recommendation problem. Given a query, HashAttention encodes keys and queries in Hamming space, capturing the required semantic similarity, using learned mapping functions. HashAttention efficiently identifies pivotal tokens for a given query using bitwise operations and computes attention using only these tokens, improving the overall attention efficiency. Trained on generic data, HashAttention reduces tokens used by up to $16\times$ with minimal quality loss, requiring only 32 bits of auxiliary memory per token. Sparsity can be further improved to $32\times$ through task-specific fine-tuning. On A100 GPU, at $32\times$ sparsity, incorporating HashAttention reduces attention latency by up to $4.3\times$ in GPT-FAST and $2.54\times$ in FlashDecode, and achieves up to $3.12\times$ higher throughput for GPT-FAST.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。