arXiv cs.AI· Hanzhi Zhang, Qiao Zhang, Qinglei Cao, Heng Fan, Yan Huang, Kewei Sha, Yunhe Feng·· 11 小时前AI 评分38
TriAttention：面向长上下文 LLM 推理的统一稀疏混合精度注意力机制
TriAttention: Unified Sparse Mixed-Precision Attention for Long-Context LLM Inference
AI 导读
TriAttention 将 query-key tile 分配至跳过、4-bit 或 8-bit 三档计算，用两条对角距离边界编码每个 head 的计算区域，使策略存储与上下文长度无关。
正文 · AI 翻译
摘要：稀疏性和量化通过跳过交互以及降低所保留交互的精度来降低注意力计算成本。然而，有用的交互在精度需求上各不相同：有些交互可以容忍低精度分数，但对移除仍然敏感。因此，我们提出了 TriAttention，它将查询-键（QK）瓦片分配为跳过、4 位或 8 位计算，使用 4 位以更低的成本保留对量化容忍的交互。为了紧凑表示，我们用两个对角距离边界来编码每个头的计算区域，使策略存储与上下文长度无关。然后，离线校准使用归一化投影输出误差和完整瓦片执行成本来评估候选边界对。利用这些测量结果，模型级分配器在共享计算预算下选择每个头的边界，以最小化校准误差。在预填充期间，选定的边界为融合内核生成按精度分组的遍历，该内核跳过省略的瓦片，并将两种 QK 精度的贡献合并为一个归一化注意力输出。预填充期间准备的量化键和值随后在解码时被复用于压缩键值（KV）缓存中，减少对提示数据的重复读取。实验表明，TriAttention 相比 Torch 注意力实现了高达 7.88 倍的端到端生成加速，相比 BF16 基线实现了 2.65 倍的加速。同时，TriAttention 在通用和长上下文任务上保持了与原始模型相当的质量。这些结果展示了紧凑的保留与精度策略如何将局部敏感性转化为注意力计算和重复的 KV 缓存访问。实现见 此 https URL。
| 主题： | 机器学习（cs.LG）；人工智能（cs.AI） |
| 引用为： | arXiv:2611.07256 [cs.LG] |
| （或 arXiv:2611.07256v1 [cs.LG] 为此版本） | |
| https://doi.org/10.48550/arXiv.2611.07256 由 DataCite 颁发的 arXiv DOI（待注册） |
提交历史
来自：Hanzhi Zhang [查看电子邮件]
[v1]
2026 年 11 月 5 日星期四 00:14:13 UTC（1,503 KB）
来源：arXiv cs.AI · arxiv.org