TileMix：让长上下文注意力按计算瓦片动态混合精度

导语

长上下文推理的预填充阶段，需要一次性处理大量输入 token。标准密集自注意力会计算查询与键之间的成对分数，序列变长后，计算量和内存访问压力都会快速上升。降低数值精度是常见的优化方向，但如果整条注意力路径统一使用 INT8，虽然有机会减少计算和数据搬运成本，也可能损害长上下文任务的表现。

TileMix 的切入点不是删掉部分 token 交互，而是决定“哪些注意力分数应该用更高精度计算”。

核心方法：以瓦片为单位路由精度

TileMix 将注意力矩阵切分成与硬件执行方式匹配的 score tiles，再把相邻的键瓦片组织为瓦片组。每个瓦片组通过紧凑的 bitmask 标记，随后在融合的、类似 FlashAttention 的内核中，被分派到 FP16 或 INT8 分数计算路径。

这一设计解决了一个实际的工程矛盾：如果精度选择过于细碎，路由元数据和控制分支会抵消低精度计算带来的收益；如果选择过于粗略，又难以保留关键交互的数值精度。TileMix 允许一个路由位控制多个相邻键瓦片，在保持硬件对齐的同时压缩元数据，尤其适合较长序列。

两条精度路径并不是各自完成一套注意力计算。无论瓦片采用 FP16 还是 INT8，都会更新同一个 online softmax 状态。这使得不同精度的局部计算可以继续参与同一份密集注意力归一化结果，避免为了混合精度而拆分成多个独立算子。

与稀疏注意力的区别

TileMix 路由的是数值精度，而不是 token 是否参与计算。只要属于合法的瓦片组，交互仍会被保留，因此它维持了密集注意力的 token 连通性，也不需要训练路由器或重新训练模型。该实现还覆盖 grouped-query attention、变长 batch，以及使用 INT8 key/value cache 的场景。

实验信号与意义

论文在 LongEval、LV-Eval 和 A100 预填充测试上，使用 LLaMA、Qwen、Vicuna 等模型进行评估。摘要给出的结论是：相较统一 INT8，TileMix 能恢复部分长上下文质量；相较 FP16，则改善预填充吞吐。不同模型和路由设置因此可以形成可调节的精度—效率前沿，而不是只能在“完整 FP16”和“统一低精度”之间二选一。

这项工作的价值主要在于把混合精度从离线量化配置变成了融合注意力内核中的空间执行决策。它并未消除长上下文注意力的二次复杂度，也没有声称通过稀疏化减少全部交互；它提供的是一种更贴近 GPU 瓦片执行方式的折中方案。实际收益仍取决于路由策略、硬件、序列长度和模型家族，因而开源实现以及更完整的性能数据，对判断其部署价值十分重要。

来源：Hugging Face Daily Papers