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 的计算区域，使策略存储与上下文长度无关。
正文 · 原文
Abstract:Sparsity and quantization reduce the cost of attention by skipping interactions and lowering the precision of those retained. However, useful interactions differ in their precision requirements: some tolerate low-precision scores while remaining sensitive to removal. We therefore introduce TriAttention, which assigns query-key (QK) tiles to skipping, 4-bit, or 8-bit computation, using 4-bit to retain quantization-tolerant interactions at lower cost. To represent compactly, we encode each head's computation regions with two diagonal-distance boundaries, keeping policy storage independent of context length. Offline calibration then evaluates candidate boundary pairs using normalized projected-output errors and complete-tile execution costs. Using these measurements, a model-wide allocator selects each head's boundaries to minimize calibration error under a shared computation budget. During prefill, the selected boundaries generate precision-grouped traversals for a fused kernel that skips omitted tiles and combines contributions from both QK precisions into one normalized attention output. The quantized keys and values prepared during prefill are then reused in a compressed key-value (KV) cache for decoding, reducing repeated reads of prompt data. Experiments show that TriAttention achieves end-to-end generation speedups of up to 7.88x over the Torch attention and 2.65x over the BF16 baseline. Meanwhile, TriAttention maintains quality comparable to the original models on general-purpose and long-context tasks. These results show how a compact retention-and-precision policy translates local sensitivities into attention computation and repeated KV-cache access. The implementation is this https URL.
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2611.07256 [cs.LG] |
| (or arXiv:2611.07256v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2611.07256 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Hanzhi Zhang [view email]
[v1]
Thu, 5 Nov 2026 00:14:13 UTC (1,503 KB)
来源：arXiv cs.AI · arxiv.org