Papers

TriAttention:稀疏混合精度注意力加速长上下文推理

TriAttention: Unified Sparse Mixed-Precision Attention for Long-Context LLM Inference

Why we picked it

一篇很实在的推理优化论文:把注意力分块按精度分档,长文本生成最高快 7.88 倍,代码已开源,做长上下文部署的可以试试。

TriAttention 将 query-key 分块分配到跳过、4-bit 或 8-bit 三种计算路径,用两条对角距离边界编码每个 head 的策略,存储开销与上下文长度无关。离线校准在共享计算预算下为整个模型选边界,最小化归一化输出误差。预填充阶段通过融合核把两种精度合并成一次归一化注意力输出,量化后的 KV cache 在解码时复用。实验显示端到端生成相比 Torch attention 提速最高 7.88x,相比 BF16 基线提速 2.65x,任务质量与原模型相当。

Original · arXiv cs.AI

TriAttention: Unified Sparse Mixed-Precision Attention for Long-Context LLM Inference

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 https://github.com/HanzhiZhang-Ulrica/TriAttention.