🤖 machine learning

Tile-Level Activation Overlap for Efficient LLM Inference

本文介绍了两个基于 CUTLASS 的专门化 SM90 核函数，它们在分块层面（tile level）将 SwiGLU 激活函数与矩阵乘法进行融合，以消除中间张量物化开销，在 NVIDIA H100 GPU 上实现了高达 2.47 倍的加速和 79.5% 的峰值利用率，同时在效率和数值精度方面均优于 PyTorch 和 cuBLAS。

原作者： Abhinav Jangda, Tyler Sorensen, Sebastian Burckhardt, Jianlan YE, Chaoyin Li, Atul Gupta

发布于 2026-07-07

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Abhinav Jangda, Tyler Sorensen, Sebastian Burckhardt, Jianlan YE, Chaoyin Li, Atul Gupta

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

想象一下，你正在经营着一家高速运转的工厂，专门制造复杂的机器人（这些机器人就是构建聊天机器人的 大语言模型）。为了制造每台机器人，你的工厂配备了一条特定的装配线，即 MLP（这是机器人大脑的核心部分）。

长期以来，这条装配线一直存在一个主要的低效问题。以下是该论文中提出的问题与解决方案，已进行简化说明。

问题：“中间人”瓶颈

在现代机器人工厂中，使用一种叫做 SwigLU 的特定步骤来让机器人变得更聪明。你可以把 SwigLU 想象成一种质量检查，它需要两个独立的计算步骤：

计算 A： 测量机器人的当前状态。
计算 B： 测量机器人的潜力。
胶水（The Glue）： 将这两个测量结果结合起来，得到最终结果。

旧的方法（瓶颈）：
在标准的工厂设置中（例如 PyTorch 中使用的设置），在完成计算 A 后，工人们必须将结果写在一块巨大的白板上（高带宽内存/HBM）。然后，他们必须走过去，阅读那块白板，进行计算 B，将那个结果写在第二块白板上，最后再走回来读取这两块白板以完成“胶水”步骤。

研究发现，对于较小的机器人（较小的 AI 模型），这种“走到白板前”的过程占据了总时间的 30% 到 37%。这就像一位厨师花费一半的时间在走动去储藏室拿单种调料，而不是在烹饪。

解决方案：两个全新的“超级厨房”

作者构建了两个全新的、定制设计的厨房（称为 算子/Kernels），它们彻底消除了白板的存在。他们不再需要写下结果并走回原处，而是让工人们将原材料保持在手中（在 寄存器/Registers 中），并以一种连续的动作完成所有工作。

他们针对不同的工厂规模设计了两种不同的策略：

1. “乒乓”厨房 (Kernel-1)

运作方式： 想象一场接力赛。当一名工人正在取下一批原料（加载数据）时，另一名工人已经在混合当前的这一批原料了。
诀窍： 他们使用了一种“乒乓式”调度方案。当机器正在忙于获取第二组原料时，工人们利用这段时间对第一组原料进行“胶水”数学运算。
最适用场景： 制造 巨型机器人（大型模型）或 同时运行许多机器人（大批次/Large Batches）的工厂。这就像一条庞大的装配线，有足够的工人来让大型机器保持满负荷运转。

2. “交错”厨房 (Kernel-2)

运作方式： 想象你正在打包箱子。你不是先打包完所有的红色物品，再打包所有的蓝色物品，而是完美地交替进行：先拿一个红色，再拿一个蓝色，如此循环。
诀窍： 他们在开始计算之前，先将两个权重矩阵（即计算的“配方”）混合在一起。这使得工人可以同时抓取“红色”和“蓝色”原料，并立即一起处理。
最适用场景： 制造 较小机器人 或 同时运行较少机器人（小批次/Small Batches）的工厂。这种方法非常高效，能确保工厂车间始终挤满了工作的工人，防止有人闲置等待。

结果：速度与精度

作者在 NVIDIA H100 芯片（目前最强大的 AI 处理器）上测试了这些新厨房，涵盖了各种机器人规模（从微型的 0.5B 模型到庞大的 72B 模型）。

巨大的加速： 对于较小的机器人，新厨房的速度比旧标准快了 2.47 倍。这就像把一个 10 分钟的任务缩短到了 4 分钟。
转移瓶颈： 旧系统是“受限于内存”（Spending too much time walking to the whiteboard），而新系统是“受限于计算”（Spending all its time actually doing the math）。他们达到了芯片最大理论速度的 79.5%。
编译器无法胜任： 作者尝试使用标准的“自动驾驶”软件（PyTorch 的 torch.compile）来自动修复这个问题。结果失败了。这个“自动驾驶”比他们的定制厨房要 慢 3 到 7 倍。这证明了对于这个特定问题，你需要人类专家来手工打造解决方案；计算机目前还无法自行解决。
更高的准确度： 出人意料的是，新的定制厨房也比标准方法更准确。标准方法在 4.5% 到 11% 的结果中存在微小的数学误差，而新厨房的 误差为零。

总结

论文表明，通过重新组织工厂车间的运作方式——特别是通过停止让工人们不断走到白板前去读写中间笔记——我们可以让 AI 模型运行得显著更快，尤其是对于用于边缘设备（如手机或笔记本电脑）的小型模型。他们证明了标准软件工具无法复制这种效率，需要专门的手写代码才能实现这些成果。

技术摘要：用于高效 LLM 推理的 Tile 级激活融合

问题陈述

现代大语言模型（LLM），包括 Qwen-2.5、LLaMA、Mistral 和 Gemma，主要使用 SwiGLU 激活函数。这种门控 MLP 结构计算 $Y = \text{SiLU}(A \times W_1) \odot (A \times W_2)$ ，需要两个独立的矩阵乘法（GeMM）以及随后的元素级门控激活。

在标准实现中，该过程会在 GeMM 与激活阶段之间在高速带宽内存（HBM）中实例化两个完整的中间张量（$Gate $和$ Up$）。在 NVIDIA H100 GPU 上的性能分析显示，这种中间张量实例化及相关的激活计算消耗了 MLP 执行总时间的 9–37%。对于较小的模型（如 Qwen-2.5 0.5B），这种开销最为严重，SwiGLU 占到了 MLP 时间的 30% 以上。随着通过量化（FP8, INT4）实现的 Tensor Core 算术密度增加，这些受限于内存的操作的相对成本也随之上升，从而形成了显著的瓶颈。

现有的编译器基础设施，特别是 PyTorch 的 torch.compile，无法解决这一问题。由于图级融合传递（graph-level fusion passes）的根本限制，它无法跨越具有不同权重矩阵的两个独立 GeMM 进行融合。实验表明，torch.compile 仅能达到 eager 模式下 PyTorch 性能的 34–94%，且显式的融合提示提供的改进微乎其微。

方法论

作者提出了两种互补的、基于 CUTLASS 的手工定制 SM90 内核，它们在 Tile 级别将 SwiGLU 融合进 GeMM，消除了中间 HBM 的物化过程。这两种内核都利用了 H100 (SM90) 上可用的张量内存加速器（TMA）和线程束特化（warp-specialized）调度。

1. 内核-1：通过 Pingpong 调度重叠同步 SwiGLU

机制： 该内核利用了 CUTLASS Pingpong 线程束特化调度，该调度将线程束划分为生产者（通过 TMA 加载数据）和消费者（执行 MMA）。
策略： 它创建一个 $[M, N]$ 线程块网格。在消费者阶段，在 Gate Tile 计算完成之后、Up Tile 数据到达之前，存在一个消费者线程束处于空闲状态的短暂窗口。该内核利用这个窗口在寄存器驻留的 Gate 累加器上计算 Swish 激活（ $\text{SiLU}(Gate)$ ）。
重叠： Swish 计算与生产者的 Up 权重 Tile 的 TMA 加载在时间上重叠。最后的元素级乘法（ $\text{SiLU}(Gate) \odot Up$ ）和存储发生在 Epilogue 阶段。
优化： 该设计针对大 Batch Size（ $M \ge 2048$ ）进行了优化，此时 $[M, N]$ 网格足以充分饱和 GPU 的流式多处理器（SM）。

2. 内核-2：通过自定义 Epilogue Visitor Tree 交织 SwiGLU

机制： 该内核引入了一个名为 PairMulStore 的自定义 Epilogue Visitor Tree (EVT) 节点。
策略： 它将权重矩阵 $W_1$ 和 $W_2$ 合并为一个单一的融合矩阵 $W_{fused} \in \mathbb{R}^{K \times 2N}$ ，其中列是交织的（偶数列用于 Gate，奇数列用于 Up）。内核在 $[M, 2N, K]$ 上启动一个标准的 GeMM，创建一个 $[M, 2N]$ 线程块网格。
融合： PairMulStore 节点在每个线程块内对 $[M, 2N]$ 累加器进行操作。它读取相邻的列对 $(D[m, 2n], D[m, 2n+1])$ ，应用 $\text{SiLU}(D[m, 2n]) \odot D[m, 2n+1]$ ，然后进行存储。
重叠： 该设计将 SwiGLU 计算与 Tile 存储交织在一起。当前 Tile 的算术运算与前一个 Tile 的异步 TMA 存储在计算与存储之间实现重叠。
优化： $2\times$ 的线程块数量增加显著提高了小 Batch Size（ $M \le 2048$ ）下的 SM 占用率，使其成为大多数配置下的首选默认方案。

核心贡献

Tile 级融合： 首个在 Tile 级别将 SwiGLU 融合进 GeMM 的实现，在每次推理步骤中消除了 4 个中间 HBM 操作（Gate 和 Up 的读/写）。
双内核策略： 一种互补的方法，使用 Kernel-1 处理大 Batch（顺序执行 Gate $\to$ Swish $\to$ Up）以及使用 Kernel-2 处理小 Batch（交织权重配合自定义 EVT），分别应对占用率和调度约束。
数值优越性： 融合内核与 FP32 参考值相比实现了零不匹配（0% 的元素相对误差 >1%），而基准 cuBLAS 实现由于在不同 Tile 间采用了激进的 BF16 部分累加，表现出 4.5–11% 的不匹配。
验证了编译器的局限性： 通过实证证据表明 torch.compile 无法复制此融合效果，其速度仍比手工定制内核慢 3–7 倍。

结果

在 H100 上对 Qwen-2.5 模型（0.5B 至 72B）进行的评估结果如下：

加速比： 这些内核相比 PyTorch eager 模式（cuBLAS + 分离的 SwiGLU）实现了高达 2.47 $\times$ 的加速。在 20 个基准测试配置中，Kernel-2 在 13 个配置中胜出，Kernel-1 在 7 个配置中胜出（主要针对高 Batch Size 下的大型模型）。
范式转移： Roofline 分析确认，融合将工作负载从受限于内存（memory-bound）转向受限于计算（compute-bound）。对于小型模型，算术强度增加了高达 247%，使运行点跨越了 Roofline 脊线。
利用率： 在 H100 上，峰值 BF16 利用率达到 79.5%，几乎是基准方案的两倍。
效率： 对于小型模型，融合效率（实际加速比 vs 理论内存节省）接近 100%。对于 Kernel-2，由于存在计算-存储重叠收益，其效率在某些配置下甚至超过了 100%。
可扩展性： 随着模型规模增大，加速比会逐渐减小（0.5B 模型为 2.47 $\times$ ，而高 Batch Size 下的 72B 模型约为 ~1.0 $\times$ ），因为 GeMM 计算主导了总时间，且中间内存流量的相对比例降低。

意义

本文认为，对于现代 LLM 推理，特别是针对边缘部署和移动端规模的模型（0.5B–14B），瓶颈已从原始计算转向由中间张量物化引起的内存流量。作者证明了手工定制、硬件感知的 Tile 级设计对于克服当前编译器基础设施的局限性是必要的。通过消除中间 HBM 访问并将激活计算与数据移动重叠，这些内核显著降低了推理延迟并提高了计算利用率，为对延迟和吞吐量有严格要求的生产级服务系统提供了关键的优化方案。