Towards Free Normalization: Fusing Normalization into GEMM and Attention Kernels

向免归一化迈进：将归一化融合到GEMM和注意力内核中 – PyTorch

特色项目

TL;DR

在本文中，我们介绍了针对LayerNorm和RMSNorm等常见归一化操作的多种创新内核融合技术，通过减少这些内存密集型内核的内存-IO开销，显著提升了性能。我们首先简要概述归一化操作在LLM和广告推荐模型中的建模重要性及性能挑战，然后介绍解决性能瓶颈的新策略，包括Lazy Pre-Norm和多CTA归一化融合。我们证明，通过与GEMM融合，这些技术可隐藏多达90%的归一化内核延迟。最后，我们提出了FlashNormAttention算法，在GDPA [1]等注意力内核周围应用多归一化融合，实现高达35%的内核加速。

本工作主要使用两种内核DSL实现：TLX，一组带有GPU执行控制硬件感知支持的Triton DSL扩展；以及Helion，一个在开发者效率、可移植性和全面自动调优方面表现优异的高级DSL。基准测试在Meta数据中心的NVIDIA B200 GPU上进行，使用bfloat16数据类型，功率限制为750 W。

代码可在以下地址获取：https://github.com/facebookresearch/ads_model_kernel_library/tree/main/multi_cta_norm_fusion 和 https://github.com/facebookresearch/ads_model_kernel_library/tree/main/gdpa_megakernel

引言

由于归一化技术在稳定训练和加速收敛方面的卓越效果，它已成为大多数深度学习架构中不可或缺的组成部分。特别是，沿最内层嵌入维度的传统归一化（如LayerNorm、RMSNorm）已成为现代大型语言模型及Meta广告模型等推荐系统模型中最常见且普及的类型。例如，在部署于Meta最大推荐系统训练基础模型Kunlun [2]上的生成广告模型（GEM）[3]中，LayerNorm/RMSNorm几乎存在于所有关键组件中，如多头注意力、分层种子池化和GDPA增强[1]的PFFN。

然而，归一化的普遍性也带来了严峻的性能挑战：其高度依赖内存且无法利用TensorCore。这阻碍了我们在模型训练中充分利用硬件计算能力。以Kunlun [2]为例，归一化操作占据了该模型总训练延迟的约20%。这意味着在未优化的情况下，我们立即损失了20%的硬件计算吞吐量。在计算密集型的典型LLM中，归一化仍可能占据总延迟的约10%。

为解决这一问题，我们必须以IO感知的方式设计归一化内核，在不牺牲计算精度的前提下，通过内核融合谨慎节省内存IO成本，同时将这些内存/CUDACore密集型操作与TensorCore密集型操作重叠。由于大多数归一化操作都跟随或先于矩阵乘法操作（如MLP、Attention），我们的工作重点在于如何高效地将归一化与矩阵乘法融合。我们首先描述多种高效融合归一化与单个GEMM的策略，最后介绍FlashNormAttention算法，将LayerNorm和RMSNorm同时融合到注意力操作中。

注意：在以下基准测试结果中，我们默认禁用逐元素仿射变换，除非另有说明。因为我们发现这些操作会带来显著的性能开销，而对模型质量的提升却非常有限。此外，除非另有说明，所提出的核优化和算法思想在存在/不存在逐元素仿射变换的情况下均适用，尽管性能结果可能有所不同。

1. 归一化融合的挑战

概述：在本节中，我们将讨论以常规方式将归一化与计算密集型核（如GEMM）融合时面临的挑战，这些挑战本质上源于不同的分块策略。随后我们将展示一种"朴素"的融合方案，该方案强制GEMM算法遵循与归一化算法相同的分块方式，并观察到该方案在N非常小时表现良好，但随着N增大，由于分块约束和效率问题，性能会变得次优甚至不可行。

与标准激活融合（如GEMM+ReLU）相比，归一化融合的根本挑战在于分块策略的差异。归一化本质上是一种需要访问整个维度数据才能计算正确结果的归约操作。特别是对于LayerNorm和RMSNorm，典型核通常沿外层维度进行分块，而非内层维度，这意味着每个CTA始终需要加载完整的数据行。相比之下，典型GEMM在两个维度上都进行分块，每个分块不跨越整个行，这使得后续的逐行归一化变得不可能。

解决这个问题最直接的方法是扩展GEMM的分块尺寸，使每个分块跨越整个内层维度。对于典型的(MxK)@(KxN) GEMM，这意味着沿N维度的分块尺寸必须大于N（通常是N的下一个2的幂）。高层算法如以下图表所示：

这种方法存在两个主要问题：

• 与纯GEMM的最优分块策略相违背，这会导致GEMM本身的性能下降，因为缓存行为、流水线行为等会变得次优。

• 对输入形状施加了硬性限制，特别是N的大小。N过大时无法放入共享内存。

以下是一些粗略估算，帮助理解N能有多大。假设使用Blackwell GPU（共享内存大小228KB），数据类型为bfloat16，且为高效流水线/重叠至少需要2个流水线阶段。进一步假设M维度和K维度的最小分块尺寸均为32。那么我们有：

2 stages x 2 bytes / element x (tile_m x tile_k + tile_k x tile_n + tile_m x tile_n) < 228KB
=> 32 x 32 + 32 x tile_n + 32 x tile_n < 228KB / 4
=> 512 < tile_n < 1024

由于分块尺寸通常应为2的幂，这限制了tile_n（进而限制N）的最大值为512，否则该核将无法运行。

我们将在后续章节讨论绕过这些限制的方法。但尽管存在这些限制，我们发现对于较小的N值，这种融合策略仍能带来显著的性能提升。在该实验中，我们选择了Helion，因为它具有高开发效率和详尽的自动调优功能，这在像这种分块尺寸受到硬性约束的非典型情况下非常有帮助。

以下是广告模型中常见输入形状的基准测试结果。请注意，延迟节省是根据torch inductor归一化内核的延迟百分比计算的。我们这样做的目的是使该指标独立于基础GEMM内核的延迟（以及其与归一化内核的对比），并天然捕捉此类融合尝试的优化空间（即100%表示理论上最佳的完全重叠归一化与GEMM的情况）。

对于64和128等小形状，这种融合策略可为LayerNorm内核带来显著的17%-32%延迟节省。然而，当K/N超过128时，收益开始消失甚至出现严重倒退。这是因为随着N增大，强制tile_n=N的设定会与未融合GEMM内核的理想分块尺寸产生更大偏差，而节省内存IO的收益逐渐被扭曲基础GEMM算法带来的损害所抵消。

2. 惰性预归一化：一种新型预归一化与线性层融合技术

概述：在本节中，我们介绍一种新型的前导融合技术，用于将预RMSNorm融合到GEMM内核中。我们将讨论此类前导融合的动机及挑战，并提出一种名为惰性预归一化（Lazy Pre-Norm）的创新算法，通过巧妙延迟部分预归一化计算到GEMM完成之后，从而有效解决这些挑战并实现性能提升。

我们提出的第一个方案通过将预归一化与后续GEMM融合（作为前导融合）解决了前述问题。尽管前导融合通常应避免，但仍有几个值得探索的原因：

• 前导融合绕过了尾部融合中遇到的分块问题，即GEMM内核中的每个CTA无法访问输出张量的完整行。相比之下，每个CTA只需通过算法就能扫描输入张量A的完整行！

• 实际上，预归一化在后归一化中变得越来越普遍，尤其是在大型语言模型中。

为了使前导融合有效，我们为预归一化的一种特殊情况（无逐元素仿射变换的RMSNorm）设计了名为惰性预归一化的优化技术。具体来说，我们希望融合以下计算：C = rmsnorm(A) @ B

其中 rmsnorm(A) = A * rstd(A)[:, None]，而 rstd(A) = rsqrt((A ** 2).sum(dim=-1) / A.shape[-1] + 1e-5)

该融合的关键难点在于：在典型的分块GEMM中，虽然最终可以访问完整行（这使我们能够计算归一化结果rstd），但需要分块访问，而实际上我们处理每个分块时都需要rstd！这会形成循环依赖：我们需要等到k循环结束才能计算rstd，但又需要rstd才能开始处理循环！

为了解决这个问题，关键观察是这两个相互依赖的部分本质上属于不同类型的计算：归约和逐元素应用。

• rstd计算部分是对整行的归约操作。

• 使用rstd进行归一化则是对A中每个元素的逐元素计算。

让我们分别处理这两个组件。对于归约部分，首先需要注意到它本身并不阻塞任何操作，这是一个很好的特性，因为它意味着我们可以并行计算它与TensorCore计算。由于每个CTA自然沿着矩阵A的内维度进行扫描，我们可以在并行于矩阵乘法的同时，沿着该维度累积A的平方和。

逐元素部分则更具挑战性，因为它依赖于归约的结果，从而导致循环依赖。这里的解决方案是基于一个关键观察的数学技巧：在无仿射变换的RMSNorm中，逐元素乘法实际上是对行的乘法，即矩阵A同一行的所有元素都乘以相同的rstd。这引出了以下关键性质：

(A * rstd[:, None]) @ B = (A @ B) * rstd[:, None]

证明：行乘法等价于M @ A，其中M是一个对角矩阵。因此我们有 (A * rstd) @ B = (M @ A) @ B = M @ (A @ B) = (A @ B) * rstd

这非常棒，因为它意味着逐元素计算可以"延迟计算"，直到整个k循环完成，从而有效成为后处理阶段！将所有内容整合起来，下面是Lazy Pre-Norm算法的核伪代码：

def GEMM_norm_fusion_kernel(A, B, C):
	compute the m_tile and n_tile of this CTA
	square_sum = zeros(m_tile)
	acc = zeros(m_tile, n_tile)
	for each k_tile:
		tile_A = A[m_tile][k_tile]
		tile_B = B[k_tile][n_tile]
		acc += tile_A @ tile_B
		square_sum += (tile_A * tile_A).sum(-1) # computed in parallel to the GEMM!
	rstd = rsqrt(square_sum / A.shape[-1] + 1e-5)
	acc *= rstd[:, None]
	C[m_tile][n_tile] = acc

请注意，虽然每个k迭代仍需要一些额外计算，但这些计算可以与矩阵乘法并行执行。通过warp专业化，核的warp划分和执行方式如下：

请注意，该算法仍然具有前处理融合的关键缺点：RMSNorm计算在多个CTA中被冗余执行（考虑所有CTA计算相同行但不同列的输出张量；更多细节见第3节）。然而，由于Lazy Pre-Norm确保了大部分计算与TensorCore完全重叠，这种冗余是可以接受的，并且仍能带来良好的性能提升。

关于Lazy Pre-Norm算法需要注意的几个限制：

• 它难以支持逐元素仿射变换，因为这些变换作为列乘法起作用。这会破坏我们对逐元素操作必须是行乘法的先决条件。

• 它不适用于LayerNorm，因为LayerNorm的逐元素部分涉及减法，不是简单的行乘法。

• 该融合的反向传播实现会很复杂，因为在正向传播中我们从未显式计算rmsnorm(A)。因此，在计算dA和dB时，我们需要从A和rstd实时重构rmsnorm(A)。

3. 多CTA范数：将后处理范数与线性层融合为后处理阶段

概述：尽管Lazy Pre-Norm前处理融合带来了良好的加速效果，但它仍存在局限性，无法推广到大多数范数使用场景。在本节中，我们将讨论一种更通用的技术，用于将后处理范数与GEMM融合，并回到后处理融合的领域，直接解决第1节中提出的分块不匹配问题，使用CTA集群和分布式共享内存。

我们借鉴了Quack的思路，并将其扩展到融合内核中。Quack的归一化内核利用CTA集群将大N划分到同一集群的不同CTA中，并通过分布式共享内存相互通信，协作完成对N的单次归约。这使我们能够让多个CTA协作处理相同的数据行，并根据归一化需求进行必要的通信，而无需承担全局内存IO的成本。

如上所述，大多数归一化操作都可以分解为归约部分（例如RMSNorm的rstd、LayerNorm的均值和方差）和后续的逐元素部分，后者利用归约结果。只有归约部分需要遍历整个N维度，而我们通过CTA集群对其进行划分和征服。由于归约结果通常较小（毕竟这是归约后的结果），因此只需要极小的分布式共享内存通信开销即可在CTA之间发送/接收该结果。

请注意，这个思路正好解决了我们与归一化融合面临的相同问题——仅仅是N太大了！（尽管N过大的原因和阈值不同）。这意味着我们只需将这个多CTA算法直接放入GEMM的后记部分，融合就完成了！

def GEMM_norm_fusion_kernel(A, B, C):
	计算该CTA的m_tile和n_tile
	acc = zeros(m_tile, n_tile)
	for each k_tile:
		tile_A = A[m_tile][k_tile]
		tile_B = B[k_tile][n_tile]
		acc += tile_A @ tile_B
	acc = multi_cta_norm(acc) # 此处发生DSMEM通信
	C[m_tile][n_tile] = acc

需要指出的是，这种融合绝非没有代价，除了引入DSMEM开销外，还对内核施加了一些限制，这可能会导致我们融合的基GEMM内核性能出现退化：

• 它对CTA调度施加了严格的限制。具体而言，集群中相邻的CTA必须共享相同的m_tile但不同的n_tile。

• 由于上述限制，配对CTA矩阵乘法变得困难

• 同样由于上述限制，无法进行瓦片超级分组[5]

• 这仍未解除N无限大的限制。我们仍然受单CTA限制（约512）与最大集群规模的乘积所约束。在Blackwell上，可移植的最大集群规模为8。这将N限制在最多4096。

尽管如此，节省大量内存IO带来的好处仍远超这些限制。我们选择了TLX作为该内核的实现方案，它在灵活性/开发效率与底层硬件控制之间取得了良好的平衡，这在本案例研究中至关重要。我们在TLX GEMM内核基础上通过线程束专业化构建了融合内核。我们针对广告建模中的一些常见形状（M = 256k, K = O(512), N = O(512))进行了基准测试，取得了以下性能结果。

请注意我们将K和N限制在2048，因为当它们达到4096时，延迟完全由GEMM主导，归一化操作所耗时间不足总延迟的5%。

反向传播呢？融合重组。

Overview: 在本小节中，我们讨论在反向传播中实现相同融合思想时面临的额外挑战：正向传播的尾部融合在反向传播中自然转变为前导融合。我们将探讨前导融合的关键问题，并提出一种新颖的解决方案，该方案在正向传播中将归一化操作与不同的GEMM融合，而在反向传播中则实现高效的尾部融合。

LayerNorm和RMSNorm的反向传播也涉及归约操作，这在正向传播中已通过类似方式解决，因此不会成为大问题。高效的反向传播计算还需要从正向传播中存储的中间归约结果，由于归约结果是一维的且IO开销极小，这同样不会造成显著问题。真正的挑战在于正向传播的尾部融合在反向传播中会转变为前导融合。

# forward formula
C = norm(A @ B)

# backward formula
dA = norm_backward(dC) @ B.T
dB = A.T @ norm_backward(dC)

请注意在反向传播中，norm_backward计算发生在GEMM之前，这可能导致潜在的前导融合。我们在第2节已讨论过前导融合的一些通用缺点，但在此特定情况下，让我们通过一个潜在的前导融合方案来更深入理解为何这会成为问题。

def GEMM_norm_bwd_fusion_kernel(dC, BT, dA):
	compute the m_tile and n_tile of this CTA
	acc = zeros(m_tile, n_tile)
	for each k_tile:
		tile_dC = dC[m_tile][k_tile]
		tile_dC = multi_cta_norm_bwd(tile_dC) # where DSMEM communication happens
		tile_BT = BT[k_tile][n_tile]
		acc += tile_dC @ tile_BT
	dA[m_tile][n_tile] = acc

此方法存在多个性能问题：

• 归一化反向计算位于关键路径上，会阻塞每次迭代的GEMM计算！这与尾部融合形成鲜明对比，尾部融合只需在主循环结束后执行一次计算。

• 归一化反向计算存在冗余。请记住，每个dC行都由不同CTA单独加载，这些CTA计算不同dA分块（共享相同m_tile但不同n_tile）。每个CTA都需要对相同dC分块执行相同的归一化反向计算。

• 该融合内核仅计算dA，但还需要计算dB，这需要再次执行norm_backward(dC)，导致更多冗余计算。

解决方案是什么？除非是第2节讨论的某些特殊情形，否则我们几乎无能为力。因此应避免前导融合。为此我们需要在融合策略上更具灵活性：由于归一化层通常位于线性层之间，为什么不将归一化操作在正向传播中与不同的线性层融合，而在反向传播中也采用相同策略？

现在很容易看出，反向传播也变成了尾部融合，归一化反向计算不再存在冗余。更重要的是，融合结构与正向传播完全一致！只需将正向传播内核中的multi_cta_norm替换为multi_cta_norm_bwd，即可获得反向传播内核。以下是与正向传播相同形状的基准测试结果。

请注意，要使这个方案奏效，线性层不一定是严格意义上的线性层。例如在LLM架构中，我们可能看到attention -> norm -> linear或相反的模式。只要归一化层两侧是计算密集型操作（且适合融合归一化操作），这些模式仍可采用相同优化技术。

下一节将讨论将归一化操作融合到注意力机制中的示例。

4. FlashNormAttention：将前归一化和后归一化融合到FlashAttention变体内核中

概述：在本节中，我们将研究如何将前述融合思想应用于GDPA内核[1]的案例研究，并介绍FlashNormAttention算法。GDPA内核在Meta广告模型中被广泛使用，特别是在Kunlun[2]架构中，它是从FlashAttention[6]重新设计的通用注意力内核。因此，下文讨论的大多数优化思想均可推广到其他注意力内核（如FlashAttention）。该算法使用与上述多CTA GEMM+归一化融合完全相同的算法技巧，但复杂度更高（注意力机制 vs GEMM，融合两个归一化 vs 一个）。为实现高性能，采用了多种优化技术，包括：

• SMEM / TMEM 重用以降低内存压力
• 寄存器子分块以降低寄存器压力
• 精细调整的warp专业化以并行化繁重的CUDACore操作
• 在反向传播中重新计算归一化以避免前向传播中保存额外张量的IO成本
• 使用TMA_REDUCE_ADD和TensorCore Accumulate等高级硬件特性以提高流水线效率

Kunlun[2]中典型的PFFN模块以GDPA内核为核心，但其周围还包含若干归一化和残差连接。下图展示了PFFN模块内部的数据流动方式。

请注意，GDPA内核前后分布着大量IO密集型操作（归约和逐元素核），这使得整体计算效率较低。我们的目标是将所有这些操作融合到一个单一内核中，我们称之为FlashNormAttention。这类似于执行模块内所有操作的“超级内核”，但与原始超级内核[7]在意图和意义上存在差异：它不仅旨在减少内核启动开销，更重要的是减少与HBM之间的数据传输总量，而后者通常是更严重的瓶颈。

此处的整体融合方案保持一致：利用CTA集群和分布式共享内存协同计算归一化。需要注意的是，此处的GDPA是多头注意力机制，归一化操作在所有头之间进行。因此，尽管在典型的GDPA/FA内核中，单个CTA可以访问完整的头维度，但仍需要其他头的数据进行归一化，这就需要多CTA归一化算法。

让我们从原始GDPA内核算法开始，其伪代码如下。为简化起见，仅包含与融合相关的最基本部分。详细且优化的算法请参见原始GDPA博客。

# 输入：Q: [batch_size, seq_len_q, H, head_dim], K/V: [batch_size, seq_len_kv, H, head_dim]
# 元参数：BLOCK_M（seq_len_q上m_tile的分块大小），BLOCK_N（seq_len_kv上n_tile的分块大小）
# 网格布局：(batch_size, seq_len_q // BLOCK_M, H)
def gdpa_fwd_kernel(Q, K, V, output):
	计算当前CTA的batch_idx、m_tile和head_idx
	q = Q[batch_idx, m_tile, head_idx, :] # [BLOCK_M, head_dim]，B和H维度已索引
	acc = zeros(BLOCK_M, head_dim)
	for n_tile 遍历整个seq_len_kv:
		k = K[batch_idx, n_tile, head_idx, :] # [BLOCK_N, head_dim]，B和H维度已索引
		v = V[batch_idx, n_tile, head_idx, :] # [BLOCK_N, head_dim]
		p = elementwise_activation(q @ k.T) # [BLOCK_M, BLOCK_N]
		acc += p @ v # [BLOCK_M, head_dim]
	output[batch_idx, m_tile, head_idx, :] = acc

在深入融合工作之前，让我们先对算法进行一个关键的调整。在 Kunlun 的 [2] 用例中，我们观察到 seq_len_q 通常较大（O(1k)），而 seq_len_kv 通常较小（O(128)）。这使得内循环的流水线非常浅，并暴露了前导和尾部的开销。为了改善这种情况的性能，我们交换了 Q 和 K/V 在内核中的角色，以 KV 为网格进行划分，并对 Q 进行循环。请注意，只有在不按 seq_len_kv 进行分块时，该算法在数值上才是正确的。

# 输入: Q: [batch_size, seq_len_q, H, head_dim], K/V: [batch_size, seq_len_kv, H, head_dim]
# 元参数: BLOCK_M (seq_len_q 上 m_tile 的分块大小)
# 在 (batch_size, H) 上进行网格划分
def gdpa_fwd_kernel_short_kv(Q, K, V, output):
	计算当前 CTA 的 batch_idx 和 head_idx
	k = K[batch_idx, :, head_idx, :] # [seq_len_kv, head_dim], B 和 H 维度已索引
	v = V[batch_idx, :, head_idx, :] # [seq_len_kv, head_dim]
	for m_tile 遍历整个 seq_len_q:
		q = Q[batch_idx, m_tile, head_idx, :] # [BLOCK_M, head_dim], B 和 H 维度已索引
		p = elementwise_activation(q @ k.T) # [BLOCK_M, BLOCK_N]
		output[batch_idx, m_tile, head_idx, :] = p @ v # [BLOCK_M, head_dim]

虽然此优化与我们当前的主题没有直接关系，但我们仍在此包含它，因为我们是在此版本的基础上构建了融合内核，并且基准测试也是针对此版本进行的。为了完整性，我们在此包含它，因为原始 GDPA 博客中并未提及这一点。

现在我们将归一化和残差项融合到内核中。其思路与上述方法相同，但使用了多 CTA 的归约。唯一需要注意的是，同一集群中的 CTA 应该共享相同的 batch_idx 并处理不同的 head_idx。同时请注意，此处的 layernorm 是前导融合。虽然流水线和矩阵乘法对其的依赖性仍然是一个问题（我们将在下文解决），但幸运的是，由于 K/V 较短，我们不会遇到冗余的 layernorm 计算问题。由于我们不会对 K/V 的长度进行分块，每个 Q 的分块只会被一个 CTA 加载和处理。

# 输入: Q: [batch_size, seq_len_q, H, head_dim], K/V: [batch_size, seq_len_kv, H, head_dim]
# 元参数: BLOCK_M (seq_len_q 上 m_tile 的分块大小)
# 在 (batch_size, H) 上进行网格划分
def gdpa_fwd_fusion_kernel_short_kv(Q, K, V, output):
	计算当前 CTA 的 batch_idx 和 head_idx
	k = K[batch_idx, :, head_idx, :]
	v = V[batch_idx, :, head_idx, :]
	for m_tile 遍历整个 seq_len_q:
		q = Q[batch_idx, m_tile, head_idx, :]
		ln_q = multi_cta_layernorm(q) # 多 CTA 归一化 1
		p = elementwise_activation(ln_q @ k.T)
		gdpa_out = p @ v
		gdpa_out += ln_q # 残差连接 1 
		out = multi_cta_rmsnorm(gdpa_out) # 多 CTA 归一化 2
		out += q # 残差连接 2
		output[batch_idx, m_tile, head_idx, :] = out

尽管这一切看起来都很不错，但在伪代码背后隐藏着两个关键问题：

• 内存压力：这种大规模融合对寄存器和共享内存的使用带来了巨大压力。我们需要保留许多以前不存在的内容，例如 ln_q 和 rmsnorm(gdpa_out)。如果执行过程是纯粹的顺序执行，这会没有问题，因为当我们为某个操作生成输出时，其输入可以立即释放。但由于残差连接的存在，这里的情况并非如此。请注意 q 和 ln_q 的生命周期跨越了大范围区域，因为我们需要保留它们以供后续的残差连接计算使用。这意味着我们一定需要为这些变量分配一些共享内存，从而增加了总内存需求。事实上，使用这种原始算法版本时，我们观察到共享内存使用量翻了一番，显著超过了限制。

• CUDA 核心主导与流水线停滞：尽管我们通过融合消除了大部分内存 I/O，但归一化和残差连接的 CUDA 核心计算仍然存在，并阻碍了 Tensor Core 的利用率。需要精细调整 warp 专业化和流水线技术，尽可能隐藏 CUDA 核心的延迟。

针对内存压力，我们应用了三个主要优化思路：

• 内存缓冲区重用：节省内存使用的关键技术是让不重叠的数据共享相同的内存缓冲区。在我们的情况下，一个很好的例子是 out 的共享内存缓冲区。在上面代码的最后一行，看起来我们是直接从寄存器将 out 存储到 HBM，但实际上发生的是我们首先将其放入 SMEM 缓冲区，然后调用 TMA 异步地从 SMEM 存储到 HBM。显然，这个缓冲区的生命周期很短，我们将其重用于在计算完 ln_q 后、在需要用于第二次残差连接之前临时存储 q。

• 利用 Tensor Memory 和 TensorCore 的累加功能：请注意 ln_q 是立即被加到 matmul p @ v 的结果中的。而不是将 ln_q 保留在 SMEM 中并在 matmul 完成后读取出来，注意这正好是 tcgen05 在 TMEM 中支持的 MMA 语义。因此，我们可以直接将 ln_q 保留在为 p @ v 分配的 TMEM 缓冲区中，并将 ln_q 的加法卸载到 TensorCore！这有助于节省 SMEM 占用空间和计算时间。

• 寄存器子分块：特别是在我们这种以 CUDA 核心为主的内核中，寄存器的稀缺程度高于 SMEM/TMEM。除了仔细调整寄存器分配外，我们还使用寄存器子分块来缓解寄存器压力。我们将 SMEM/TMEM 中的张量分割成块，每次加载一个块到寄存器中进行归一化计算（包括归约和逐元素计算）。这有助于防止寄存器溢出，这会导致显著的性能下降。

针对流水线停滞，我们使用了以下技术来提高流水线效率：

• 线程束专业化：在原始GDPA设计中，我们有4个主要的专用线程束分区（加载、mma、激活和尾声）。在FlashNormAttention中，我们将RMSNorm计算放在激活线程束上，同时新增一个专用于前导层归一化计算的第五个分区，以便更好地与TensorCore以及其他CUDA Core操作（如上一次迭代的RMSNorm计算）重叠。执行流程如下：我们为激活分区使用8个线程束（0-7），为层归一化使用4个线程束（8-11）。我们为激活线程束最大化寄存器分配，同时通过寄存器子分块技术限制层归一化线程束的寄存器分配，以达到寄存器使用与延迟的最优平衡。

• 寄存器预加载：执行流程中的关键瓶颈是残差连接引入的复杂数据依赖性。q和ln_q都需要长时间保留在内存中，阻碍了这些张量在下一次迭代中的预取和预计算。由于TMA预取使用SMEM，上述关于内存压力的优化2有效将ln_q移出了关键路径。对于q，我们让激活线程束在q就绪后立即从SMEM预加载q，并一直保留到第二次残差连接发生时。通过这种方式，我们可以立即释放q的SMEM，使下一次迭代的q加载可以在处理当前迭代时进行。由于我们为激活线程束分配了最大数量的寄存器，寄存器压力仍然保持在合理范围内。

以下是我们在典型GDPA形状上的基准测试结果。K/V是长度精确为128的密集序列。Q是平均稀疏度为0.5且最大长度可变的稀疏序列。批处理大小为768。头维度设置为128，头数量也根据不同的归一化维度和不同的CTA集群大小进行调整，以反映不同的性能表现。由于复杂的融合操作，我们以总基线延迟的百分比形式展示延迟节省，而不是仅展示归一化/逐元素内核的延迟。基线使用inductor编译。

反向实现

关于反向传播，为简洁起见我们略去详细算法，仅说明有趣的优化点以及与正向传播的相似性/差异性。

• 算法：首先注意到反向传播的融合与正向传播非常相似，因为归一化等约减操作的反向传播也是约减操作，残差连接的反向传播也是残差连接。

• 重新计算：原始GDPA反向传播内核通过重新计算q@k来节省正向传播的内存IO。我们为融合采用相同思路。首先，在我们的内核中q@k变为ln(q)@k，因此我们首先重新计算ln(q)。我们通过在正向传播中存储均值和方差（由于它们是一维的，因此成本很低）来避免DSMEM开销，这样反向传播可以轻松使用它们推导出ln(q)。其次，我们需要rmsnorm的输出来进行反向计算，我们通过rmsnorm_out = kernel_out – q重新计算。最后，我们还需要rmsnorm的输入用于残差连接的反向传播，我们通过rmsnorm_in = rmsnorm_out / rstd重新计算。

• Warp 特化：原始 GDPA 反向传播内核包含 4 个 warp 分区：mma、激活、加载和归约（用于原子添加 dQ）。我们沿用相同结构，将除 Layernorm 反向传播外的所有新计算放入激活分区。我们将 Layernorm 反向传播计算放入归约分区，以实现更高效的流水线（类似于前向传播中将前导 Layernorm 放入独立分区的做法）。

• 内存压力：与前向传播相比，反向传播面临更严重的内存压力，这仅仅是因为需要存储和计算的数据量更大。为缓解这一问题，我们像前向传播一样采用了激进的 SMEM/TMEM 重用策略，但也必须缩小分块尺寸。

• 流水线效率：与前向传播类似，残差块的长生命周期阻碍了流水线效率。对于内部残差，我们仍可使用 TMEM 和 TensorCore 累加来打破生命周期。但对外部残差，由于内存压力更高，我们无法在反向传播中进行寄存器预加载。相反，我们采用 TMA_REDUCE_ADD 操作，一旦残差就绪就直接从 SMEM 添加到 HBM。乍看之下这种做法似乎违背了 FlashNormAttention 减少内存 IO 的初衷，但实际上非常合理。通过融合操作，内存 IO 不再是内核的瓶颈，因此我们愿意以增加少量内存 IO 为代价，来缓解当前的瓶颈——计算和流水线停顿。

下图展示了反向传播融合后的性能提升。与前向传播使用完全相同的图表形状。

致谢

感谢 Tri Dao、Markus Hoehnerbach、Jay Shah、Ted Zadouri、Vijay Thakkar、Wentao Guo 在 Flash Attention 和 Quack 中的开源工作，这些工作为本文讨论的许多优化技术奠定了基础并提供了灵感。感谢 PyTorch 和 Triton 团队对 Helion 和 TLX 的开发维护，这些工具使本文提出的理念探索变得高效且富有成果。

参考文献

[1] Generalized Dot-Product Attention: Tackling Real-World Challenges in GPU Training Kernels.

https://pytorch.org/blog/generalized-dot-product-attention-tackling-real-world-challenges-in-gpu-training-kernels/

[2] Kunlun: Establishing Scaling Laws for Massive-Scale Recommendation Systems through Unified Architecture Design.

https://arxiv.org/abs/2602.10016

[3] Meta’s Generative Ads Model (GEM): The Central Brain Accelerating Ads Recommendation AI Innovation.

https://engineering.fb.com/2025/11/10/ml-applications/metas-generative-ads-model-gem-the-central-brain-accelerating-ads-recommendation-ai-innovation/

[4] Quack: Getting Memory-bound Kernels to Speed-of-Light.

https://github.com/Dao-AILab/quack/blob/main/media/2025-07-10-membound-sol.md

[5] Triton Tutorials: 03 Matrix Multiplication.

https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html

[6] FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision.

https://arxiv.org/pdf/2407.08608

[7] Look Ma, No Bubbles! Designing a Low-Latency Megakernel for Llama-1B.

https://hazyresearch.stanford.edu/blog/2025-05-27-no-bubbles

/post-content

/inner-wrap