TriAttention：面向长上下文 LLM 推理的统一稀疏混合精度注意力

稀疏化和量化通过跳过交互并降低所保留交互的计算精度来减少注意力计算的成本。然而，有效交互对精度的要求各不相同：一些交互能够容忍低精度分数，却仍对移除操作敏感。因此，我们提出 TriAttention，将查询—键（QK）分块分配为跳过、4 位计算或 8 位计算，利用 4 位计算以更低成本保留可容忍量化的交互。为实现紧凑表示，我们使用两个对角线距离边界对每个注意力头的计算区域进行编码，使策略存储量与上下文长度无关。随后，离线校准利用归一化投影输出误差和完整分块执行成本来评估候选边界对。基于这些测量结果，一个模型级分配器为每个注意力头选择边界，在共享计算预算约束下最小化校准误差。在预填充阶段，所选边界为融合内核生成按精度分组的遍历方案，该内核跳过被省略的分块，并将两种 QK 精度的计算贡献合并为一个归一化注意力输出。预填充阶段准备好的量化键和值随后存入压缩的键值（KV）缓存，在解码时复用，从而减少对提示数据的重复读取。实验表明，相较于 Torch 注意力和 BF16 基线，TriAttention 的端到端生成速度最高分别提升至 7.88 倍和 2.65 倍。同时，TriAttention 在通用任务和长上下文任务上保持了与原始模型相当的质量。这些结果表明，紧凑的保留与精度策略能够将局部敏感性转化为对注意力计算和重复 KV 缓存访问的安排。实现代码见 https://github.com/HanzhiZhang-Ulrica/TriAttention。