🤖 AI

TileMix: Tile-Centric Mixed-Precision Attention for LLM Inference Acceleration

TileMix 是一种无需训练、以分块为中心的混合精度算子，它通过在融合的稠密注意力操作中，将硬件对齐的注意力评分分块动态路由至 FP16 和 INT8 路径，从而在加速长上下文大语言模型推理的同时，恢复了统一低精度方法所损失的精度并提升了吞吐量。

原作者： Hanzhi Zhang, Qiao Zhang, Qinglei Cao, Heng Fan, Yan Huang, Kewei Sha, Yunhe Feng

发布于 2026-08-19

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Hanzhi Zhang, Qiao Zhang, Qinglei Cao, Heng Fan, Yan Huang, Kewei Sha, Yunhe Feng

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

大语言模型已成为新一代人工智能背后的引擎，它们能够总结整本书籍、回答来自长文档的复杂问题，并进行跨越数千词的对话。为了实现这一点，这些模型依赖于一种被称为“注意力”（attention）的机制，该机制允许它们权衡句子中每个词相对于其他所有词的重要性。当文本较短时，这个过程非常迅速。但随着上下文增长到包含整个章节或法律合同，计算成本会呈爆炸式增长。模型必须为每一对可能的词对计算一个得分，从而产生一个巨大的数据网格，这需要极高的内存和处理能力。这种瓶za（瓶颈）使得在标准硬件上高效运行这些强大的工具变得困难，尤其是在处理法律分析或医疗记录审查等现实世界任务所需的长序列时。

研究人员曾试图通过简化数学运算来解决这个问题。一种常见的方法是降低计算机使用的数字精度，即从高精度计算切换到更快速、低精度的计算。这就像是从用实验室天平测量食材，切换到使用厨房勺子；虽然速度快得多，但如果你在复杂的食谱中每一步都这样做，最终的成品味道可能会出错。另一种方法是完全跳过某些计算，忽略那些看似不重要的词。虽然这节省了时间，但却面临着切断模型理解故事所需关键连接的风险。挑战在于，如何找到一种既能利用低精度数学的速度，又不牺牲高精度数学的准确性，同时还能保持模型观察全局的能力。

一个研究团队推出了一种名为 TileMix 的新方法，它通过将注意力过程视为不是一个单一、统一的工作块，而是一组较小的、可管理的“磁砖”（tiles）来解决这一问题。想象一下那个巨大的词对得分网格是一个大型马赛克。TileMix 不再是用一种类型的油漆涂抹整个马赛克，而是将其划分为小的、与硬件对齐的正方形区域。对于每一个这样的区域，系统都会做出一个瞬时的决定：这一组特定的词对是需要高精度的慢速计算，还是可以采用快速的低精度计算？这个决定是基于一个预先规划的地图进行的，该地图将这些磁砖组合在一起，使得计算机可以在不停止重新组织数据的情况下，在同一次操作中切换高精度和低精度。

这项创新的核心是一个充当计算机处理器“交通指挥员”的路由系统。它将这些决策打包成紧凑的代码，本质上是一串比特流，告诉处理器每个磁砖应该走哪条路径。如果一个磁砖被标记为高精度，处理器就会使用其最精确的数学单元；如果它被标记为低精度，则切换到最快速、最高能效的单元。至关重要的是，两条路径都会更新一个共享的内存状态，以追踪整体结果，确保最终输出的一致性。这使得系统能够保留模型的完整连通性——这意味着词与词之间的交互从未被忽略——同时仍能从那些可以容忍低精度的计算部分中获取速度优势。

在实验中，研究人员在包括 LLaMA、Qwen 和 Vicuna 在内的几种流行大语言模型上测试了这种方法，使用的序列长度从 16,000 到 64,000 个词不等。他们将这种方法与标准的全文高精度方法以及一种全量使用低精度数学的版本进行了对比。结果显示，对整个计算过程使用低精度数学往往会导致准确度显著下降，导致模型丢失细节或在检索任务中失败。然而，TileMix 能够恢复大部分损失的质量。通过精心路由特定的磁砖组进入高精度路径，该系统能够保持非常接近全高精度基准的准确度水平，同时实现了更快的处理速度。

研究还探索了决定哪些磁砖应获得高精度处理的不同模式。他们发现，排列方式至关重要；某些模式（即将高精度计算保持在词网格的特定空间区域内）的效果比其他模式更好，具体取决于任务。例如，某些为局部词交互保留高精度的布局在事实召回任务中表现更好，而另一些布局在通用问答任务中则更具鲁棒性。研究人员证明，这种方法可以在无需重新训练模型的情况下应用，这意味着它可以立即部署在现有系统上。他们还展示了该方法在处理变长批次和不同模型架构时表现良好，表明它是一个提高效率的灵活工具。

最终，TileMix 提供了一种在速度与准确度之间的可控平衡。它表明，高效的长上下文处理的未来并不需要在这两者之间做选择。相反，通过智能地混合两者，系统可以以接近硬件理论极限的速度处理长文档，而不会带来使用全量低精度数学所导致的严重准确度惩罚。这种方法为在大规模数据集实时分析或需要即时理解长且复杂上下文的交互式工具等对速度和精度要求极高的场景中部署大语言模型，提供了一条切实可行的路径。

技术摘要：TileMix：用于 LLM 推理加速的以 Tile 为中心的混合精度注意力机制

问题陈述

长文本大语言模型（LLM）的推理在预填充（prefill）阶段面临着显著的瓶颈，这是由于密集自注意力机制具有二次方计算复杂度（ $O(L^2)$ ）。现有的加速策略通常分为三类，但各自存在局限性：

低精度量化： 如 INT8 量化等方法虽然提高了算术和内存效率，但通常在整个注意力矩阵中应用统一的精度路径。这往往导致空间精度路由无法在融合注意力循环之外进行，且统一的低精度执行可能会降低长文本质量。
基于稀疏性的方法： 通过选择活跃的 Token 交互（例如滑动窗口、稀疏模式）来减少计算，但这会改变注意力图的密集连通性，可能丢失关键的 Token 交互。
IO 感知融合注意力： 像 FlashAttention 这样的算子将注意力划分为硬件对齐的 Tile 以优化数据移动并融合操作（分数计算、在线 Softmax、数值聚合）。然而，这些算子传统上维持统一的分数计算路径，错失了在单次执行中将特定空间区域路由到不同精度水平的机会。

核心挑战在于：如何在不破坏密集 Token 连通性、不引入显著元数据开销的前提下，将异构精度（FP16 和 INT8）集成到单个融合注意力算子中，并确保不同的算术路径能够正确收敛到共享的在线 Softmax 状态。

方法论：TileMix

TileMix 引入了一种以 Tile 为中心的精度路由算子，将数值精度视为融合密集注意力计算中各组分数 Tile 的空间决策。

核心机制

Tile 组精度路由（Tile-Group Precision Routing）：
- 将注意力矩阵（ $L_q \times L_k$ ）划分为大小为 $BLOCK_M \times BLOCK_N$ 的硬件对齐计算 Tile。
- TileMix 并非按单个 Tile 分配精度，而是沿 Key 维度对相邻的 Key Tile 进行分组。一个单一的路由位控制一组 $g$ 个 Key Tile（其中 $g$ 为分组因子）。
- 这使得算子可以在同一个流式循环内，将每个 Tile 组分发至 FP16 路径（标准矩阵乘法）或 INT8 路径（使用带有 INT32 累加的 INT8 Tensor Core MMA）。
紧凑的位掩码路由（Compact Bitmask Routing）：
- 路由决策被打包进 64 位整数（位掩码）中，其中每一位对应一个 Key Tile 组。
- 这实现了在内层循环中使用移位和掩码操作进行常数时间查找（ $R_{hk,m,g_j} = ((b_{hk,m} \gg g_j) \& 1)$ ）。
- 元数据开销为 $O(H_k T_m)$ ，其中 $H_k$ 是 KV 头数， $T_m$ 是 Query Tile 数，即使随着上下文长度增加，该开销依然保持紧凑。
共享状态的异构执行（Shared-State Heterogeneous Execution）：
- FP16 和 INT8 路径都会更新一个共享的在线 Softmax 状态（行向最大值、归一化因子以及未归一化输出累加器）。
- INT8 路径计算 $Q_8 K_8^\top$ ，应用块缩放（block scales），并将结果重缩放到一个共同的浮点域，以便更新共享状态。
- 这种设计在保持密集注意力图（计算所有合法的 Token 交互）的同时，根据空间路由策略改变算术精度。
策略与布局（Policy and Layouts）：
- 路由策略由静态、无数据的结构化模板定义（例如：Band、Global、Row-Random、Aligned Sparse、BigBird、Sparse Transformer）。
- 这些模板根据目标 INT8 覆盖率（例如 25%、50%、75%）来决定注意力矩阵的哪些空间区域被路由至 FP16（高精度）以及哪些被路由至 INT8（低精度）。
- 系统支持分组查询注意力（GQA）、变长批处理以及 INT8 KV 缓存。

核心贡献

以 Tile 为中心的精度路由： 本文引入了将 Tile 组精度作为融合密集注意力的空间执行抽象，实现了在不稀疏化图形的前提下，对所有合法 Token 交互进行细粒度的 FP16/INT8 分配。
共享状态的异构分数执行： 设计了一个融合算子，使 FP16 和 INT8 分数路径对齐到一个共同的分数域，并通过单一的在线 Softmax 递归进行集成。
紧凑且可扩展的内核原生路由： 作者开发了带有常数时间内层循环查找的紧凑位掩码路由，以及可扩展的精度分组技术，在保持硬件对齐计算 Tile 的同时支持长上下文。
实用的长上下文推理： 该实现支持 GQA、变长批处理和 INT8 KV 缓存，并在多个模型家族（LLaMA、Qwen、Vicuna）及基准测试中得到了验证。

实验结果

评估是在 NVIDIA A100 GPU 上针对 LLaMA 3.2 3B、Qwen 2/2.5 7B 和 Vicuna 7B 进行的。

精度恢复： 统一 INT8 注意力（“One”）在上下文长度增加时，往往会导致检索和问答性能显著下降。TileMix 通过保留选定 Tile 组的 FP16 精度，成功恢复了这部分损失的质量。例如，在 LongEval 基准测试中，混合路由配置（如 SpTrans 或 BigBird 布局）的表现始终优于统一 INT8，通常能接近或匹配 FP16 基准线，同时保持较高的 INT8 覆盖率。
吞吐量效率： TileMix 提升了相对于标准 FP16 FlashAttention 的预填充吞吐量。在 4k tokens 时，像 SpTrans75（75% INT8）这样的配置达到了约 31.8 K tokens/s，而 FlashAttention 仅为约 14.3 K tokens/s。TileMix 配置在端到端吞吐量上通常优于或匹配统一 INT8 基准线，证明了路由和混合精度处理带来的开销相对于从 INT8 Tensor Core 中获得的收益而言是微不足道的。
数值稳定性： 对比固定 FP16 参考值的平均绝对偏差分析显示，偏差随 INT8 覆盖率的增加而增加，但仍处于可控范围内。路由布局（FP16 的空间分布）的选择对精度有显著影响，某些布局（如 SpTrans）在不同模型和任务中表现出较强的鲁棒性。
重击者保护（Heavy-Hitter Protection）： 研究表明，像 SpTrans 这样的静态路由布局能自然地保护高重要性的注意力交互（Heavy Hitters），即使在名义上的 INT8 覆盖率很高的情况下，也仅会对极小比例的高质量交互进行低精度路由。

意义与主张

本文声称 TileMix 确立了分数-Tile组精度路由作为一种内核原生的长上下文注意力抽象。其主要意义在于：

解耦连通性与精度： 它允许在保持密集 Token 连通性（不同于稀疏注意力）的同时，实现异构精度执行（不同于统一量化）。
可控的精度-效率前沿： 它提供了一种通过调整 INT8 覆盖率和路由布局来权衡精度与效率的实用机制，为不同模型家族提供了“可控的精度-效率前沿”。
无需训练的部署： 该方法无需重新训练，依靠静态空间模板来引导精度路由。
内核级优化： 通过将精度路由直接集成到融合注意力算子中并使用常数时间查找，它避免了因调用层级精度选择或外部调度带来的延迟惩罚。

作者总结道，空间精度控制是优化密集注意力算子的一个实用且有效的维度，它能在保持长上下文质量的同时，解决统一低精度方法带来的问题，并显著加速推理。