TileMix：用于LLM推理加速的以分块为中心的混合精度注意力

# 面向LLM推理加速的以分块为中心的混合精度注意力  
来源：https://arxiv.org/html/2608.17336  

## TileMix：面向LLM推理加速的以分块为中心的混合精度注意力  

乔张、曹庆雷、范恒  
机构：北得克萨斯大学计算机科学与工程系 LLaVi 实验室  
黄燕  
机构：北得克萨斯大学计算机科学与工程系 LLaVi 实验室  
沈可为  
机构：圣路易斯大学计算机科学系；北得克萨斯大学数据科学系  
\{hanzhi.zhang, heng.fan, yan.huang, kewei.sha, yunhe.feng\}@unt.edu  
\{qiao.zhang, qinglei.cao\}@slu.edu  
冯云河  
机构：北得克萨斯大学计算机科学与工程系 LLaVi 实验室  

###### 摘要  
大语言模型（LLM）中的长上下文预填充会带来巨大的计算与内存流量，因为密集自注意力需要计算二次方的查询-键分数。现有方法要么采用统一的低精度路径，要么选择性地连接词元交互，导致基于硬件对齐的分数分块空间精度路由无法融入融合的密集注意力中。我们提出TileMix——一种以分块为中心的精度路由内核，它将数值精度转化为在融合密集注意力内部对分数分块组进行的可执行空间决策。TileMix将注意力矩阵划分为硬件对齐的分数分块，将路由决策打包为紧凑的位掩码，并通过FP16或INT8分数计算路径分派每个分块组，同时两条路径更新共享的在线softmax状态。可扩展的精度分组使得每个路由位可以控制多个相邻的键分块，从而在长上下文场景下保持硬件对齐的计算分块和紧凑的元数据。通过路由所有合法的分块组，TileMix保留了密集的词元连接性，无需训练，并支持分组查询注意力、变长批处理和INT8键/值缓存。在LLaMA、Qwen和Vicuna模型的LongEval、LV-Eval及A100预填充基准测试中，TileMix恢复了在均匀INT8量化下损失的长上下文质量，并在FP16基础上提升了预填充吞吐量，从而在不同模型家族中实现了可控的精度-效率前沿。代码实现可通过 https://github.com/HanzhiZhang-Ulrica/TileMix 获取。  

## 1 引言  
Transformer模型越来越依赖长上下文进行文档摘要[7] (https://arxiv.org/html/2608.17336#bib.bib34)、多页问答[30] (https://arxiv.org/html/2608.17336#bib.bib35) 和检索增强生成[20] (https://arxiv.org/html/2608.17336#bib.bib36)，使得高效长序列处理成为实际LLM推理的核心。在预填充阶段，密集自注意力计算所有查询与键词元之间的交互，产生O(L²)的分数计算量（L为序列长度）。这种二次方计算使注意力成为长文档及其他上下文密集型工作负载的主要执行瓶颈。  

图1：注意力效率策略对比。颜色表示执行状态：蓝色为高精度，灰色为低精度，白色为被移除的交互。*量化*降低了权重/激活内存，但通常将注意力softmax和累加保持在较高精度。*稀疏/分块注意力*通过结构化或动态模式选择词元交互子集。*TileMix*保留所有合法的词元交互，并在一个融合注意力内核内部将分数分块组路由至FP16或INT8路径。  

现有加速方法主要优化数值格式、词元连接性或IO调度，如图1 (https://arxiv.org/html/2608.17336#S1.F1) 中的量化、稀疏性和IO感知融合注意力所示。(1) *低精度量化*（如INT8[33] (https://arxiv.org/html/2608.17336#bib.bib5); [38] (https://arxiv.org/html/2608.17336#bib.bib26)）提升了包括权重、激活和注意力在内的模型算子的算术与内存效率。量化注意力内核通常在每次调用或每个阶段使用单一算术路径，使得基于L×L分数分块的空间精度路由位于流式循环之外。(2) *基于稀疏性的方法*[35] (https://arxiv.org/html/2608.17336#bib.bib37); [37] (https://arxiv.org/html/2608.17336#bib.bib15); [3] (https://arxiv.org/html/2608.17336#bib.bib13) 通过选择活跃的词元交互来减少计算量。(3) *IO感知融合注意力*[11] (https://arxiv.org/html/2608.17336#bib.bib9) 将注意力划分为硬件对齐的分块，并融合分数计算、在线softmax和值聚合。这些内核利用分块进行数据移动和工作划分，但算术精度在每次内核调用或注意力阶段通常是固定的[25] (https://arxiv.org/html/2608.17336#bib.bib11); [4] (https://arxiv.org/html/2608.17336#bib.bib12)。  

这些方向共同提示使用硬件对齐的分数分块作为空间精度单位。融合内核可以在一次流式计算中通过多个算术路径路由分数分块组，同时保留完整的注意力图。现代融合注意力内核通过规则的Tensor Core分块以及协调的工作划分、在线softmax和跨GPU内存层次的数据移动实现高利用率[12] (https://arxiv.org/html/2608.17336#bib.bib10); [25] (https://arxiv.org/html/2608.17336#bib.bib11)。分块组精度路由必须协调FP16和INT8 Tensor Core的不同路径（包括INT8重缩放），然后两条路径更新共享的行最大值、归一化器和输出累加器[6] (https://arxiv.org/html/2608.17336#bib.bib40); [4] (https://arxiv.org/html/2608.17336#bib.bib12)。因此，精度分派落在延迟关键的内循环中，使得紧凑的分块对齐路由对于规则的长上下文执行至关重要，并带来了超出调用级精度选择的内核设计挑战。  

为解决这些挑战，我们提出TileMix——一种以分块为中心的精度路由内核，将异构分数算术集成到单个FlashAttention风格的执行流程中。对于每个查询分块行，TileMix加载打包的路由字，使用常数时间位操作解码每个键分块组的决策，并将QK计算分派到FP16或INT8 Tensor Core路径。重缩放后，两条路径更新共享的在线softmax状态，保持密集流式执行。可扩展的精度分组使得每个路由位可以控制相邻的键分块，同时随着上下文长度增长保持底层硬件对齐的计算分块和紧凑的元数据。该内核结合了全注意力的密集连接性与混合精度的算术灵活性（无需训练）。它还支持分组查询注意力、变长批处理和INT8键/值缓存，并暴露了FP16与均匀INT8注意力之间的可控精度-效率前沿。  

我们的贡献如下：  
1. **以分块为中心的密集注意力精度路由**：我们将分块组精度作为融合密集注意力的空间执行抽象，支持在所有合法词元交互上进行细粒度的FP16/INT8分配。  
2. **共享状态的异构分数执行**：我们设计了一种融合内核，将FP16和INT8分数路径对齐到统一的分数域，并通过一个在线softmax递推集成它们。  
3. **紧凑且可扩展的内核原生路由**：我们开发了带常数时间内循环查找和O(H_k T_m)元数据的打包位掩码路由，以及精度分组，可在保留硬件对齐计算分块的同时将路由扩展到长上下文。  
4. **实用的长上下文推理与评估**：我们实现了分组查询注意力、变长批处理和INT8键/值缓存支持，并通过长上下文检索、问答、预填充效率和数值分析，在LLaMA、Qwen和Vicuna模型上验证了TileMix。  

## 2 相关工作  
Transformer加速主要遵循三个方向：低精度量化、IO感知的分块注意力和结构化稀疏性。这些方法通过数值压缩、数据移动优化或选择性词元连接来降低长上下文推理成本。  

**低精度量化**通过以INT8[31] (https://arxiv.org/html/2608.17336#bib.bib1); [27] (https://arxiv.org/html/2608.17336#bib.bib2) 和INT4[40] (https://arxiv.org/html/2608.17336#bib.bib3) 等格式表示权重和激活来降低内存和算术成本。量化感知和训练后方法通过校准、激活变换、异常值处理和块级缩放来提高鲁棒性[34] (https://arxiv.org/html/2608.17336#bib.bib4); [33] (https://arxiv.org/html/2608.17336#bib.bib5); [24] (https://arxiv.org/html/2608.17336#bib.bib7)。最近的量化注意力内核将低精度分数计算、值聚合和数值近似集成到融合执行中[4] (https://arxiv.org/html/2608.17336#bib.bib12); [16] (https://arxiv.org/html/2608.17336#bib.bib23); [38] (https://arxiv.org/html/2608.17336#bib.bib26); [26] (https://arxiv.org/html/2608.17336#bib.bib39)。这些设计通常在张量、算子、注意力阶段或量化块级别分配格式，而二维分数分块组在流式循环内遵循单一算术路径[14] (https://arxiv.org/html/2608.17336#bib.bib6); [31] (https://arxiv.org/html/2608.17336#bib.bib1); [18] (https://arxiv.org/html/2608.17336#bib.bib8)。  

**IO感知注意力内核**[11] (https://arxiv.org/html/2608.17336#bib.bib9); [12] (https://arxiv.org/html/2608.17336#bib.bib10); [13] (https://arxiv.org/html/2608.17336#bib.bib38) 在SRAM驻留分块中处理注意力，并融合分数计算、在线softmax和值聚合，避免在HBM中完整材料化注意力矩阵。这些分块控制数据移动和工作划分，而算术精度通常在每次内核调用或注意力阶段是固定的[25] (https://arxiv.org/html/2608.17336#bib.bib11); [4] (https://arxiv.org/html/2608.17336#bib.bib12)。  

**稀疏注意力方法**通过滑动窗口、步进、扩张或混合局部-全局模式选择词元交互来减少计算[3] (https://arxiv.org/html/2608.17336#bib.bib13); [5] (https://arxiv.org/html/2608.17336#bib.bib14); [37] (https://arxiv.org/html/2608.17336#bib.bib15); [2] (https://arxiv.org/html/2608.17336#bib.bib16)。其他系统使用基于内容的选择和聚类[23] (https://arxiv.org/html/2608.17336#bib.bib17); [32] (https://arxiv.org/html/2608.17336#bib.bib18); [17] (https://arxiv.org/html/2608.17336#bib.bib19)、位置机制[10] (https://arxiv.org/html/2608.17336#bib.bib20); [39] (https://arxiv.org/html/2608.17336#bib.bib21)、自适应架构[8] (https://arxiv.org/html/2608.17336#bib.bib22)或动态构建的活跃块集[15] (https://arxiv.org/html/2608.17336#bib.bib32); [19] (https://arxiv.org/html/2608.17336#bib.bib33); [35] (https://arxiv.org/html/2608.17336#bib.bib37)。这些方法使用空间结构来选择执行的词元交互，在改变注意力连接性的同时保持精度位于选择决策之外。  

## 3 预备知识  
我们简要回顾 (i) 带在线softmax的FlashAttention风格分块注意力，以及 (ii) 用于低精度矩阵乘法的块级量化。它们的交互定义了TileMix解决的核心内核挑战：FP16和INT8分数分块遵循不同的算术路径，但贡献于一个共享的流式softmax状态。附录表4 (https://arxiv.org/html/2608.17336#A1.T4) 总结了本文中使用的符号。  

### 3.1 带在线softmax的分块注意力  
考虑一个自注意力头，其中 Q, K, V ∈ ℝ^{L×d}，L 为序列长度，d 为头维度。注意力计算分数矩阵 S ∈ ℝ^{L×L}、注意力权重 P ∈ ℝ^{L×L} 和输出 O ∈ ℝ^{L×d}：  
S = QKᵀ/√d，P = softmax(S)，O = PV。  
在GPU高带宽内存（HBM）中将 S 和 P 完整材料化需要 O(L²) 的中间内存占用，并在长上下文长度下产生大量 HBM 读写流量。FlashAttention流式处理键/值分块，同时保持分数和概率分块在片上，并仅将归一化输出写入HBM，为TileMix的分块组精度路由提供了执行基底。该内核将 Q 划分为 T_m = ⌈L/b_q⌉ 个分块 {Q_m}_{m=1}^{T_m}，将 (K, V) 划分为 T_n = ⌈L/b_{kv}⌉ 个分块 {(K_n, V_n)}_{n=1}^{T_n}。对于每个 Q_m，它在 n=1,...,T_n 上流式处理，同时维护共享的在线softmax状态 (m̃_mⁿ, z̃_mⁿ, Ō_mⁿ)，包括行最大值和归一化器 m̃_mⁿ, z̃_mⁿ ∈ ℝ^{b_q}，以及未归一化的累加器 Ō_mⁿ ∈ ℝ^{b_q × d}。状态初始化为 m̃_m⁰ = -∞, z̃_m⁰ = 0, Ō_m⁰ = 0，并更新如下：  
m̃_mⁿ = max{ m̃_mⁿ⁻¹, rowmax(S_mⁿ) },  
z̃_mⁿ = e^{m̃_mⁿ⁻¹ - m̃_mⁿ} z̃_mⁿ⁻¹ + rowsum( e^{S_mⁿ - m̃_mⁿ} ),  
Ō_mⁿ = e^{m̃_mⁿ⁻¹ - m̃_mⁿ} Ō_mⁿ⁻¹ + e^{S_mⁿ - m̃_mⁿ} V_n,  
S_mⁿ = Q_m K_nᵀ/√d。  
这里，rowmax(·) 和 rowsum(·) 在 S_mⁿ ∈ ℝ^{b_q × b_{kv}} 的 b_{kv} 列上进行归约。在分块 T_n 之后，将 Ō_m^{T_n} 除以 z̃_m^{T_n} 即可得到 Q_m 的输出。  

### 3.2 块级量化  
对于矩阵乘积 C=AB，块级量化独立对操作数块应用 ψ(·)，生成低精度表示及其缩放因子：(Â, δ_A) = ψ(A), (B̂, δ_B) = ψ(B), C ≈ δ_A δ_B (Â B̂)。表示 Â 和 B̂ 可使用INT8、FP8或其他低精度格式。我们使用INT8实例化 ψ(·)，因为NVIDIA A100 GPU提供了优化的INT8 Tensor Core原语以实现高吞吐量矩阵乘法。这些Tensor Core执行MMA（*矩阵乘加*）指令，在INT32寄存器中乘以INT8分片并累加部分和。将量化块与注意力对齐允许相应的缩放因子 δ_A 和 δ_B 在融合执行期间高效应用。  

FP16和INT8分数分块表现出不同的舍入、累加和重缩放行为。因为每个分数分块都贡献于共享状态 (m̃_mⁿ, z̃_mⁿ, Ō_mⁿ)，特定路径的数值效应通过行最大值跟踪、归一化和输出传播。