【算子开发】算子融合与Softmax_LayerNorm实现

算子融合的价值

在讨论 Softmax 与 LayerNorm 的具体实现之前，有必要先回答一个更根本的问题：为什么不直接按数学公式逐行实现，而要引入"融合"（fusion）这一层设计？ 答案藏在一个简单的观察里：深度学习算子的计算强度（compute intensity）正在逐年下降，而模型中对内存带宽的渴望从未降低。换句话说，瓶颈早已不是算力，而是数据搬运。

访存开销：被低估的隐藏成本

一个算子的执行时间由两部分构成：计算时间 和访存时间 。现代 GPU（如 A100 或 H100）的浮点算力动辄数十 TFLOPS，而全局内存（HBM）带宽只有 1-3 TB/s。以 Softmax 为例，假设输入规模为 N=220N=2^{20}N=220（约 100 万元素，float32），一次完整的前向计算需要一个 max 归约、一个 sum-exp 归约和一次 归一化 逐元素操作。如果把这三个步骤写成三个独立的 kernel，每个 kernel 都要经历一次完整的 加载 → 计算 → 写回 流程：

阶段	访存路径	数据量（读 + 写）
kernel 1（求 max）	全局内存 → 寄存器 → 全局内存	2N×42N \times 42N×4 字节
kernel 2（求 sum-exp）	全局内存 → 寄存器 → 全局内存	2N×42N \times 42N×4 字节
kernel 3（归一化）	全局内存 → 寄存器 → 全局内存	2N×42N \times 42N×4 字节

总计 6N×46N \times 46N×4 字节 ≈ 24 MB 的全局内存流量。而如果通过算子融合，将三个步骤合并为单个 kernel，数据只需从全局内存读入一次、写回一次，流量降为 2N×42N \times 42N×4 字节 = 8 MB。访存量降低三分之二。

这个数字代表什么？以 A100 的 HBM 带宽 1.6 TB/s 计算：

• 非融合版本：24 MB/1.6 TB/s=15μs24 \text{ MB} / 1.6 \text{ TB/s} = 15 \mu s24 MB/1.6 TB/s=15μs 的纯访存时间
• 融合版本：8 MB/1.6 TB/s=5μs8 \text{ MB} / 1.6 \text{ TB/s} = 5 \mu s8 MB/1.6 TB/s=5μs 的纯访存时间

对于单个操作来说，10 微秒的差距看似微不足道，但一个深度模型中可能包含 200+ 次 Softmax/LayerNorm 调用（Transformer 中的每个 attention 层各一次）。累积下来就是毫秒级的额外开销------这等同于整个推理过程的数十倍延迟增长。

更关键的是计算与访存的比例失衡 ：Softmax 的计算复杂度为 O(N)O(N)O(N)，访问复杂度同为 O(N)O(N)O(N)，算术强度（arithmetic intensity）大约为 1 FLOP/byte。而 A100 的转折点 （两个操作耗时相等的强度阈值）约为 100 FLOP/byte。这意味着一半以上的时间花在等待数据到达，而非执行计算。任何能减少全局内存往返的改动，都会直接转化为可见的性能提升。

性能提升：减去的不只是访存------kernel 启动也减少了

除了全局内存流量的减少，融合还带来第二重收益：kernel 启动开销的消除。

每次内核启动（kernel launch）都有固有成本：CPU 需要向 GPU 的命令队列提交工作描述符，GPU 需要完成上下文切换。在 CUDA 中，这个开销大约为 3-10 微秒，看似不起眼，但有两个放大因素：

1. 小 kernel 场景：当单个 kernel 的执行本身只有几微秒时（例如小 batch 下的 Softmax），启动开销可能占据总耗时的 30%-50%。
2. 串联依赖 ：非融合版本中，kernel 1 必须在 kernel 2 开始前完全结束，因为后者依赖前者的输出。这种同步屏障（memory fence）强制 GPU 流水线排空，阻碍了不同 kernel 之间的指令级并行。

算子融合将多个 kernel 的启动合并为一次，同时消除了中间 kernel 间的同步屏障。结合前文访存减少带来的收益，在真实模型（如 BERT-base 的 12 层 Transformer）中，通过将 Softmax/LayerNorm/残差连接融合，端到端推理速度可提升 20%-40%（具体数字取决于实现质量与硬件平台）。

融合的原则：什么该融、什么不该融

任何内容都可以融合，但值得融合的算子符合一个简单标准：

上游产生一个大的中间张量；下游只做逐元素（elementwise）操作。

具体来说，识别可融合算子的三步法：

1. 张量形状分析 ：检查算子间传递的数据是否存在 shape 不变的依赖关系（即逐元素操作或多对一归约）。
2. 数据复用度：中间张量被使用的次数。只被使用一次 → 高融合潜力；被使用多次 → 需要权衡（如残差连接中被引用的分支）。
3. 内存占用：中间张量的大小是否远大于输入/输出。若中间张量巨大（例如 feature map），融合收益最大。

以下情形不适合融合：

• 输出形状发生剧烈变化（如池化层的 kernel 尺寸变化）
• 下游算子需要跨区块数据（如 batch normalization 的全局统计量）
• 算子间的数据依赖过于复杂，超出 kernel 的承载能力

实际工程中，成熟的编译器（如 TVM、XLA）会自动执行这类融合分析，但理解其基本原理，才能在手动优化原始 CUDA kernel 时做出正确的取舍。

预热：融合思路如何指导 Softmax 实现

回到 Softmax。将上述原则应用到 Softmax(xi)=exp⁡(xi)∑jexp⁡(xj)\text{Softmax}(x_i) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}Softmax(xi)=∑jexp(xj)exp(xi)，我们立刻发现三个操作（max 求取、sum-exp 求取、归一化）之间存在数据依赖链：归一化需要 sum-exp 的结果，sum-exp 需要 max 的结果。如果严格按公式分三步走，每个步骤之间就存在一整轮全局内存往返------这正是上一节分析的可优化之处。

后文将展示：在线 Softmax（online softmax） 算法如何在不做两次单独归约的前提下，仅用一次遍历同时完成 max 和 sum 的更新。它不仅将访存流量从三层降至一层，更重要的是，它为 LayerNorm 的融合实现提供了同样的方法论基础------当你能在一次 kernel 内完成归约与归一化时，算子融合就不再是一个编译器优化，而成为了算法设计的一部分。

Softmax数值稳定原理

上一节用"数据搬运"的视角解释了为何要融合算子，但融合之前必须先解决一个更隐蔽的数学问题：朴素的Softmax实现会在真实数据上直接崩溃。这个问题与GPU架构无关，却决定了我们必须采用特定的计算顺序。先看一个最直觉的Softmax实现：

yi=exi∑j=1Nexjy_i = \frac{e^{x_i}}{\sum_{j=1}^{N} e^{x_j}}yi=∑j=1Nexjexi

逐行翻译这句话，第一行 exp(x_i) 就可能让程序直接返回 inf 或 nan。假设某个输入特征 xi=1000x_i = 1000xi=1000，则 e1000≈10434e^{1000} \approx 10^{434}e1000≈10434------远远超过 float32 能表示的最大值（约 3.4×10383.4 \times 10^{38}3.4×1038）。这不是极端假设，在 transformer 模型中，经过多层残差连接后的 logits 数值分布往往横跨数百甚至上千的尺度，偶尔出现几百以上的数值是常态。

最大值偏移：一个平凡的数学恒等式

解决方案出人意料的简单。对任意常数 CCC，Softmax 具有平移不变性：

softmax(xi)=exi∑jexj=exi−C⋅eC∑jexj−C⋅eC=exi−C∑jexj−C\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}} = \frac{e^{x_i - C} \cdot e^{C}}{\sum_j e^{x_j - C} \cdot e^{C}} = \frac{e^{x_i - C}}{\sum_j e^{x_j - C}}softmax(xi)=∑jexjexi=∑jexj−C⋅eCexi−C⋅eC=∑jexj−Cexi−C

分子分母同时乘以 eCe^{C}eC，结果完全不变。因此，取 C=max⁡j(xj)C = \max_j(x_j)C=maxj(xj)，每个输入减去这个最大值后再做指数运算，所有指数项都被压缩到 (−∞,0](-\infty, 0](−∞,0] 区间内：

• 最大项变为 e0=1e^0 = 1e0=1，安全
• 其余项 e负数∈(0,1)e^{负数} \in (0, 1)e负数∈(0,1)，安全

这就是**最大值偏移（max-shift）**技巧。它不改变数学结果，但彻底消除了溢出的可能。代码如下：

import numpy as np

def softmax_stable(x):
    x_shifted = x - np.max(x)      # 平移：最大值为0，其余为负
    exp_x = np.exp(x_shifted)      # 全部落在(0, 1]区间，无溢出风险
    return exp_x / np.sum(exp_x)   # 归一化

这一版实现已经解决了数值溢出问题。但读者如果将其应用到 transformer 的 attention 层，会发现仍有隐患：np.max(x) 是对整个向量的一次完整归约，需要先遍历一遍数据。这意味着：读取全部数据求最大值 → 再读一遍计算 exp 和 sum → 再读一遍做除法。三轮完整的数据遍历，每轮都耗费全局内存带宽。

分阶段计算与两次遍历

为了更清晰地分析，将稳定的 Softmax 拆解为三个顺序执行的阶段：

1. 求最大值 ：m=max⁡j(xj)m = \max_j(x_j)m=maxj(xj)，一次归约遍历
2. 求指数和 ：s=∑jexj−ms = \sum_j e^{x_j - m}s=∑jexj−m，一次逐元素 + 归约遍历
3. 归一化 ：yi=exi−m/sy_i = e^{x_i - m} / syi=exi−m/s，一次逐元素遍历

每个阶段都必须等前一个阶段完成才能开始------阶段 2 需要阶段 1 的 mmm，阶段 3 需要阶段 2 的 sss。因此，即使不考虑 kernel 启动开销，全局内存也要被读取三次。当一个 attention 层要计算多个头的 Softmax 时，这个"三遍走"的成本会精确地按头数线性放大。

事实上，这正好呼应了前文所述的分阶段归约策略------第 1 节的讨论已指出，凡是需要全局信息的算子（如 max、sum），都必须将归约过程分为"局部计算 + 跨 block 合并"两个子阶段。Softmax 中的 max 和 sum-exp 正是两个典型的归约操作，它们的 sequential 依赖关系构成了性能优化的核心约束。

单遍算法：Online Softmax

能否把三次遍历压缩成一次？关键在于：能不能在不知道全局最大值的情况下，边读取数据边维护一个"局部正确"的 Softmax？

答案是肯定的，核心思想是迭代修正 。设 mkm_kmk 为前 kkk 个元素的最大值，dkd_kdk 为前 kkk 个元素修正后的 exp 和（即 ∑j=1kexj−mk\sum_{j=1}^{k} e^{x_j - m_k}∑j=1kexj−mk）。每读取一个新元素 xkx_kxk，分两种情况处理：

• 若 xk≤mk−1x_k \le m_{k-1}xk≤mk−1：现有最大值不变，只需累加 exk−mk−1e^{x_k - m_{k-1}}exk−mk−1
• 若 xk>mk−1x_k > m_{k-1}xk>mk−1：最大值更新为 xkx_kxk，此前所有元素都需要修正 ------它们之前除以的是 mk−1m_{k-1}mk−1，现在必须乘以 emk−1−xke^{m_{k-1} - x_k}emk−1−xk 才能换算成相对新最大值的指数

写成递推公式：

mk=max⁡(mk−1,xk)m_k = \max(m_{k-1}, x_k)mk=max(mk−1,xk)

dk=dk−1⋅emk−1−mk+exk−mkd_k = d_{k-1} \cdot e^{m_{k-1} - m_k} + e^{x_k - m_k}dk=dk−1⋅emk−1−mk+exk−mk

当整个向量扫描完毕，mNm_NmN 就是全局最大值，dNd_NdN 就是全局的 sum-exp。最后再用 第二遍 （注意不是第三遍）计算最终输出 yi=exi−mN/dNy_i = e^{x_i - m_N} / d_Nyi=exi−mN/dN。

这样总共只需两次遍历 ，而不再是三次。这就是著名的 online softmax 或 streaming softmax ，由 Milakov 与 Gimelshein 在 2018 年提出。其核心价值在于：在需要对大规模数据做分块处理（如 GPU 上数据无法一次性装入 SRAM）时，可以在流式读取的过程中持续维护 mkm_kmk 和 dkd_kdk，不必等待全局归约完成后才能开始计算------归约与逐元素计算在时间上重叠了。

与 log-sum-exp 的联系

另一条通用的数值稳定路径是 log-sum-exp （简称 LSE）。从 log 视角看，Softmax 的分母 log⁡(∑jexj)\log(\sum_j e^{x_j})log(∑jexj) 是一个在数学与统计学中频繁出现的量。直接计算它同样可能溢出，但利用与最大值偏移完全相同的恒等式：

log⁡∑jexj=m+log⁡∑jexj−m\log\sum_j e^{x_j} = m + \log\sum_j e^{x_j - m}logj∑exj=m+logj∑exj−m

其中 m=max⁡j(xj)m = \max_j(x_j)m=maxj(xj)。这等价于前一节"最大值为 0"的偏移，只不过在 log 空间表达。online softmax 中的 dNd_NdN 正是 eLSEe^{\text{LSE}}eLSE，两者互为表里------online softmax 维护的是线性空间的修正和，LSE 则直接在 log 空间工作。在后续实现中，最终归一化需要的是线性空间的 dNd_NdN，因此采用前一种形式更直接；而遇到需要 log 概率的场景（如某些损失函数），LSE 形式则更自然。

从空间到时间的延伸

将上面的讨论与第 1 节的访存分析结合，可以发现一个更本质的规律：数值稳定性的需求（减去最大值）恰好在天然也要求对全局数据的归约；而归约的顺序依赖（先求最大、再求和）决定了多遍遍历的必要性。online softmax 的价值不在于改变了这一依赖------它依然需要先有全局最大值才能做最终归一化------而在于将"求最大值"与"求 exp 和"合并为一遍遍历，让数据只在全局内存与计算单元之间往返两次。

至此，Softmax 的数学性质已经清晰。但这样的两步往返真的能压缩为一步吗？这就取决于能否将归一化过程拆分为局部输出 + 跨 block 修正 ------这正是下一节 基于分块归约的 Softmax 融合 要解决的问题：让 max 与 sum 不再需要单独的全局归约 kernel，而是嵌入到融合 kernel 的内部。

基于Reduction实现Softmax

有了数值稳定的基础，下一个问题自然浮现：稳定的Softmax应当如何落实到GPU上？ 第1节的访存分析已经给出了方向------用融合而非逐算子串联的方式消除中间结果的全局内存往返。但融合并非凭空而来，它是对"朴素多kernel方案"中每一步访存开销的逐一消除。因此，先拆解多kernel方案的成本结构，才能理解融合究竟省下了什么。

三个Kernel，两次全局内存往返

按照第2节推导出的稳定计算顺序，最直接的多kernel实现分三步：求 max、求 sum-exp、归一化。每一步对应一个独立的kernel，每个kernel都遵循加载 → 计算 → 写回 的模式。以批大小 B=256B=256B=256、序列长度 N=1024N=1024N=1024、隐藏维度 D=512D=512D=512 的输入张量为例（B×N×DB \times N \times DB×N×D 共约1.34亿个元素，float32 下约 512MB），整个流程的访存量可以精确量化：

Kernel	读取	写入	访存量
1. 求 max	512 MB（完整张量）	1 MB（每行一个最大值）	513 MB
2. 求 sum-exp	512 MB（完整张量）	1 MB（每行一个和）	513 MB
3. 归一化	512 MB（完整张量） + 1 MB（max）+ 1 MB（sum）	512 MB（结果）	1026 MB

总访存量约 2052 MB ------读取了三次整个张量，写回一次结果。而输入数据本身只有 512MB，这意味着数据在全局内存中被搬运了4次。

这还不是全部。每个kernel的启动本身也有开销：GPU kernel launch 的延迟通常在 3-10 微秒，虽然单次看似微不足道，但在推理场景中，Softmax 往往只是注意力层的一部分，而注意力层在解码阶段会被逐token调用数千次。三次启动叠加的延迟，叠加在访存延迟之上，共同构成了"多kernel方案"的真实成本基线。

中间量：融合的核心障碍

多kernel方案之所以效率低，根源在于它把 max 和 sum 这两个中间量 暴露给了全局内存。观察第二个kernel：它计算 ∑exi−m\sum e^{x_i - m}∑exi−m 时需要读取 mmm；第三个kernel归一化时需要同时读取 mmm 和 sss。这些中间量虽然不是完整张量（每行仅一个标量），但它们的存在强制了kernel之间的数据依赖，也强制了全局内存的往返。

融合的目标，就是让这些中间量只存在于片上（on-chip） 。GPU的共享内存（shared memory）和寄存器就是为此而生的------它们的带宽比全局内存高一个数量级，且无需跨越PCIe或HBM总线。理想情况下，max 和 sum 应该被计算出来、被后续步骤使用、然后被丢弃，全程不出芯片。

但这里有一个矛盾：融合要求一个kernel完成所有步骤，而标准CUDA的执行模型是------block内的线程通过共享内存协作，block之间无法直接通信 。Softmax的归约（reduction）发生在行内，而行通常远大于一个block能容纳的线程数。这就要求设计者决定：行的归约如何切分到block上，切分后中间量放在哪里。这正是第2节提到的分阶段归约在工程上的具体形态。

引用第10篇：归约的模式

此处的归约模式与第10篇文章中讨论的归约策略一脉相承。第10篇的核心结论是：归约的本质是将 NNN 个元素合并为更少的值（通常是1个），而高效的归约必须最大化片上数据复用、最小化跨线程通信。Softmax的三个步骤------max、sum-exp、归一化------恰好是归约的三种变体：

1. max ：使用 fmaxf 的归约，不涉及加法，但需要跨线程比较
2. sum-exp ：先计算 exi−me^{x_i - m}exi−m 再做加法归约。注意这里 mmm 是上一步的结果，形成了依赖链
3. 归一化 ：每个元素除以 sss，是 elementwise 操作而非归约，但需要广播 sss 到所有线程

第10篇中强调的"二分归约"（逐步合并部分和）在这里同样适用------无论是求最大值还是求和，都应当将行内元素划分给多个线程，每个线程先做局部归约，再通过共享内存在block内合并。但Softmax的特殊之处在于：第三步归一化不是归约，而是elementwise 。这决定了Softmax的kernel设计必须同时容纳两种范式------这恰恰是算子融合的典型特征：在一个kernel内，既有reduction，又有elementwise，且二者通过片上中间量连接。

从分步到融合：代价与收益

若将三步合并为一个kernel（暂且忽略行跨block的复杂性），访存将变为：

操作	访存量
读取完整张量	512 MB
写回结果	512 MB
总计	1024 MB

对比多kernel方案的 2052 MB，访存量恰好减半------这个结论与第1节"访存量降低三分之二"的表述略有差异，原因在于第1节的例子中假设了中间量也是完整张量（即未做数值稳定处理的朴素实现），而这里的三步方案已经将中间量压缩到了行级标量。但即便如此，融合仍然消除了一次完整的全局内存读和一次完整的全局内存写，从数据搬运的视角看，收益是质变级的。

更精确地说，融合的真实收益来自两个层面：

• 访存层面：从"读3次+写1次"变为"读1次+写1次"，全局内存流量减半
• 调度层面：三次kernel启动缩减为一次，GPU的空闲窗口（kernel间切换、内存延迟隐藏的空隙）被压缩

这背后隐含着一个深度原则：融合的本质是让中间量的生命周期（lifetime）终止于片上。中间量一旦被写入全局内存，就必然产生一次完整的读写往返；而如果能在片上直接消费掉，这部分开销就归零。

融合的边界与本文的承诺

以上分析假设了"融合一个kernel"的可行性，但现实中有一道必须直面的墙：当一行数据超过一个block的承载能力时，单kernel无法完成标准的三步归约 。此时面临两个选择------要么拆分为两个kernel（如FlashAttention的split-K方案），要么引入第2节提到的迭代修正策略，让归约不必等待完整max就能开始计算。后一条路指向online softmax，它正是下一节将要展开的核心算法。

在此之前，先建立一个清晰的心智模型：Softmax的融合实现 = 一个kernel内完成 max → sum-exp → 归一化，中间量max和sum全程驻留片上，全局内存只发生一次读和一次写。这个模型虽然朴素，但它是理解一切高级融合策略（包括online softmax）的基准点------后续的每一步优化，都能在这个基准上量化收益。

LayerNorm基础与实现

Softmax 的融合方案揭示了 reduction 与 elementwise 算子协同的基本范式：先归约出全局统计量，再按统计量逐元素修正 。LayerNorm 恰好是这一范式的另一经典实例------只不过它的统计量从单一的 max 变成了 mean（均值） 与 variance（方差） 两个，且归一化后还多了一个可学习的 affine（仿射） 变换。正因如此，LayerNorm 的融合实现不能照搬 Softmax 的三阶段结构，而要在归约策略上做出关键调整。

方差数值稳定：从原始定义到平移不变式

先看 LayerNorm 的数学定义。给定输入向量 x∈RNx \in \mathbb{R}^Nx∈RN，LayerNorm 的输出为：

yi=xi−μσ2+ϵ⋅γi+βiy_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma_i + \beta_iyi=σ2+ϵ xi−μ⋅γi+βi

其中 μ\muμ 是均值，σ2\sigma^2σ2 是方差，ϵ\epsilonϵ 是防止除零的小常数（通常取 10−510^{-5}10−5），γ\gammaγ 和 β\betaβ 是可学习的缩放与偏移参数。

直觉上，方差可以按原始定义直接计算：

σ2=1N∑i=1N(xi−μ)2(1)\sigma^2 = \frac{1}{N}\sum_{i=1}^{N}(x_i - \mu)^2 \tag{1}σ2=N1i=1∑N(xi−μ)2(1)

但这个公式在数值上有隐患。展开平方项：

σ2=1N∑i=1Nxi2−μ2(2)\sigma^2 = \frac{1}{N}\sum_{i=1}^{N}x_i^2 - \mu^2 \tag{2}σ2=N1i=1∑Nxi2−μ2(2)

公式 (2) 在数学上等价，但在浮点运算中却可能灾难性地失效。当 xix_ixi 的均值很大（比如 μ=106\mu = 10^6μ=106）而方差很小（比如 σ2=1\sigma^2 = 1σ2=1）时，公式 (2) 需要计算 1N∑xi2≈1012\frac{1}{N}\sum x_i^2 \approx 10^{12}N1∑xi2≈1012，然后减去 μ2≈1012\mu^2 \approx 10^{12}μ2≈1012，两个巨大的数相减得到一个小数 1。在 float32 精度（约 7 位有效数字）下，101210^{12}1012 量级的数字只能分辨出 10510^5105 的变化------结果就是方差被直接算成 0 ，归一化因分母为 ϵ\epsilonϵ 而失效。

这正是 Softmax 中 max-shift 问题的翻版：先平移、再计算，才能保住精度 。公式 (1) 虽然绕开了大数相减，但要求先算出 μ\muμ 才能逐元素计算 (xi−μ)2(x_i - \mu)^2(xi−μ)2，这引入了一步额外的归约。更糟糕的是，在融合 kernel 中这意味着要两遍扫描数据：第一遍求均值，第二遍求方差。对于必须驻留在寄存器中的片上数据来说，两遍扫描意味着要么重复加载全局内存，要么将整个向量缓存到共享内存。

共享内存：两遍扫描的片上缓存方案

两遍扫描的本质是：第一遍归约的中间结果，需要在第二遍扫描时再次访问 。如果数据只存在于全局内存，那么第二遍就必须重新从 HBM 读取，访存量翻倍。但如果把整个向量缓存在共享内存（shared memory） 中，两遍扫描就都只需访问片上存储。

共享内存是 GPU 上的一块可编程片上缓存，由线程块（block）内所有线程共享，访问延迟约为全局内存的 1/20 到 1/30，带宽则高出近一个数量级。对于 LayerNorm 而言，一个 block 处理一行（或一个 token）的数据，该行的大小通常不超过 1024 或 4096 个元素（对应模型隐藏维度），完全放得进共享内存。

以下代码展示了一个融合的 LayerNorm kernel：单个 kernel 内完成均值/方差归约、归一化与 affine 变换，中间量（均值、方差、缩放后的输入）只存在于片上。

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# 使用 CUDA 实现融合 LayerNorm
cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>

// 每个 block 处理一行数据，BLOCK_SIZE 为行长度（需为 32 的倍数）
template<int BLOCK_SIZE>
__global__ void layernorm_fused_kernel(
    const float* __restrict__ input,   // [num_rows, BLOCK_SIZE]
    const float* __restrict__ gamma,   // [BLOCK_SIZE]
    const float* __restrict__ beta,    // [BLOCK_SIZE]
    float* __restrict__ output,        // [num_rows, BLOCK_SIZE]
    float eps,
    int num_rows) {

    // 1. 每个线程加载一个元素到寄存器
    int row = blockIdx.x;
    int col = threadIdx.x;
    int idx = row * BLOCK_SIZE + col;

    float x = input[idx];

    // 2. 两遍扫描的中间数据缓存在共享内存
    __shared__ float shared_x[BLOCK_SIZE];

    // 第一遍归约：计算均值（block 内所有线程协作）
    __shared__ float sum_buf;
    __shared__ float mean_val;

    shared_x[col] = x;
    __syncthreads();  // 确保 shared_x 全部写入

    // block 内归约求和（树形归约）
    for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
        __syncthreads();
        if (col < stride) {
            shared_x[col] += shared_x[col + stride];
        }
    }

    // block 归约后，线程 0 持有总和
    if (col == 0) {
        sum_buf = shared_x[0];
        mean_val = sum_buf / BLOCK_SIZE;
    }
    __syncthreads();  // 广播 mean

    // 3. 第二遍扫描：利用共享内存中的原始 x 计算方差
    // （无需重新从全局内存加载 x）
    float diff = shared_x[col] - mean_val;
    float sq_diff = diff * diff;

    __syncthreads();  // 重置 shared_x 以复用做方差求和
    shared_x[col] = sq_diff;
    __syncthreads();

    for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
        __syncthreads();
        if (col < stride) {
            shared_x[col] += shared_x[col + stride];
        }
    }

    __shared__ float var_val;
    if (col == 0) {
        var_val = shared_x[0] / BLOCK_SIZE;
    }
    __syncthreads();  // 广播方差

    // 4. 归一化 + affine 变换
    float inv_std = rsqrtf(var_val + eps);
    float norm_x = (x - mean_val) * inv_std;
    output[idx] = norm_x * gamma[col] + beta[col];
}

torch::Tensor layernorm_fused(torch::Tensor input,
                              torch::Tensor gamma,
                              torch::Tensor beta,
                              double eps) {
    int num_rows = input.size(0);
    int block_size = input.size(1);

    auto output = torch::empty_like(input);
    const float* in_ptr = input.data_ptr<float>();
    const float* g_ptr = gamma.data_ptr<float>();
    const float* b_ptr = beta.data_ptr<float>();
    float* out_ptr = output.data_ptr<float>();

    // 只支持 block_size = 1024 的模板实例化
    layernorm_fused_kernel<1024><<<num_rows, 1024>>>(
        in_ptr, g_ptr, b_ptr, out_ptr, (float)eps, num_rows);

    return output;
}
'''

cpp_source = r'''
torch::Tensor layernorm_fused(torch::Tensor input,
                              torch::Tensor gamma,
                              torch::Tensor beta,
                              double eps);
'''

# 编译并测试
layernorm_fused = load_inline(
    name="layernorm_fused",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["layernorm_fused"],
    verbose=False
)

# 验证正确性
torch.manual_seed(42)
x = torch.randn(4, 1024, dtype=torch.float32, device="cuda")
gamma = torch.randn(1024, dtype=torch.float32, device="cuda")
beta = torch.randn(1024, dtype=torch.float32, device="cuda")

# 参考实现
ref = torch.nn.functional.layer_norm(x, (1024,), gamma, beta, 1e-5)

# 融合实现
fused = layernorm_fused(x, gamma, beta, 1e-5)

# 验证数值一致性
print("Max absolute error:", (ref - fused).abs().max().item())
# 输出: Max absolute error: ~1e-6 量级（float32 精度可预期）

这段 kernel 的关键设计有三个：

1. 共享内存的复用 ：第一次 shared_x 缓存原始输入 x，供第二遍计算方差时直接复用，避免了从全局内存重新加载 x。随后同一块共享内存被重置，转而做方差求和。整个 kernel 对全局内存只做了一次读（原始输入）和一次写（输出）。
2. 两遍归约：第一遍求均值，第二遍求方差。虽然多了一遍归约，但由于是片上操作，开销远小于全局内存往返。
3. 寄存器驻留 ：x 从共享内存读出后保存在寄存器中，最终归一化时无需再次访问共享内存。

这个实现相比三 kernel 分离方案（全局归约均值 → 全局归约方差 → 逐元素归一化），访存量从 3 次读 + 1 次写降为 1 次读 + 1 次写，减少了一半的全局内存流量。与 Softmax 融合节省三分之二访存量相比有所逊色，原因在于 LayerNorm 需要两个统计量、天然涉及两遍归约------但片上缓存让第二遍归约近乎免费。

从 LayerNorm 到训练反向

以上实现覆盖了 LayerNorm 的前向（inference）路径。在训练场景中，反向传播还需要计算均值和方差对输入的梯度，这意味着 反向 kernel 同样需要 μ\muμ 和 σ2\sigma^2σ2 。这正是融合的另一个好处：前向 kernel 计算出的统计量可以保留在片上或寄存器中传给反向 kernel，而分离实现则必须将它们写回全局内存再重新读取。这一深化留待探讨反向融合时再展开，而下一节将讨论另一个重要的归一化算子------BatchNorm 与 LayerNorm 在融合策略上的本质差异，以及这种差异如何影响 kernel 设计。

融合实现：Online Softmax与Block内归一化

第 3 节的多 kernel 方案暴露了一个结构性矛盾：每个 kernel 都要完整地读一遍数据、写一遍中间结果 。即便每个 kernel 内部的计算都极其高效，全局内存上的往返依然无法避免。Softmax 对每一行数据需要先求出 max，再求 sum-exp，最后归一化------所以三个 kernel 就是三次全量往返。能不能把这三个阶段收进同一个 kernel？答案藏在两个关键操作中：warp shuffle 与 online softmax。

从"全局中间数组"到"片上中间量"

第 1 节分析过：多 kernel 方案之所以慢，核心在于中间量被写回全局内存。以一张 [B, N] 的输入为例，多 kernel 方案在全局内存中至少需要为 max 和 sum-exp 各分配一份 [B] 的中间数组。假设 B=65536B = 65536B=65536，N = 4096，则在单精度下，仅中间数组就占用 65536×4 bytes×2=512 KB65536 \times 4 \text{ bytes} \times 2 = 512\text{ KB}65536×4 bytes×2=512 KB------这还只是存储开销，真正的开销在于每个中间量的写入与再次读取各产生一次全局内存事务。

融合方案的目标只有一个：让所有中间量只存在于片上（on-chip） 。具体而言，将 max、sum-exp 以及最终的归一化结果全部保留在寄存器（register）与共享内存（shared memory）中，直到最后一个元素计算完毕，才把结果写回全局内存。数据从全局内存中读取一次，产生一个最终输出，全局内存流量从"读三遍、写三遍"压缩到"读一遍、写一遍"。

但这引出一个新的问题：第 3 节的朴素归约方案中，max 和 sum-exp 是两个独立的数据扫描过程。如果把它们合并为一次扫描，sum-exp 需要一个全局的 max 才能开始累加 ------而扫描到一半时，全局的 max 尚未出现，怎么办？

Online Softmax：边扫描边修正

这正是 online softmax 的用武之地。它的核心思想是：不等待完整的 max 再计算 sum-exp，而是在扫描过程中维护一个"运行中的最大值"，并随时修正已经累加的部分和。

算法维护两个状态变量：当前最大值 mmm，以及未归一化的指数和 sss。每读取一个新元素 xix_ixi，执行以下三步：

m′=max⁡(m,xi) m' = \max(m, x_i) m′=max(m,xi)

s′=s⋅em−m′+exi−m′ s' = s \cdot e^{m - m'} + e^{x_i - m'} s′=s⋅em−m′+exi−m′

m←m′,s←s′ m \leftarrow m', \quad s \leftarrow s' m←m′,s←s′

第二步是精妙所在：当新的 xix_ixi 比当前最大值 mmm 更大时（即 m′>mm' > mm′>m），此前所有累加进 sss 的项都以旧的 mmm 为基准做了偏移，现在基准变成了 m′m'm′，因此需要乘以一个修正系数 em−m′e^{m - m'}em−m′ 来统一到新的基准 ；而当前元素 xix_ixi 直接以 m′m'm′ 为偏移量加入。扫描结束后，sss 即为以 max⁡\maxmax 为偏移的指数和，只需一步 yi=exi−m/sy_i = e^{x_i - m}/syi=exi−m/s 即可归一化。

这个"边扫描边修正"的策略，把第 3 节的两遍扫描压缩为一遍。以一批 B=65536B=65536B=65536 行、每行 N=4096N=4096N=4096 的输入为例，多 kernel 方案需要两次全量读取（求 max、求 sum-exp）+ 一次全量读取（归一化），共三次；online softmax 只需一次读取 + 一次写回。全局内存流量直接从 6BN6BN6BN 字节降至 2BN2BN2BN 字节------对 512 MB 的数据量，就是 2052 MB 与 684 MB 的差距。

Warp Shuffle：线程间的片上归约

解决了扫描次数的矛盾，下一个问题是：如何让一个 warp（32 个线程）高效地协同完成行内归约 ？两个线程若想交换数据，最直接的方式是经过共享内存，但共享内存存在访存延迟与 bank 冲突的风险。GPU 提供了一条更轻量的路径：warp shuffle 指令（__shfl_xor_sync 等），允许一个 warp 内任意两个线程直接交换寄存器中的值，不经过任何中间存储。

对于长度为 NNN 的一行数据，假设 32 个线程各持有 N/32N/32N/32 个连续元素。每个线程先对本地元素做部分归约（求局部的 max 和局部的 sum-exp），得到一个局部值。接下来的关键是对 32 个局部值做归约------这一步完全可以用 warp shuffle 完成：

// 假设 lane_id 为线程在 warp 中的编号，val 为线程持有的局部值
for (int offset = 16; offset > 0; offset >>= 1) {
    // 与相距 offset 个 lane 的线程交换数据，
    // 取两者中的较大者（求 max 的归约）
    val = fmaxf(val, __shfl_xor_sync(0xffffffff, val, offset));
}
// 循环结束后，warp 内所有线程的 val 均等于全局最大值

这段代码的精髓在于 __shfl_xor_sync：它将当前线程寄存器中的 val 与 lane 编号异或 offset 后的线程的 val 交换。经过 5 轮迭代（offset = 16, 8, 4, 2, 1），一个 warp 内的 32 个值以 O(log 32) = 5 步完成归约。每一步只依赖上一步的结果，无锁、无同步、无共享内存访问。

以 A100 为例，一次 __shfl_xor_sync 的延迟约为几纳秒，而一次共享内存访问（含同步开销）往往需要数十纳秒。当 N 足够大、行数足够多时，warp shuffle 的累计收益是数量级的访存节省 。更重要的是，shuffle 使中间值始终停留在寄存器中，天然契合 online softmax 的"片上修正"策略------m 与 s 在归约过程中不断被 shuffle 更新，却从未离开过寄存器。

融合的完整蓝图：一张数据流图

将上述三个要点组合起来，单 kernel 的完整数据流如下：
渲染错误: Mermaid 渲染失败: Parse error on line 6: ...个线程执行归一化
y = exp(x - max) / sum-exp] -----------------------^ Expecting 'SQE', 'DOUBLECIRCLEEND', 'PE', '-)', 'STADIUMEND', 'SUBROUTINEEND', 'PIPE', 'CYLINDEREND', 'DIAMOND_STOP', 'TAGEND', 'TRAPEND', 'INVTRAPEND', 'UNICODE_TEXT', 'TEXT', 'TAGSTART', got 'PS'

每个 block 处理一行（或多行），行内归约由 warp shuffle 完成；若一行跨多个 warp / block，则在 warp 之间再用共享内存做一次合并。无论哪种情况，中间量都不离开 SM（流式多处理器） ------融合的收益由此兑现：全局内存往返从三次降为一次，中间的 max 与 sum-exp 数组被彻底消除，kernel 启动次数从三个降为一个。

---

至此，Softmax 的融合版本已经从"多 kernel 串联"进化为"单 kernel 一次扫描"。而这一数据流------归约出统计量、广播、逐元素修正------正是第 4 节 LayerNorm 融合实现将要复用的骨架。区别在于，LayerNorm 的统计量是 mean 与 variance，且它多了一个可学习的 affine 变换。下一节，我们将顺着这条融合路径，回答 LayerNorm 的方差应该如何高效且数值稳定地在 block 内归约，同时处理好那个"两遍扫描"的老问题。

案例：合并多个 Activations

前几节的分析始终围绕一个核心命题：减少数据搬运比减少计算量更能带来真实的性能提升 。Online Softmax 与 LayerNorm 的融合实现已经给出了范式------将多个依赖的算子合并为单个 kernel，让中间数据停留在片上。但这两个案例都属于"一类输入、逐行归约"的结构。实践中更常见、也更能体现融合通用性的场景，是 MatMul + Bias + ReLU 这条在各类神经网络中反复出现的数据通路。下面用一个具体案例展示融合如何落地，并揭示它与前面 Softmax/LayerNorm 案例在数据流层面的差异。

数据流共享：融合的根本依据

先看未融合的朴素实现中，数据是如何流动的。设输入矩阵为 X∈RM×K\mathbf{X} \in \mathbb{R}^{M \times K}X∈RM×K，权重矩阵为 W∈RK×N\mathbf{W} \in \mathbb{R}^{K \times N}W∈RK×N，偏置向量为 b∈RN\mathbf{b} \in \mathbb{R}^{N}b∈RN。三个算子的数学定义依次为：

Z=XW,Z′=Z+b,Y=max⁡(Z′,0) \mathbf{Z} = \mathbf{X} \mathbf{W}, \quad \mathbf{Z}' = \mathbf{Z} + \mathbf{b}, \quad \mathbf{Y} = \max(\mathbf{Z}', 0) Z=XW,Z′=Z+b,Y=max(Z′,0)

逐算子执行时，中间结果 Z\mathbf{Z}Z 和 Z′\mathbf{Z}'Z′ 都必须写回全局内存 ，因为下一个 kernel 无法直接读取上一个 kernel 的片上数据。这就是前文反复强调的"加载 → 计算 → 写回"循环------每个中间量在全局内存中被写入一次、读取一次，合计产生 2×M×N2 \times M \times N2×M×N 个元素的两次额外访存事务。

数据流共享 指的是：这三个算子的输入输出之间存在严格的依赖链------Z\mathbf{Z}Z 的每个元素只被 Bias 算子使用一次，Bias 的输出 Z′\mathbf{Z}'Z′ 也只被 ReLU 使用一次。没有任何中间结果被多个消费者共享 。这意味着它们不需要物化到全局内存中，完全可以"接力"在片上完成。换言之，融合的本质是将数据流从"全局内存中转"改为"寄存器/共享内存直传"。

对比前文的 Softmax 案例，这里的差异值得注意：Softmax 需要两次归约（先求 max 再求 sum-exp），中间量跨越了 kernel 边界；而 MatMul + Bias + ReLU 的中间量 Z\mathbf{Z}Z 只是逐元素传递 ，不涉及归约。因此它的融合难度更低，但数据流共享的原理完全一致------中间数据什么时候必须离开片上？只有当它需要被全局归约、或需要被另一个 block 消费时，才被迫落到全局内存。

融合思路：从 loop 嵌套到单 kernel

融合的思路并不神秘：在一个 kernel 内，用循环展开替代 kernel 边界 。具体而言，每个线程负责计算输出矩阵 Y\mathbf{Y}Y 中的一个或多个元素。对于每个输出元素 yijy_{ij}yij，它需要的输入是 X\mathbf{X}X 的第 iii 行与 W\mathbf{W}W 的第 jjj 列------这些数据通过共享内存或寄存器缓存后，线程在同一个 kernel 内依次完成乘加、加偏置、ReLU 三个步骤：

# 伪代码：每个线程计算一个输出元素 y[m][n]
def fused_matmul_bias_relu_kernel(X, W, b, Y, M, N, K):
    m, n = get_thread_coordinates()  # 当前线程负责的输出坐标
    
    # Step 1: 在寄存器中累加 MatMul 结果（不写回全局内存）
    acc = 0.0
    for k in range(K):
        acc += X[m][k] * W[k][n]
    
    # Step 2: 加偏置 ------ 直接修改寄存器值
    acc += b[n]
    
    # Step 3: ReLU ------ 同样在寄存器中完成
    acc = max(acc, 0.0)
    
    # 仅在最终结果产生后写回一次
    Y[m][n] = acc

这段伪代码的访存模式与三 kernel 版本完全不同：

方案	全局内存写次数	全局内存读次数	说明
三 kernel	2 次（Z、Z'）	3 次（X、W/Z/Z'各一次）	中间结果各往返一次
单 kernel	1 次（Y）	2 次（X、W）	中间量仅存于寄存器

从表中的对比可以直接读出融合的收益：全局内存写次数从 2 次降为 1 次，读次数从 3 次降为 2 次 。对于 M=N=K=1024M=N=K=1024M=N=K=1024 的典型规模，三 kernel 方案在全局内存上搬运约 2×3×10242=6.32 \times 3 \times 1024^2 = 6.32×3×10242=6.3 MB 数据（含中间结果），单 kernel 方案仅搬运 2×2×10242=4.22 \times 2 \times 1024^2 = 4.22×2×10242=4.2 MB------访存量降低 33%。若考虑中间结果 Z 和 Z' 占用的额外分配开销，实际收益还会更高。

更关键的是，这个融合模式可以自由扩展 。只需在 ReLU 之后继续加激活函数（如 GELU）、加残差连接（acc += residual[m][n]）、甚至加 LayerNorm------每个扩展都只是在寄存器链上追加一步操作，不增加任何全局内存往返。这正是"算子融合是架构设计而非技巧"的含义：它改变的是数据流的基本模式，而非某个算子的实现细节。

回到本文的叙事主线：Softmax 的融合解决的是归约跨 block 的通信问题 ，LayerNorm 的融合解决的是两遍扫描的片上缓存问题 ，而本节 MatMul + Bias + ReLU 的融合解决的则是逐元素流水线的数据直传问题 。三者共同指向一个统一的结论：任何中间数据，只要不存在跨 block 的依赖，就应当留在片上。下一节将把这一原则推广到更系统的视野------对比 BatchNorm 与 LayerNorm 在融合策略上的差异，看看当统计量的计算范围从"一个样本的一个特征"扩展到"一个 batch 的全部样本"时，融合的设计空间又会发生怎样的变化。