🏛️ 第23讲：为什么Attention必须先算Max？——从 Naive Reduce 到 Fused Online Softmax 深度解构

主讲人：👓 Ringi（大厂 AI Infrastructure 工程师）
所属模块：Module 02: CUDA 编程与高性能算子优化
篇章范式：⚡ CUDA 编程与高性能算子优化篇（Kernel & Operator Optimization Paradigm）
核心导读：在深度学习全栈工程中，Softmax 往往被看作一个平平无奇的激活函数。但正是这个简单的 $\frac{e^{x_i}}{\sum e^x}$ ，却卡死了无数大模型推理与训练的吞吐上限。初学者写 Softmax，第一步就被 IEEE-754 浮点溢出教做人；加上数值稳定保护后，又不得不为了求 Max、求 Sum、做归一化而在高昂的 HBM 显存与计算核心之间来回往返扫描 3 次（3-Pass）！本讲我们将从并行计算经典基石——树形规约（Reduce）的七级性能跃迁讲起，彻底拆解 Warp Shuffle 跨线程寄存器直通原语；进而深入剖析 2018 年 NVIDIA 提出的 Online Softmax 在线动态修正算法，手算代数结合律证明；最后推开现代大模型加速圣殿的大门，揭示它究竟是如何演化为颠覆时代的 FlashAttention 底座核心原语的。

========================================================================================================================
                                Ringi 3D 架构工坊 · Reduce 与 Fused Online Softmax 核心全景
========================================================================================================================
[ 传统 3-Pass Safe Softmax (高昂 HBM 往返深渊) ]
   Global Memory (HBM) ───[Pass 1 读: 4N]───► 求全局 Max (m) ───[写回: 4B]───► Global Memory
   Global Memory (HBM) ───[Pass 2 读: 4N]───► 算 Exp 求和 (d) ──[写回: 4B]───► Global Memory
   Global Memory (HBM) ───[Pass 3 读: 4N]───► 算 Exp(x-m)/d ────[写回: 4N]───► Global Memory
   ★ 显存总流量: 3次读 + 1次写 = 16N 字节 (对于 4096 维张量，带宽利用率惨不忍睹，ALU 长期饥饿等待)
------------------------------------------------------------------------------------------------------------------------
                                                ▼  算法与硬件架构重构演进
------------------------------------------------------------------------------------------------------------------------
[ 现代 1-Pass Fused Online Softmax (片上寄存器/共享内存极速闭环) ]
   Global Memory (HBM) ───[唯一 1 次读: 4N]───► [ SM 片上寄存器缓存 Reg Cache ]
                                                      │
                                                      ├──► 动态维护局部 (m_local, d_local)
                                                      │    ★ 修正因子递推: d = d * exp(m_old - m_new) + exp(x - m_new)
                                                      │
                                                      ├──► Warp 级跨 Lane 寄存器规约: __shfl_down_sync (5步折叠)
                                                      │
                                                      ├──► Block 级共享内存快速交换 (Shared Memory)
                                                      │
                                                      └──► 寄存器原地归一化直接写回 ───► Global Memory [唯一 1 次写: 4N]
   ★ 显存总流量: 1次读 + 1次写 = 8N 字节 (显存访问流量暴降 50%，算子吞吐直接翻倍！)
------------------------------------------------------------------------------------------------------------------------
                                                ▼  赋能大模型注意力机制
------------------------------------------------------------------------------------------------------------------------
[ 终极形态: FlashAttention 核心计算原语 ]
   SRAM 分块维护局部 (m_i, d_i) ──► 动态缩放输出累加张量: O_new = diag(exp(m_old - m_new)) * O_old + P_ij * V_j
   ★ 彻底终结 O(N^2) 显存存储瓶颈，实现无限长上下文注意力极速流式计算！
========================================================================================================================

📑 目录导航

0. Ringi 开场：生产真实现场与痛点冲突
1. 规约（Reduction）体系结构：从交错寻址到 Warp Shuffle 的七级跳跃
2. 数值稳定性第一性原理：为什么 Softmax 必须做“Safe”保护？
3. Online Softmax 算法推导：如何在一遍扫描中同时求 Max 和 Sum？
4. 工业级工程实现：从 2-Pass Online Softmax 到 1-Pass Fused Softmax
5. FlashAttention 核心前置：Online Softmax 是如何成就大模型注意力革新的？
6. 动手实战与代码实验室（Minimal Runnable Code）
7. Ringi 避坑指南与生产黄金准则
- 7.1 避坑表格：7 大常见小白错误理解 vs 大厂 AI Infra 正确认知
- 7.2 生产性能工程黄金 Checklist
8. Ringi 5 点核心速记口诀、自我检验清单与课后深度思考题
9. 📚 参考资料与核心源码/经典论文指引
附录：Appendix A — 大厂硬核高频面试题与白板推导（Interview Drill）
🎨 配图工坊生图 Prompt 暂存区

0. Ringi 开场：生产真实现场与痛点冲突

0.1 真实工程矛盾：为什么长文本一开，Softmax 成了显存吞噬兽？

在 Transformer 架构中，自注意力机制（Self-Attention）的核心计算公式天下皆知：

\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V

当序列长度（Sequence Length

N

）从 2K、8K 扩展到长上下文的 32K、128K 乃至 1M 时，一个极其残酷的算力矛盾暴露无遗：计算矩阵乘

S = QK^T

是典型的 Compute-Bound（计算密集型） 任务，在 NVIDIA A100/H100 的 Tensor Core 上可以跑出 300~1000 TFLOPS 的恐怖峰值算力；然而紧接着的 Softmax 算子，却是一个典型的 Memory-Bound（访存密集型） 任务！让我们算一笔真实的显存账本：对于一个包含

M

行、每行

N

个元素的中间注意力得分矩阵

S

：

传统的 Safe Softmax 算法为了防止指数爆炸，必须先扫描一遍数据求出每行的最大值 $m = \max(x)$ ；
接着必须再次从 HBM 全局显存把这批数据读进 SM 核心，计算 $\sum e^{x_i - m}$ ，再把分母和写回 HBM；
最后，第三次从 HBM 读出输入数据，执行 $y_i = \frac{e^{x_i - m}}{d}$ ，并将结果写入全局显存供后续与 $V$ 做矩阵乘。

这意味着对同一个张量，全局显存来回读了 3 遍、写了 1 遍！ 当

N = 32768

（32K 长文本）时，单头单层的注意力矩阵大小为

32768 \times 32768 \times 2 \text{B (FP16)} = \mathbf{2 \text{ GB}}

！仅仅这一个中间 Softmax 算子，在 GPU 显存总线上就要搬运整整

2 \text{ GB} \times (3 + 1) = \mathbf{8 \text{ GB}}

的数据！在多头、多层并发下，GPU 哪怕有 2 TB/s 的 HBM 带宽，也瞬间被这些毫无意义的数据搬运彻底榨干，Tensor Core 被迫长期处于“断粮发呆”的严重饥饿状态。

Ringi 导师解构：从 3-Pass 显存往返到 1-Pass 寄存器流水的显存流量对比

0.2 线上事故复盘：某多模态大模型 FP16 数值溢出（NaN 灾难）与三遍扫描往返墙

2024 年春，某大厂在将一款视觉-语言多模态大模型（VLM）从 FP32 训练迁移至 FP16 生产推理集群时，线上偶发性出现整句输出全部变为乱码符号甚至直接崩溃返回 HTTP 500 的事故。监控显示，模型推理内部某层 Softmax 的输出张量突然变成了全 NaN（Not a Number）。资深 Infra 工程师下场排查后发现，某位算法开发同学在手写底层融合算子时，认为“既然模型在 FP16 下权重数值都挺小，何必费劲先算一遍最大值？直接算

\exp(x)

性能还能快 30%”！他写出了如下看似极简的代码：

// ❌ 致命隐患：未做 Safe 保护的 Naive Softmax
float sum = 0.0f;
for(int i = 0; i < N; ++i) {
    float ex = __expf(x[i]); // 当 x[i] 稍微偏大时瞬间爆炸！
    sum += ex;
}

这一行代码成了致命的地雷！ 在 IEEE-754 半精度浮点数（FP16）规范中，最大能表示的正实数仅为

65504

。而在反向推导中：

\ln(65504) \approx 11.0898

这意味着：只要注意力得分矩阵中有任何一个位置的数值超过了 11.1， $\exp(x_i)$ 就会立刻溢出为浮点正无穷（+Inf）！ 接下来，任何数值除以 +Inf，或者 +Inf 减去 +Inf，整个张量就会像瘟疫一样瞬间被染成全是 NaN！为了解决这个问题，算法团队不得不退回“先求最大值”的 Safe Softmax。然而，换上 Safe Softmax 后，原本的 P99 延迟直接从 22ms 暴涨至 58ms！因为他们使用了三个独立的 CUDA Kernel 分别算 Max、算 Sum、算 Normalize，触发了臭名昭著的 “三遍全局显存往返墙（3-Pass Round-Trip Wall）”。直到 Infra 团队重构接入了 Online Softmax 原生融合算子，将 3 遍读取骤降为 1 遍片上直通闭环，才在彻底保证 FP16 绝对数值稳定的前提下，将延迟重新压回了 14ms，不仅排除了炸数值故障，还带来了 4.1 倍的端到端吞吐提升！

0.3 AI Infra 规约与 Softmax 各演进阶段速查表

演进阶段 (Paradigm)	核心算法与硬件机制	全局显存访问次数 (Memory Passes)	显存总搬运流量 (FP32, 大小 $N$ )	Warp 分支发散 (Divergence)	共享内存 Bank Conflict	A100 实测有效带宽利用率	典型工业应用场景
Stage 1: Naive Interleaved Reduce	树形折半，取模 `% (2*s)` 寻址	1 读 1 写 (两阶段多 Block)	$\sim 4N$ 字节	严重 (活跃线程每轮减半且分散)	严重 (多线程同一 Bank)	$\sim 12\%$	早期教科书教学代码
Stage 2: Sequential Addressing Reduce	两端对折折叠，`tid < s` 连续寻址	1 读 1 写 (结合共享内存)	$\sim 4N$ 字节	无发散 (前 $N/32$ 轮完全对齐)	0 冲突 (完美消除)	$\sim 54\%$	通用单精度向量求和
Stage 3: Warp Shuffle Reduce	寄存器直通 `__shfl_down_sync`	1 读 1 写 (寄存器级零 SMEM)	$\sim 4N$ 字节	0 发散 (硬件 SIMT 锁步)	无 SMEM 开销	$\sim 88\%$	高性能 LayerNorm / RMSNorm
Stage 4: 传统 Safe Softmax (3-Pass)	分离 Kernel：Max $\rightarrow$ Sum $\rightarrow$ Norm	3 读 1 写 (全局来回倒腾)	$16N$ 字节 (极大浪费)	取决于具体 Reduce 实现	规约阶段可能存在	$\sim 35\%$	早期未融合 PyTorch 原生算子
Stage 5: Online Softmax (2-Pass)	动态维护 $(m, d)$ ，在线乘法修正	2 读 1 写 (合并 Max+Sum)	$12N$ 字节 (减少 25% 流量)	结合 Warp Shuffle 极低	0 冲突	$\sim 70\%$	现代高性能 Softmax 基础版
Stage 6: Fused 1-Pass Online Softmax	寄存器缓存输入 + Online 递推 + 归一化	1 读 1 写 (物理极限！)	$8N$ 字节 (流量减少整整 50%)	0 发散	0 冲突	$\sim 85\% - 92\%$	生产级长文本 Attention Softmax
Stage 7: FlashAttention 融合注意力	Online Softmax + GEMM 分块流式更新	0 显存落盘 (中间得分全在片上)	仅读 $Q,K,V$ ，写 $O$	极低	依靠 Swizzle / TMA 规避	硬件算力打满	现代大模型训练与推理标准引擎

1. 规约（Reduction）体系结构：从交错寻址到 Warp Shuffle 的七级跳跃

为了在深入具体细节前建立完整的物理心智模型，下方给出了规约树演进、Safe Softmax 显存往返墙、Online Softmax 动态修正代数递推与 1-Pass 算子融合的工业级全景架构拓扑：

并行规约演进、Safe Softmax 存储墙与 Fused Online Softmax 全景图

1.1 规约操作在 AI 算子中的统治地位

在并行计算领域，规约（Reduction） 的定义是将一个数组中的

N

个元素，通过一个满足结合律的二元操作符

\oplus

（如加法、乘法、求最大值、求最小值），逐步聚合为一个单一标量标量的过程：

y = x_0 \oplus x_1 \oplus x_2 \oplus \dots \oplus x_{N-1}

在大模型底层体系结构中，规约是出现频次仅次于矩阵乘（GEMM）的第二大类算子：

Softmax 算子：需要求行最大值 $\max(x_i)$ 和指数和 $\sum e^{x_i}$ ；
LayerNorm / RMSNorm 算子：需要求特征维度的均值 $\mu = \frac{1}{d}\sum x_i$ 与方差 $\sigma^2 = \frac{1}{d}\sum x_i^2$ ；
Cross-Entropy Loss 算子：需要在全局 Batch 上做损失求和；
Gradient AllReduce：在多卡分布式训练中，不同 GPU 之间的核心通信原语就是跨节点的梯度向量规约。

可以说：规约写得好不好，直接决定了整个深度学习系统在 Memory-Bound 算子上的吞吐下限。

1.2 Mark Harris 经典七级优化脉络全景解构

2007 年，NVIDIA 体系结构科学家 Mark Harris 发表了传世之作《Optimizing Parallel Reduction in CUDA》，系统性地展示了如何将一个简陋的规约算子逐步提速数十倍。这套七级进化体系，至今仍是所有 GPU 工程师必经的技术洗礼。

====================================================================================================
                        Mark Harris 规约树 (Tree Reduction) 演进拓扑
====================================================================================================
[ 传统交错寻址 (Interleaved) - 活跃线程分散，严重分支发散与 Bank 冲突 ]
Thread ID:   0   1   2   3   4   5   6   7   8   9  10  11  12  13  14  15 ...
Step 1:      └───┬───┘   └───┬───┘   └───┬───┘   └───┬───┘   └───┬───┘   └───┬───┘   (stride=1, 只有偶数线程活跃)
Step 2:      └───────┬───────┘       └───────┬───────┘       └───────┬───────┘       (stride=2, 只有4的倍数活跃)
Step 3:      └───────────────┬───────────────────────┘                               (stride=4, 只有8的倍数活跃)
----------------------------------------------------------------------------------------------------
[ 现代连续寻址 (Sequential) - 活跃线程两端对折，前段完全对齐，0 Bank 冲突 ]
Thread ID:   0   1   2   3   4   5   6   7   8   9  10  11  12  13  14  15 ...
Step 1:      ├───┼───┼───┼───┼───┼───┼───┼───┘   |   |   |   |   |   |   |       (stride=8, 线程 0~7 连续活跃)
             │   │   │   │   │   │   │   │
Step 2:      ├───┼───┼───┼───┘   |   |   |                                       (stride=4, 线程 0~3 连续活跃)
Step 3:      ├───┼───┘                                                               (stride=2, 线程 0~1 连续活跃)
====================================================================================================

Kernel 0：朴素交错寻址（Interleaved Addressing）

最直观的写法：每次迭代步长翻倍，用取模运算判断活跃线程：

// K0: 充满硬件缺陷的代码
for (unsigned int s = 1; s < blockDim.x; s *= 2) {
    if ((tid % (2 * s)) == 0) { // ❌ 致命的分支发散与取模开销！
        sdata[tid] += sdata[tid + s]; // ❌ 随着 s 增大，严重的 Bank Conflict！
    }
    __syncthreads();
}

缺陷 1：分支发散（Warp Divergence）。在第一轮迭代中，只有偶数线程干活，奇数线程休眠，Warp 内部 50% 算力被废除；第二轮只有 25% 活跃……SIMT 核心极度空转。
缺陷 2：Bank 冲突。当 $s \ge 32$ 时，活跃线程访问的共享内存地址相差 32 的倍数，所有线程撞击在同一个 Bank 上！

Kernel 1 & 2：消除发散与连续寻址（Sequential Addressing）

将寻址方式彻底颠覆为“两端对折”：步长从一半开始逐步减半，所有活跃线程紧密排布在 Block 的前半截：

// K2: 消除发散与 Bank 冲突
for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
    if (tid < s) { // ✅ 线程 0 到 s-1 连续活跃！
        sdata[tid] += sdata[tid + s]; // ✅ 连续线程访问连续地址，0 Bank 冲突！
    }
    __syncthreads();
}

由于一个 Warp 是 32 线程连续排布的，当

s \ge 32

时，前几个 Warp 是 100% 全员满载工作，后几个 Warp 则是 100% 全员休眠，硬件完全不会产生任何 Warp 内部的分支发散！

Kernel 3 & 4：展开最后一个 Warp 与完全循环展开（Loop Unrolling）

当步长缩小到

s &lt; 32

时，所有的工作已经收敛到了 Warp 0 的前 32 个线程。在 SIMT 物理微架构中，同一个 Warp 内的 32 个线程是硬件时钟严格同步锁步（Lock-Step）执行的！因此，最后 5 次迭代根本不需要任何 __syncthreads() 同步指令！直接在代码里手动展开（Unroll），可以省下昂贵的片上屏障同步开销。

1.3 Warp Shuffle 硬件微架构原语：跨 Lane 寄存器直通网络

上述优化无论再怎么精巧，数据始终需要在 SM 片上的 Shared Memory（共享内存） 里存入、读取、同步。而在 NVIDIA Kepler（SM 3.0）架构以后，硬件工程师在 SM 内部的 32 个 Lane（寄存器切片）之间，铺设了一条专属的高速双向 Crossbar 硬件数据网络——Warp Shuffle 原语。

核心原语家族：

__shfl_sync(mask, val, srcLane)：向指定编号的 Lane 索要寄存器中的 val；
__shfl_up_sync(mask, val, delta)：向 Lane ID 较小的邻居线程索要数据；
__shfl_down_sync(mask, val, delta)（规约核心原语）：当前线程从 laneId + delta 的邻居线程直接读取寄存器 val，在 1 个时钟周期内完成数据跨线程传递，完全不走 Shared Memory，0 访存延迟，0 Bank 冲突，天然无需同步！

5 步完成 32 线程 Warp 级规约的极简魔法：

__device__ inline float warpReduceSum(float val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val; // Lane 0 此时保存了整个 Warp 32 线程的累加总和！
}

每轮迭代折半（16

\rightarrow

\rightarrow

\rightarrow

\rightarrow

1），只需 5 条 SASS 级 SHFL 汇编指令，耗时仅几个时钟周期，便可完成一个 Warp 内的完整规约！

1.4 Ringi 工程师五问：规约计算视角下的 Shape 与 Cost

📐 Shape 是什么：输入张量的行宽 $N$ 与 Batch 行数 $M$ 分别是多少？ $N$ 是小于 1024（单 Block 搞定）、小于 32（单 Warp 搞定）还是上万（多 Block 层次规约）？
💰 Cost 花在哪里：算术强度只有不到 0.25 FLOP/Byte，时间 90% 以上花在 HBM 读写和片内等待数据搬运上。
⚙️ Machine 怎么跑：SM 内部是走多级规约（Warp Shuffle $\rightarrow$ Shared Memory $\rightarrow$ Warp 0 Shuffle），还是多个 Block 跨 SM 做原子操作（atomicAdd）？
🔍 Evidence 在哪里：Nsight Compute 中是否出现大量的 sm__sass_lsu_write_bytes_mem_shared？Warp Shuffle 的占比是否超过 80%？
🏭 Production 怎么选：在生产环境中，单行规约通常使用单个 Block 处理，利用多 Block 覆盖外层的行维度 $M$ ，最大化网格级并行（Grid-Level Concurrency）。

2. 数值稳定性第一性原理：为什么 Softmax 必须做“Safe”保护？

2.1 IEEE-754 浮点数的物理边界：FP16 与 FP32 的溢出悬崖

在标准数学定义中：

\text{Softmax}(x)_i = \frac{e^{x_i}}{\sum_{j=1}^N e^{x_j}}

在理想数学世界中，这个公式完美无瑕。但在由硅片晶体管构筑的有限精度浮点世界（IEEE-754 标准）中，它是一座极其危险的活火山。让我们检视硬件存储的真实物理极限：

浮点数格式 (IEEE-754)	符号位 (Sign)	指数位 (Exponent)	尾数位 (Mantissa)	最大可表示正实数 ( $MaxVal$ )	触发 $\exp(x)$ 上溢的临界阈值 ( $x_{\text{overflow}}$ )	触发 $\exp(x)$ 下溢为 0 的临界阈值 ( $x_{\text{underflow}}$ )
FP32 (单精度)	1 bit	8 bits	23 bits	$\approx 3.4028 \times 10^{38}$	$x > 88.7228$	$x < -103.972$
FP16 (半精度)	1 bit	5 bits	10 bits	$65504$	$x > 11.0898$ (极易触碰！)	$x < -17.3287$
BF16 (Brain Float)	1 bit	8 bits	7 bits	$\approx 3.3895 \times 10^{38}$	$x > 88.7228$	$x < -88.7228$

在大模型计算

QK^T / \sqrt{d_k}

时，未经归一化的点积绝对值很容易达到 15 到 30。 在 FP16 精度下，只要某个注意力得分达到 12.0，其指数运算结果就会直接变成 +Inf！ 一旦分母出现 +Inf，或者分子分母同时出现 +Inf，算子输出直接崩溃为 NaN。这就是为什么在工业级生产中，绝不允许直接对输入张量执行朴素 Softmax。

2.2 No Naked Formula 2.0：平移不变性与 Safe Softmax 模型

为了驯服狂暴的指数运算，我们严格走完 No Naked Formula 2.0 五步穿透法：

① 为什么需要算它？

消除指数运算的上溢爆炸，将所有输入数值安全收敛到浮点表示范围的最健壮区间。

② Mental Model（物理直觉比喻）

想象全班同学去称体重。秤的最大量程只有 100 公斤，但班上有个 150 公斤的巨汉，直接把秤踩爆了（上溢 +Inf）。怎么办？班长先扫视全班，找出全班最重的人（假设正是 150 公斤），然后让所有人站在秤上之前，都从口袋里掏出一个标称“减去 150 公斤”的负重块。此时，最重的人体重变成了

150 - 150 = 0

公斤；其余所有人体重全是负数（

-10, -20, \dots

）。由于任何非正数的指数

\exp(\le 0) \in (0, 1]

，所有数值被严严实实地锁定在

(0, 1]

的安全量程内，秤永远不可能被踩爆！最后算比例时，由于分子分母都被同等缩放，最终归一化概率分毫不差！

③ Tiny Calculator（极简数字手算）

设输入向量只有 3 个小数字：

X = [2.0, 4.0, 1.0]

。

第一步：求最大值

m = \max(2.0, 4.0, 1.0) = 4.0

第二步：平移输入向量

\tilde{X} = X - m = [2 - 4, 4 - 4, 1 - 4] = [-2.0, 0.0, -3.0]

第三步：求指数与和

e^{\tilde{x}_0} = e^{-2} \approx 0.1353, \quad e^{\tilde{x}_1} = e^0 = 1.0000, \quad e^{\tilde{x}_2} = e^{-3} \approx 0.0498

d = \sum e^{\tilde{x}_i} = 0.1353 + 1.0000 + 0.0498 = 1.1851

第四步：归一化

y = \left[ \frac{0.1353}{1.1851}, \frac{1.0000}{1.1851}, \frac{0.0498}{1.1851} \right] \approx [0.1142, 0.8438, 0.0420]

检查总和：

0.1142 + 0.8438 + 0.0420 = 1.0000

。数值完全正确！

④ Formal Model（标准公式）

定义 Safe Softmax 标准数学模型：

m = \max_{1 \le k \le N} x_k

\text{SafeSoftmax}(x)_i = \frac{e^{x_i - m}}{\sum_{j=1}^N e^{x_j - m}}

⑤ Sanity Check（代数恒等证明）

我们证明该平移操作不改变数学本质：

\frac{e^{x_i - m}}{\sum_{j=1}^N e^{x_j - m}} = \frac{e^{x_i} \cdot e^{-m}}{\sum_{j=1}^N (e^{x_j} \cdot e^{-m})} = \frac{e^{x_i} \cdot e^{-m}}{e^{-m} \cdot \sum_{j=1}^N e^{x_j}} = \frac{e^{x_i}}{\sum_{j=1}^N e^{x_j}} = \text{Softmax}(x)_i

证毕！ 数学上严格恒等，物理上彻底杜绝上溢。

2.3 传统 Safe Softmax 的三遍扫描之殇（3-Pass Memory Wall）

Safe Softmax 解决了数值稳定性，但却将硬件推入了另一个痛苦的泥潭——三遍扫描显存墙：

总数据读取量： $3 \times 4N = 12N$ 字节；
总数据写出量： $1 \times 4N = 4N$ 字节（忽略标量 $m$ 和 $d$ ）；
总显存流量： $16N$ 字节！对于一个简单的逐元素归一化算子，每个元素需要被反复搬运 4 次。这就引出了系统架构师终极的灵魂拷问： 为什么分母 $d = \sum e^{x_i - m}$ 一定要等待全量 $m$ 算完才能动工？能不能在求 $m$ 的同时，把分母 $d$ 也一起算了？

3. Online Softmax 算法推导：如何在一遍扫描中同时求 Max 和 Sum？

3.1 核心洞察：动态修正因子（Rescale Factor）的代数美学

2018 年，NVIDIA 科学家 Maxim Milakov 与 Natalia Gimelshein 在论文《Online normalizer calculation for softmax》中首次提出了震惊业界的 Online Softmax。他们的核心洞察极其优美：我们在流式遍历一个数组时，分母之所以不能提前算，是因为当前已见的最大值可能会在后面被推翻。设我们在处理前

k

个元素时，当前的最大值是

m_{\text{old}}

，累加的指数和是：

d_{\text{old}} = \sum_{j=1}^k e^{x_j - m_{\text{old}}}

如果在读到第

k+1

个元素

x_{k+1}

时，突然发现它比历史最大值还要大（

x_{k+1} > m_{\text{old}}

），此时新的最大值变成了：

m_{\text{new}} = x_{k+1}

按照传统思维，前面

k

个元素全算错了，必须推倒重来。 但且慢！真的需要重算吗？ 让我们观察如果用新的

m_{\text{new}}

来衡量历史总和，历史总和应该变成什么：

d_{\text{correct}} = \sum_{j=1}^k e^{x_j - m_{\text{new}}} = \sum_{j=1}^k e^{(x_j - m_{\text{old}}) + (m_{\text{old}} - m_{\text{new}})} = \left(\sum_{j=1}^k e^{x_j - m_{\text{old}}}\right) \cdot e^{m_{\text{old}} - m_{\text{new}}}

请屏住呼吸盯着这个公式：括号里的东西，不正是我们刚才已经累加好的 $d_{\text{old}}$ 吗？！这意味着：面对新的更大值，历史上的分母根本不需要重新计算，只需要乘以一个动态缩放因子（Rescale Factor）：

\alpha = e^{m_{\text{old}} - m_{\text{new}}}

然后再加上新元素的贡献

e^{x_{k+1} - m_{\text{new}}}

，就得到了最新的总分母！

3.2 No Naked Formula 2.0：单元素增量递推模型

我们再次执行 No Naked Formula 2.0，手算验证单元素在线递推模型：

① 为什么需要算它？

消除第一遍与第二遍扫描的串行依赖，使 Max 与 Sum 能够在单次循环中完全流式融合。

② Mental Model（物理直觉）

还是全班称体重的比喻。班长不再提前通读全名册，而是让同学们一个一个排队进门。进门第 1 个人体重 60 公斤，班长记录当前最高分 60，调整分和为

e^{60-60} = 1

；进门第 2 个人体重 50 公斤，未破纪录，班长直接把他的调整分

e^{50-60} = e^{-10}

加到总和里；进门第 3 个人体重 80 公斤！新纪录诞生！原本以为最高是 60，现在变成了 80。班长不需要把前两个人叫回来重新称，只需掏出计算器，把刚才记在账本上的总和乘以 $e^{60 - 80} = e^{-20}$ ，再加上第 3 个人的

e^{80-80} = 1

。账本瞬间更新完毕！

③ Tiny Calculator（手算 3 个数字）

继续使用刚才的数组：

X = [2.0, 4.0, 1.0]

。初始状态设为：

m_0 = -\infty, d_0 = 0.0

。

处理元素 $x_1 = 2.0$ ：

m_1 = \max(-\infty, 2.0) = 2.0

d_1 = 0.0 \cdot e^{-\infty - 2.0} + e^{2.0 - 2.0} = 0 + 1.0 = 1.0

状态： $m_1 = 2.0, d_1 = 1.0$ 。

处理元素 $x_2 = 4.0$ （出现更大值！）：

m_2 = \max(2.0, 4.0) = 4.0

d_2 = d_1 \cdot e^{m_1 - m_2} + e^{x_2 - m_2} = 1.0 \cdot e^{2.0 - 4.0} + e^{4.0 - 4.0} = e^{-2} + 1.0 \approx 0.1353 + 1.0 = 1.1353

状态： $m_2 = 4.0, d_2 = 1.1353$ 。

处理元素 $x_3 = 1.0$ （小于当前最大值）：

m_3 = \max(4.0, 1.0) = 4.0

d_3 = d_2 \cdot e^{4.0 - 4.0} + e^{1.0 - 4.0} = 1.1353 \cdot 1.0 + e^{-3} \approx 1.1353 + 0.0498 = \mathbf{1.1851}

状态： $m_3 = 4.0, d_3 = 1.1851$ 。 对比检验：在 2.2 节用传统三遍法算出来的分母正是 $1.1851$ ！仅用了一次循环，我们在求出全局最大值

4.0

的同一瞬间，分母 $1.1851$ 也分毫不差地同时算出来了！

④ Formal Model（递推公式）

对于序列中的任意新元素

x_k

：

m_k = \max(m_{k-1}, x_k)

d_k = d_{k-1} \cdot e^{m_{k-1} - m_k} + e^{x_k - m_k}

⑤ Sanity Check（数值安全性）

因为 $m_k \ge m_{k-1}$ ，所以指数项 $m_{k-1} - m_k \le 0$ ；
缩放因子 $e^{m_{k-1} - m_k} \in (0, 1]$ ，永远是一个小于等于 1 的衰减系数，绝对不可能产生上溢爆炸！

3.3 树形规约结合律代数推导（Two-Block Merge）

在单线程上，我们可以串行流式处理；但在 GPU 上，必须成百上千个线程并行计算。假设线程 A 处理了前半截数据，得到局部状态

(m_A, d_A)

；线程 B 处理了后半截数据，得到局部状态

(m_B, d_B)

。 我们能否直接将这两个状态合并为一个整体状态 $(m_{AB}, d_{AB})$ ？

严格代数推导：

设数据集合

A

的元素为

x_i

，集合

B

的元素为

x_j

。定义：

m_A = \max_{i \in A} x_i, \quad d_A = \sum_{i \in A} e^{x_i - m_A}

m_B = \max_{j \in B} x_j, \quad d_B = \sum_{j \in B} e^{x_j - m_B}

对于合并后的全集

C = A \cup B

：

合并最大值：

m_C = \max(m_A, m_B)

合并总分母：

d_C = \sum_{k \in C} e^{x_k - m_C} = \sum_{i \in A} e^{x_i - m_C} + \sum_{j \in B} e^{x_j - m_C}

将

m_A

和

m_B

拆解代入：

d_C = \sum_{i \in A} \left(e^{x_i - m_A} \cdot e^{m_A - m_C}\right) + \sum_{j \in B} \left(e^{x_j - m_B} \cdot e^{m_B - m_C}\right)

提公因式：

d_C = \left(\sum_{i \in A} e^{x_i - m_A}\right) \cdot e^{m_A - m_C} + \left(\sum_{j \in B} e^{x_j - m_B}\right) \cdot e^{m_B - m_C}

代入

d_A, d_B

：

d_{\text{merged}} = d_A \cdot e^{m_A - m_{\text{merged}}} + d_B \cdot e^{m_B - m_{\text{merged}}}

这个公式具有神圣的对称性与结合律！ 它证明了：无论你是单线程增量处理，还是用 32 个线程做 Warp Shuffle 规约，亦或是跨 Warp 做 Block 级树形合并，都可以直接套用这个二元合并算子！

4. 工业级工程实现：从 2-Pass Online Softmax 到 1-Pass Fused Softmax

4.1 核心架构：Warp Shuffle 双值规约 + Warp 间共享内存交换

在真实的 CUDA Kernel 中，我们将上述结合律直接翻译为高性能的 Warp Shuffle 规约原语：

// 在 Warp 内同时完成 (m, d) 的双值在线规约
__device__ inline void warpReduceOnline(float& m, float& d) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        // 跨 Lane 交换邻居的 m 和 d (仅需 2 条寄存器 Shuffle 指令)
        float m_other = __shfl_down_sync(0xffffffff, m, offset);
        float d_other = __shfl_down_sync(0xffffffff, d, offset);

        // 套用代数结合律进行在线合并
        float m_new = fmaxf(m, m_other);
        d = d * __expf(m - m_new) + d_other * __expf(m_other - m_new);
        m = m_new;
    }
}

完整的两级规约架构：

线程级累积：每个线程负责处理连续多个元素（如步长循环），在寄存器中维护单线程的 local_m 与 local_d；
第一级（Warp 内）：调用 warpReduceOnline，由 32 个线程在寄存器层面瞬时聚合成每个 Warp 的局部结果；
Warp 间交换：每个 Warp 的 Lane 0 将本 Warp 的 (m, d) 写入大小仅为 32 的微型共享内存数组：__shared__ float s_m[32], s_d[32];；
第二级（Block 内）：Warp 0 的前几个线程从共享内存读出所有 Warp 的代表值，再次调用一次 warpReduceOnline；
单周期广播：Lane 0 将最终全局的 row_max 和 row_sum 写入单值共享内存，广播给整个 Block。

4.2 终极一跃：1-Pass Fused Softmax（寄存器缓存消除最后一次全局读）

在传统的 Online Softmax 中，虽然 Max 和 Sum 被合并成了一次循环（第 1 遍读数据），但最终算归一化写出时，依然需要从全局显存把输入数据读进 SM 算一遍

\exp(x - m)/d

（第 2 遍读数据）。这被称为 2-Pass Online Softmax。 能不能把第二遍读取也彻底消灭掉？做到真正的 1-Pass 极致性能？

====================================================================================================
            1-Pass Fused Softmax 核心突破: 利用线程局部寄存器数组充当“超光速缓存”
====================================================================================================
每个线程私有寄存器: float reg_cache[MAX_ELEMS_PER_THREAD]; (在 A100 上高达 255 个寄存器/线程)

[ 唯一的一遍扫描 (Pass 1) ]
  从全局显存读取 x_i ──► 存入 reg_cache[k] (完全不写 HBM，驻留寄存器堆)
                      └──► 顺便执行 Online 递推更新 (local_m, local_d)

[ 规约阶段 (片内零延迟) ]
  Block 级两级 Shuffle 规约 ──► 瞬间得出最终 (row_max, inv_sum)

[ 归一化写出阶段 (Pass 2 原地从寄存器读取) ]
  从私有寄存器读取 reg_cache[k] ──► 算 exp(val - row_max) * inv_sum ──► 写入全局显存

★ 彻底消除第二次 HBM 显存读取！全局显存访问降至纯粹的 1 读 1 写 (8N 字节)！
====================================================================================================

当矩阵的行宽

N

在大模型常见隐藏层维度内（例如

N \le 4096

或

N \le 8192

）时，如果每个 Block 分配 256 个线程：

\text{每个线程需要处理的元素数} = \frac{4096}{256} = 16 \text{ 个 float}

16 个 float 仅仅消耗每个线程 16 个 32-bit 寄存器！而 A100 每个线程拥有高达 255 个寄存器可用。我们完全可以声明一个局部数组 float reg_cache[16]。在第 1 遍从全局显存读取时，顺手将数据保存在寄存器数组中；规约完成后，第 2 遍直接从寄存器中取数写出！ 全局显存的总访问量被压到了物理理论下限：读 1 次输入，写 1 次输出！ 显存流量相比传统 Safe Softmax 的 16N 字节暴降整整 50%，算子性能直接打到硬件 Roofline 的极限顶峰！

4.3 向量化与多行网格跨步调度（Grid-Stride Multi-Row Parallelism）

为了将 1-Pass Fused Softmax 封装为工业级通用算子，我们还需要解决两个工程细节：

向量化加载（float4）：单线程每次处理 4 个元素，使用 LDG.128 指令，进一步压低循环展开与指令分发开销；
多行网格跨步（Grid-Stride Multi-Row）：当输入矩阵非常庞大（例如 Batch 很大，行数 $M = 32768$ $M = 32768$ ）时，Block 数量可能小于行数。采用网格跨步循环：
for (int row = blockIdx.x; row < M; row += gridDim.x) { // 处理当前行 row... }
允许以固定的 Block 规模平滑处理任意规模的张量，避免因为动态申请过多 Block 导致硬件调度溢出。

5. FlashAttention 核心前置：Online Softmax 是如何成就大模型注意力革新的？

Ringi 导师解构：Online Softmax 动态 Rescale 修正因子与 FlashAttention 融合工坊

5.1 为什么标准 Attention 必须保存庞大的 $S = QK^T$ 矩阵到显存？

如果不理解 Online Softmax，你就永远无法真正看懂大模型基础设施领域最具革命性的工作——FlashAttention（Tri Dao et al., 2022）。在标准多头注意力机制中，算法流程是串行的：

$S = QK^T \in \mathbb{R}^{N \times N}$ （写入 HBM 全局显存）；
$P = \text{Softmax}(S) \in \mathbb{R}^{N \times N}$ （从 HBM 读 $S$ ，算完写回 $P$ ）；
$O = PV \in \mathbb{R}^{N \times d}$ （从 HBM 读 $P$ 和 $V$ ，算完写回 $O$ ）。

在长文本下，

N \times N

的中间矩阵

S

和

P

的体积是按序列长度的平方级 $O(N^2)$ 爆炸式增长的！当

N = 64\text{K}

时，单头注意力矩阵需要占用 8 GB 显存！不仅显存瞬间 OOM，而且反复读写这几十 GB 的中间大矩阵，让计算管线全部被 HBM 访存卡死。为什么以前的工程师不敢直接将 Softmax 和后面的矩阵乘

PV

融合（Fuse）在一起？ 就是因为传统的 Softmax 要求必须先见识过整整一整行的全部数据，才能算出最大值 $m$ 和分母 $d$ ！ 只要你必须先见识全量行，你就不得不把长达

N

的完整行落盘在全局显存中。

5.2 FlashAttention-1/2 的核心数学解构：输出张量的动态缩放递推

FlashAttention 的核心奇迹，正是 Online Softmax 分块思想与 GEMM 的完美联姻！ Tri Dao 等人的思路极其震撼：既然输入矩阵太长放不下，我们就把

K

和

V

切分成一个个可以完全放进片上 SRAM（Shared Memory）的小 Block（比如大小为

B_c \times d

）。当我们加载第 1 块 Key/Value 计算得到局部的

S^{(1)} = Q K_1^T

时：

我们利用 Online Softmax 算出局部的最大值 $m^{(1)}$ 和分母 $d^{(1)}$ ；
并且直接用局部的注意力概率乘以此刻的 $V_1$ ，计算出局部的输出累计量 $O^{(1)}$ 存放在片上寄存器里！

当加载第 2 块 Key/Value 计算出新的局部得分

S^{(2)} = Q K_2^T

时：

计算新块的最大值 $m^{(2)}$ ，并根据结合律更新全局最大值：

m^{\text{new}} = \max(m^{(1)}, m^{(2)})

计算新的分母：

d^{\text{new}} = d^{(1)} \cdot e^{m^{(1)} - m^{\text{new}}} + d^{(2)} \cdot e^{m^{(2)} - m^{\text{new}}}

最关键的神来之笔——如何更新已经算出来的输出矩阵 $O$ ？ 利用完全相同的 Rescale 因子：

O^{\text{new}} = \text{diag}\left(e^{m^{(1)} - m^{\text{new}}}\right) O^{(1)} + P^{(2)} V_2

在遍历完所有分块后，只需在最终做一次全局除法：

O = O / d^{\text{final}}

！

5.3 体系结构级洞见：用 SRAM 乘法换取 HBM 流量清零

在现代计算机体系结构中，算力的摩尔定律（每代提升 2~3 倍）远快于内存总线带宽的物理提升（每代提升 30%~50%）。 算力是廉价的，而显存搬运是极其昂贵的。 FlashAttention 和 Online Softmax 的灵魂，就在于宁可在片上 SRAM 和寄存器中多做几次乘法修正（Rescale Multiply），也绝对不向慢速的全局显存写出哪怕一个中间字节！ 通过这种算法与体系结构的深度协同，中间注意力得分矩阵被彻底“抹杀”在片上缓存中，显存占用从

O(N^2)

骤降到

O(N)

，大模型长文本处理从此彻底告别了显存爆炸的时代。

6. 动手实战与代码实验室（Minimal Runnable Code）

本章提供 4 个生产级微基准测试程序。代码严格遵循 Full-Output Enforcement 原则，绝无任何省略号或未实现函数，自带 checkCuda 错误校验与微秒级计时，可以直接使用 nvcc 编译运行并输出清晰的对比证据链。

实验 1：树形规约演进基准（`reduction_evolution_benchmark.cu`）

本实验完整对比 Mark Harris 经典演化中的 3 个标志性版本：

K0（交错寻址 Interleaved，含取模与发散）；
K2（连续寻址 Sequential，两端对折，0 Bank 冲突）；
K6（Warp Shuffle 规约，跨 Lane 寄存器直通网络）。

/**
 * File: reduction_evolution_benchmark.cu
 * Compile: nvcc -O3 -arch=native -o reduction_evolution_benchmark reduction_evolution_benchmark.cu
 * Run: ./reduction_evolution_benchmark
 */

#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>

#define checkCuda(ans) { gpuAssert((ans), __FILE__, __LINE__); }
inline void gpuAssert(cudaError_t code, const char *file, int line) {
    if (code != cudaSuccess) {
        fprintf(stderr, "CUDA Error: %s in %s at line %d\n", cudaGetErrorString(code), file, line);
        exit(code);
    }
}

// K0: 朴素交错寻址 (取模 + 严重发散 + Bank冲突)
__global__ void reduce_k0_interleaved(const float* __restrict__ g_in, float* __restrict__ g_out, int n) {
    extern __shared__ float sdata[];
    int tid = threadIdx.x;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    sdata[tid] = (idx < n) ? g_in[idx] : 0.0f;
    __syncthreads();

    for (unsigned int s = 1; s < blockDim.x; s *= 2) {
        if ((tid % (2 * s)) == 0) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    if (tid == 0) g_out[blockIdx.x] = sdata[0];
}

// K2: 连续寻址对折 (无分支发散 + 0 Bank冲突)
__global__ void reduce_k2_sequential(const float* __restrict__ g_in, float* __restrict__ g_out, int n) {
    extern __shared__ float sdata[];
    int tid = threadIdx.x;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    sdata[tid] = (idx < n) ? g_in[idx] : 0.0f;
    __syncthreads();

    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    if (tid == 0) g_out[blockIdx.x] = sdata[0];
}

// K6: Warp Shuffle 寄存器直通规约
__inline__ __device__ float warp_reduce_sum(float val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__global__ void reduce_k6_shuffle(const float* __restrict__ g_in, float* __restrict__ g_out, int n) {
    __shared__ float warp_sums[32];
    int tid = threadIdx.x;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int lane = tid % 32;
    int wid = tid / 32;

    float val = (idx < n) ? g_in[idx] : 0.0f;
    val = warp_reduce_sum(val);

    if (lane == 0) {
        warp_sums[wid] = val;
    }
    __syncthreads();

    // 最后一个 Warp 完成块内汇总
    val = (tid < (blockDim.x / 32)) ? warp_sums[lane] : 0.0f;
    if (wid == 0) {
        val = warp_reduce_sum(val);
        if (lane == 0) {
            g_out[blockIdx.x] = val;
        }
    }
}

int main() {
    int dev = 0;
    checkCuda(cudaSetDevice(dev));
    cudaDeviceProp prop;
    checkCuda(cudaGetDeviceProperties(&prop, dev));

    printf("================================================================================\n");
    printf("实验 1: 并行规约 (Reduction) 三级演进基准测试\n");
    printf("设备: %s\n", prop.name);
    printf("================================================================================\n");

    const int N = 32 * 1024 * 1024; // 32M 个 float (128 MB)
    size_t bytes = N * sizeof(float);
    const int threads_per_block = 256;
    const int num_blocks = (N + threads_per_block - 1) / threads_per_block;

    float *d_in = NULL, *d_out = NULL;
    checkCuda(cudaMalloc(&d_in, bytes));
    checkCuda(cudaMalloc(&d_out, num_blocks * sizeof(float)));
    checkCuda(cudaMemset(d_in, 0x1, bytes));

    cudaEvent_t start, stop;
    checkCuda(cudaEventCreate(&start));
    checkCuda(cudaEventCreate(&stop));
    const int runs = 20;

    // 1. 测试 K0
    checkCuda(cudaEventRecord(start));
    for (int i = 0; i < runs; ++i) {
        reduce_k0_interleaved<<<num_blocks, threads_per_block, threads_per_block * sizeof(float)>>>(d_in, d_out, N);
    }
    checkCuda(cudaEventRecord(stop));
    checkCuda(cudaEventSynchronize(stop));
    float ms_k0; checkCuda(cudaEventElapsedTime(&ms_k0, start, stop)); ms_k0 /= runs;

    // 2. 测试 K2
    checkCuda(cudaEventRecord(start));
    for (int i = 0; i < runs; ++i) {
        reduce_k2_sequential<<<num_blocks, threads_per_block, threads_per_block * sizeof(float)>>>(d_in, d_out, N);
    }
    checkCuda(cudaEventRecord(stop));
    checkCuda(cudaEventSynchronize(stop));
    float ms_k2; checkCuda(cudaEventElapsedTime(&ms_k2, start, stop)); ms_k2 /= runs;

    // 3. 测试 K6
    checkCuda(cudaEventRecord(start));
    for (int i = 0; i < runs; ++i) {
        reduce_k6_shuffle<<<num_blocks, threads_per_block>>>(d_in, d_out, N);
    }
    checkCuda(cudaEventRecord(stop));
    checkCuda(cudaEventSynchronize(stop));
    float ms_k6; checkCuda(cudaEventElapsedTime(&ms_k6, start, stop)); ms_k6 /= runs;

    double giga_bytes = (double)bytes / 1e9;
    printf("  [K0 交错寻址 (取模分支发散)] 耗时: %7.3f ms | 有效带宽: %7.2f GB/s\n", ms_k0, giga_bytes / (ms_k0 * 1e-3));
    printf("  [K2 连续对折 (消除发散与冲突)] 耗时: %7.3f ms | 有效带宽: %7.2f GB/s\n", ms_k2, giga_bytes / (ms_k2 * 1e-3));
    printf("  [K6 Warp Shuffle (寄存器直通)] 耗时: %7.3f ms | 有效带宽: %7.2f GB/s\n", ms_k6, giga_bytes / (ms_k6 * 1e-3));
    printf("  --> Shuffle 相比原始 K0 加速比: %.2fx (片内寄存器直通的巨大飞跃！)\n", ms_k0 / ms_k6);
    printf("================================================================================\n\n");

    checkCuda(cudaFree(d_in));
    checkCuda(cudaFree(d_out));
    checkCuda(cudaEventDestroy(start));
    checkCuda(cudaEventDestroy(stop));
    return 0;
}

实验 2：经典 3-Pass Safe Softmax 基准（`safe_softmax_3pass_benchmark.cu`）

本实验模拟生产环境中未融合的标准 Safe Softmax 实现，由三个连续调用的 Kernel 组成，精确测量 3 遍读写全局显存所产生的性能开销。

/**
 * File: safe_softmax_3pass_benchmark.cu
 * Compile: nvcc -O3 -arch=native -o safe_softmax_3pass_benchmark safe_softmax_3pass_benchmark.cu
 * Run: ./safe_softmax_3pass_benchmark
 */

#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>
#include <float.h>

#define checkCuda(ans) { gpuAssert((ans), __FILE__, __LINE__); }
inline void gpuAssert(cudaError_t code, const char *file, int line) {
    if (code != cudaSuccess) {
        fprintf(stderr, "CUDA Error: %s in %s at line %d\n", cudaGetErrorString(code), file, line);
        exit(code);
    }
}

__inline__ __device__ float warp_reduce_max(float val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
    }
    return val;
}

__inline__ __device__ float warp_reduce_sum(float val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

// Pass 1: 求每行最大值 Max
__global__ void pass1_row_max(const float* __restrict__ input, float* __restrict__ row_max, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid = tid / 32;
    __shared__ float s_warp[32];

    const float* row_in = input + row * N;
    float local_max = -FLT_MAX;

    for (int col = tid; col < N; col += blockDim.x) {
        local_max = fmaxf(local_max, row_in[col]);
    }
    local_max = warp_reduce_max(local_max);

    if (lane == 0) s_warp[wid] = local_max;
    __syncthreads();

    if (wid == 0) {
        local_max = (lane < (blockDim.x / 32)) ? s_warp[lane] : -FLT_MAX;
        local_max = warp_reduce_max(local_max);
        if (lane == 0) row_max[row] = local_max;
    }
}

// Pass 2: 计算 exp(x - max) 并规约分母和 Sum
__global__ void pass2_row_sum(const float* __restrict__ input, const float* __restrict__ row_max, float* __restrict__ row_sum, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid = tid / 32;
    __shared__ float s_warp[32];

    const float* row_in = input + row * N;
    float m = row_max[row];
    float local_sum = 0.0f;

    for (int col = tid; col < N; col += blockDim.x) {
        local_sum += __expf(row_in[col] - m);
    }
    local_sum = warp_reduce_sum(local_sum);

    if (lane == 0) s_warp[wid] = local_sum;
    __syncthreads();

    if (wid == 0) {
        local_sum = (lane < (blockDim.x / 32)) ? s_warp[lane] : 0.0f;
        local_sum = warp_reduce_sum(local_sum);
        if (lane == 0) row_sum[row] = local_sum;
    }
}

// Pass 3: 执行最终归一化并写回
__global__ void pass3_normalize(const float* __restrict__ input, const float* __restrict__ row_max, const float* __restrict__ row_sum, float* __restrict__ output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    const float* row_in = input + row * N;
    float* row_out = output + row * N;

    float m = row_max[row];
    float inv_d = 1.0f / row_sum[row];

    for (int col = tid; col < N; col += blockDim.x) {
        row_out[col] = __expf(row_in[col] - m) * inv_d;
    }
}

int main() {
    int dev = 0;
    checkCuda(cudaSetDevice(dev));

    printf("================================================================================\n");
    printf("实验 2: 经典 3-Pass Safe Softmax (分离 Kernel 显存往返墙测试)\n");
    printf("================================================================================\n");

    const int M = 8192; // 8K 行
    const int N = 4096; // 每行 4K 列
    size_t total_bytes = (size_t)M * N * sizeof(float);

    float *d_in = NULL, *d_out = NULL, *d_max = NULL, *d_sum = NULL;
    checkCuda(cudaMalloc(&d_in, total_bytes));
    checkCuda(cudaMalloc(&d_out, total_bytes));
    checkCuda(cudaMalloc(&d_max, M * sizeof(float)));
    checkCuda(cudaMalloc(&d_sum, M * sizeof(float)));
    checkCuda(cudaMemset(d_in, 0x1, total_bytes));

    int block_size = 256;
    dim3 grid(M);

    cudaEvent_t start, stop;
    checkCuda(cudaEventCreate(&start));
    checkCuda(cudaEventCreate(&stop));
    const int runs = 20;

    checkCuda(cudaEventRecord(start));
    for (int i = 0; i < runs; ++i) {
        pass1_row_max<<<grid, block_size>>>(d_in, d_max, M, N);
        pass2_row_sum<<<grid, block_size>>>(d_in, d_max, d_sum, M, N);
        pass3_normalize<<<grid, block_size>>>(d_in, d_max, d_sum, d_out, M, N);
    }
    checkCuda(cudaEventRecord(stop));
    checkCuda(cudaEventSynchronize(stop));

    float ms_total; checkCuda(cudaEventElapsedTime(&ms_total, start, stop)); ms_total /= runs;
    double traffic_gb = ((double)total_bytes * 4.0) / 1e9; // 3次读 + 1次写 = 4x 数据量

    printf("  矩阵规模: (%d, %d) | 数据量: %.2f MB\n", M, N, total_bytes / 1024.0 / 1024.0);
    printf("  [3-Pass Safe Softmax] 平均总耗时: %7.3f ms\n", ms_total);
    printf("  理论显存往返流量: %7.2f GB | 实际总线带宽占用: %7.2f GB/s\n", traffic_gb, traffic_gb / (ms_total * 1e-3));
    printf("================================================================================\n\n");

    checkCuda(cudaFree(d_in));
    checkCuda(cudaFree(d_out));
    checkCuda(cudaFree(d_max));
    checkCuda(cudaFree(d_sum));
    checkCuda(cudaEventDestroy(start));
    checkCuda(cudaEventDestroy(stop));
    return 0;
}

实验 3：高性能 2-Pass / 1-Pass Online Softmax 算子（`online_softmax_benchmark.cu`）

本实验实现纯正的 Online Softmax 算法，展示利用结合律融合 Max 与 Sum（2-Pass），以及利用私有寄存器数组实现终极 1-Pass Fused Softmax 的完整工程代码。

/**
 * File: online_softmax_benchmark.cu
 * Compile: nvcc -O3 -arch=native -o online_softmax_benchmark online_softmax_benchmark.cu
 * Run: ./online_softmax_benchmark
 */

#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>
#include <float.h>

#define checkCuda(ans) { gpuAssert((ans), __FILE__, __LINE__); }
inline void gpuAssert(cudaError_t code, const char *file, int line) {
    if (code != cudaSuccess) {
        fprintf(stderr, "CUDA Error: %s in %s at line %d\n", cudaGetErrorString(code), file, line);
        exit(code);
    }
}

// Warp 级双值 Online 合并
__inline__ __device__ void warp_reduce_online(float& m, float& d) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        float m_other = __shfl_down_sync(0xffffffff, m, offset);
        float d_other = __shfl_down_sync(0xffffffff, d, offset);

        float m_new = fmaxf(m, m_other);
        d = d * __expf(m - m_new) + d_other * __expf(m_other - m_new);
        m = m_new;
    }
}

// 1. 2-Pass Online Softmax (第1遍在线算 Max+Sum, 第2遍归一化)
__global__ void online_softmax_2pass_kernel(const float* __restrict__ input, float* __restrict__ output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid = tid / 32;

    __shared__ float s_m[32], s_d[32];
    const float* row_in = input + row * N;
    float* row_out = output + row * N;

    // Pass 1: 在线递推
    float local_m = -FLT_MAX;
    float local_d = 0.0f;

    for (int col = tid; col < N; col += blockDim.x) {
        float x = row_in[col];
        float m_new = fmaxf(local_m, x);
        local_d = local_d * __expf(local_m - m_new) + __expf(x - m_new);
        local_m = m_new;
    }

    warp_reduce_online(local_m, local_d);
    if (lane == 0) {
        s_m[wid] = local_m;
        s_d[wid] = local_d;
    }
    __syncthreads();

    if (wid == 0) {
        local_m = (lane < (blockDim.x / 32)) ? s_m[lane] : -FLT_MAX;
        local_d = (lane < (blockDim.x / 32)) ? s_d[lane] : 0.0f;
        warp_reduce_online(local_m, local_d);
        if (lane == 0) {
            s_m[0] = local_m;
            s_d[0] = local_d;
        }
    }
    __syncthreads();

    float final_m = s_m[0];
    float inv_d = 1.0f / s_d[0];

    // Pass 2: 重新读取并写出
    for (int col = tid; col < N; col += blockDim.x) {
        row_out[col] = __expf(row_in[col] - final_m) * inv_d;
    }
}

// 2. 终极 1-Pass Fused Softmax (寄存器缓存消除最后一次读)
#define MAX_ELEMS_PER_THREAD 16

__global__ void online_softmax_1pass_fused_kernel(const float* __restrict__ input, float* __restrict__ output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid = tid / 32;

    __shared__ float s_m[32], s_d[32];
    const float* row_in = input + row * N;
    float* row_out = output + row * N;

    // 私有寄存器数组缓存
    float reg_cache[MAX_ELEMS_PER_THREAD];
    int count = 0;

    float local_m = -FLT_MAX;
    float local_d = 0.0f;

    // 唯一一遍读取全局显存
    for (int col = tid; col < N; col += blockDim.x) {
        float x = row_in[col];
        if (count < MAX_ELEMS_PER_THREAD) {
            reg_cache[count++] = x;
        }
        float m_new = fmaxf(local_m, x);
        local_d = local_d * __expf(local_m - m_new) + __expf(x - m_new);
        local_m = m_new;
    }

    warp_reduce_online(local_m, local_d);
    if (lane == 0) {
        s_m[wid] = local_m;
        s_d[wid] = local_d;
    }
    __syncthreads();

    if (wid == 0) {
        local_m = (lane < (blockDim.x / 32)) ? s_m[lane] : -FLT_MAX;
        local_d = (lane < (blockDim.x / 32)) ? s_d[lane] : 0.0f;
        warp_reduce_online(local_m, local_d);
        if (lane == 0) {
            s_m[0] = local_m;
            s_d[0] = local_d;
        }
    }
    __syncthreads();

    float final_m = s_m[0];
    float inv_d = 1.0f / s_d[0];

    // 原地从寄存器取数据归一化并写回！(0 次全局显存重读)
    int idx = 0;
    for (int col = tid; col < N; col += blockDim.x) {
        float x = (idx < count) ? reg_cache[idx++] : row_in[col];
        row_out[col] = __expf(x - final_m) * inv_d;
    }
}

int main() {
    int dev = 0;
    checkCuda(cudaSetDevice(dev));

    printf("================================================================================\n");
    printf("实验 3: 2-Pass Online Softmax vs 1-Pass Fused Online Softmax 对比\n");
    printf("================================================================================\n");

    const int M = 8192;
    const int N = 4096; // 4096 / 256 = 16 元素/线程，完美匹配寄存器缓存
    size_t total_bytes = (size_t)M * N * sizeof(float);

    float *d_in = NULL, *d_out = NULL;
    checkCuda(cudaMalloc(&d_in, total_bytes));
    checkCuda(cudaMalloc(&d_out, total_bytes));
    checkCuda(cudaMemset(d_in, 0x1, total_bytes));

    int block_size = 256;
    dim3 grid(M);

    cudaEvent_t start, stop;
    checkCuda(cudaEventCreate(&start));
    checkCuda(cudaEventCreate(&stop));
    const int runs = 20;

    // 测试 2-Pass
    checkCuda(cudaEventRecord(start));
    for (int i = 0; i < runs; ++i) {
        online_softmax_2pass_kernel<<<grid, block_size>>>(d_in, d_out, M, N);
    }
    checkCuda(cudaEventRecord(stop));
    checkCuda(cudaEventSynchronize(stop));
    float ms_2pass; checkCuda(cudaEventElapsedTime(&ms_2pass, start, stop)); ms_2pass /= runs;

    // 测试 1-Pass Fused
    checkCuda(cudaEventRecord(start));
    for (int i = 0; i < runs; ++i) {
        online_softmax_1pass_fused_kernel<<<grid, block_size>>>(d_in, d_out, M, N);
    }
    checkCuda(cudaEventRecord(stop));
    checkCuda(cudaEventSynchronize(stop));
    float ms_1pass; checkCuda(cudaEventElapsedTime(&ms_1pass, start, stop)); ms_1pass /= runs;

    double bytes_gb = (double)total_bytes / 1e9;
    printf("  [2-Pass Online Softmax] 耗时: %7.3f ms | 等效吞吐: %7.2f GB/s\n",
           ms_2pass, (bytes_gb * 3.0) / (ms_2pass * 1e-3));
    printf("  [1-Pass Fused Softmax ] 耗时: %7.3f ms | 等效吞吐: %7.2f GB/s\n",
           ms_1pass, (bytes_gb * 2.0) / (ms_1pass * 1e-3));
    printf("  --> 1-Pass Fused 相比 2-Pass 加速比: %.2fx (总耗时压低 %.1f%%)\n",
           ms_2pass / ms_1pass, ((ms_2pass - ms_1pass) / ms_2pass) * 100.0);
    printf("================================================================================\n\n");

    checkCuda(cudaFree(d_in));
    checkCuda(cudaFree(d_out));
    checkCuda(cudaEventDestroy(start));
    checkCuda(cudaEventDestroy(stop));
    return 0;
}

实验 4：端到端吞吐压测与显存流量对比（`softmax_benchmark_harness.cu`）

本实验将 3-Pass、2-Pass 和 1-Pass Fused 三者放在统一压测套件中，并进行严格的数值对齐与精度误差比对（验证差值小于

10^{-6}

），输出综合对比汇总表。

/**
 * File: softmax_benchmark_harness.cu
 * Compile: nvcc -O3 -arch=native -o softmax_benchmark_harness softmax_benchmark_harness.cu
 * Run: ./softmax_benchmark_harness
 */

#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <float.h>

#define checkCuda(ans) { gpuAssert((ans), __FILE__, __LINE__); }
inline void gpuAssert(cudaError_t code, const char *file, int line) {
    if (code != cudaSuccess) {
        fprintf(stderr, "CUDA Error: %s in %s at line %d\n", cudaGetErrorString(code), file, line);
        exit(code);
    }
}

// 引用实验 2 与 3 的核心内核逻辑做集成对比
// 此处定义核函数
__inline__ __device__ void warp_reduce_online_harness(float& m, float& d) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        float m_other = __shfl_down_sync(0xffffffff, m, offset);
        float d_other = __shfl_down_sync(0xffffffff, d, offset);
        float m_new = fmaxf(m, m_other);
        d = d * __expf(m - m_new) + d_other * __expf(m_other - m_new);
        m = m_new;
    }
}

#define MAX_ELEMS 16
__global__ void fused_online_softmax_kernel(const float* __restrict__ input, float* __restrict__ output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid = tid / 32;

    __shared__ float s_m[32], s_d[32];
    const float* row_in = input + row * N;
    float* row_out = output + row * N;

    float reg_cache[MAX_ELEMS];
    int count = 0;
    float local_m = -FLT_MAX;
    float local_d = 0.0f;

    for (int col = tid; col < N; col += blockDim.x) {
        float x = row_in[col];
        if (count < MAX_ELEMS) reg_cache[count++] = x;
        float m_new = fmaxf(local_m, x);
        local_d = local_d * __expf(local_m - m_new) + __expf(x - m_new);
        local_m = m_new;
    }

    warp_reduce_online_harness(local_m, local_d);
    if (lane == 0) { s_m[wid] = local_m; s_d[wid] = local_d; }
    __syncthreads();

    if (wid == 0) {
        local_m = (lane < (blockDim.x / 32)) ? s_m[lane] : -FLT_MAX;
        local_d = (lane < (blockDim.x / 32)) ? s_d[lane] : 0.0f;
        warp_reduce_online_harness(local_m, local_d);
        if (lane == 0) { s_m[0] = local_m; s_d[0] = local_d; }
    }
    __syncthreads();

    float final_m = s_m[0];
    float inv_d = 1.0f / s_d[0];

    int idx = 0;
    for (int col = tid; col < N; col += blockDim.x) {
        float x = (idx < count) ? reg_cache[idx++] : row_in[col];
        row_out[col] = __expf(x - final_m) * inv_d;
    }
}

// CPU 精度基准验证
void cpu_safe_softmax(const float* in, float* out, int M, int N) {
    for (int r = 0; r < M; ++r) {
        float m = -FLT_MAX;
        for (int c = 0; c < N; ++c) m = fmaxf(m, in[r * N + c]);
        float d = 0.0f;
        for (int c = 0; c < N; ++c) d += expf(in[r * N + c] - m);
        float inv_d = 1.0f / d;
        for (int c = 0; c < N; ++c) out[r * N + c] = expf(in[r * N + c] - m) * inv_d;
    }
}

int main() {
    int dev = 0;
    checkCuda(cudaSetDevice(dev));
    cudaDeviceProp prop;
    checkCuda(cudaGetDeviceProperties(&prop, dev));

    printf("================================================================================\n");
    printf("实验 4: Softmax 全景性能压测套件与 CPU 精度校验\n");
    printf("测试平台: %s\n", prop.name);
    printf("================================================================================\n");

    const int M = 1024;
    const int N = 4096;
    size_t bytes = M * N * sizeof(float);

    float *h_in = (float*)malloc(bytes);
    float *h_out_gpu = (float*)malloc(bytes);
    float *h_out_cpu = (float*)malloc(bytes);

    // 随机生成带大数的测试数据，考察数值稳定性
    srand(42);
    for (int i = 0; i < M * N; ++i) {
        h_in[i] = ((float)rand() / RAND_MAX) * 50.0f - 10.0f; // [-10, 40] 之间，若不做 Safe 必炸
    }

    float *d_in = NULL, *d_out = NULL;
    checkCuda(cudaMalloc(&d_in, bytes));
    checkCuda(cudaMalloc(&d_out, bytes));
    checkCuda(cudaMemcpy(d_in, h_in, bytes, cudaMemcpyHostToDevice));

    // 运行 GPU 1-Pass Fused
    fused_online_softmax_kernel<<<M, 256>>>(d_in, d_out, M, N);
    checkCuda(cudaDeviceSynchronize());
    checkCuda(cudaMemcpy(h_out_gpu, d_out, bytes, cudaMemcpyDeviceToHost));

    // 运行 CPU 验证前 16 行
    cpu_safe_softmax(h_in, h_out_cpu, 16, N);

    float max_err = 0.0f;
    for (int i = 0; i < 16 * N; ++i) {
        float err = fabsf(h_out_gpu[i] - h_out_cpu[i]);
        if (err > max_err) max_err = err;
    }

    printf("  [数值正确性校验] 前 16 行最大绝对浮点误差: %.2e\n", max_err);
    if (max_err < 1e-5) {
        printf("  >>> 精度校验通过！Online Softmax 与 CPU 理论值完全对齐！<<<\n");
    } else {
        printf("  >>> 精度警告！误差过大！<<<\n");
    }
    printf("--------------------------------------------------------------------------------\n\n");

    free(h_in); free(h_out_gpu); free(h_out_cpu);
    checkCuda(cudaFree(d_in)); checkCuda(cudaFree(d_out));
    return 0;
}

7. Ringi 避坑指南与生产黄金准则

7.1 避坑表格：7 大常见小白错误理解 vs 大厂 AI Infra 正确认知

序号	❌ 常见小白错误理解	✅ 大厂 AI Infra 正确物理认知与一线工程军规
1	“Softmax 只是一层激活函数，直接写一个简单的公式就行”	绝非如此！未经平移保护的 Softmax 在 FP16 下只要输入大于 11.1 就会产生指数溢出，导致整网彻底报废为 `NaN`。
2	“求最大值 $m$ 和求分母 $d$ 存在前置依赖，无论如何都必须先扫一遍数据”	这是经典串行思维。通过 Milakov & Gimelshein 2018 的动态修正因子公式，Max 与 Sum 完全可以在单次遍历中以结合律优雅流式融合。
3	“Warp Shuffle 每次只能传一个浮点数，没法同时做复杂的聚合”	双值甚至多值规约完全可以在同一个循环中调用多条 `__shfl_down_sync` 分别传输，并在 Lane 内部结合代数公式原地完成二元合并。
4	“规约求和最好让线程 0 依次加其他线程，代码最简单”	这种做法会将 $O(\log N)$ 的规约树彻底串行化退化为 $O(N)$ ，且触发极度严重的单线程锁步等待，性能暴跌上百倍。
5	“Warp 内的线程规约也必须像 Block 一样写 `__syncthreads()` 同步”	严禁在 Warp 内乱加同步指令！Warp 内 32 线程由 SIMT 硬件发射器严格按时钟锁步执行，Shuffle 指令硬件天然保序，同步纯属画蛇添足。
6	“只要算子写得足够快，显存流量多几次无所谓”	Softmax 是极度严苛的 Memory-Bound 算子。算力从来不是瓶颈，显存总线的往返次数直接决定生死。1-Pass 相比 3-Pass 减少 50% 流量，提速接近翻倍。
7	“FlashAttention 是一个极其复杂的黑盒，跟基础算子没关系”	FlashAttention 的数学本质就是 Online Softmax 的分块动态 Rescale 机制与 GEMM 的片上融合。搞懂 Online Softmax，FlashAttention 迎刃而解。

7.2 生产性能工程黄金 Checklist

1. 【绝对数值防溢出校验】：任何生产级 Softmax Kernel，首要步骤必须执行平移保护减去局部/全局最大值 $m$ ；在 FP16 模式下，输入必须严密钳位。
2. 【规约严禁使用取模运算】：在树形规约中，全面禁止使用 tid % (2 * s) 寻址；强制使用连续对折寻址（tid < s）消除 Warp 分支发散与 Bank 冲突。
3. 【Warp 级一律改用 Shuffle 原语】：对于最后 32 个线程的规约，全面废弃共享内存中转，强制采用 __shfl_down_sync 寄存器直通交换，消除 __syncthreads() 开销。
4. 【全面淘汰 3-Pass 分离实现】：严禁在生产中使用独立的 Max Kernel + Sum Kernel + Norm Kernel；全面升级为 Online Softmax 原生融合实现。
5. 【适度启用寄存器缓存 1-Pass 融合】：当行宽 $N \le 4096$ 时，评估线程局部寄存器数组用量，优先启用 1-Pass Fused Softmax，将全局显存访问彻底压到 1 读 1 写。
6. 【多行网格跨步覆盖】：在外层 Block 调度中，使用 Grid-Stride 循环处理 $M$ 维度，确保张量规模弹性缩放时硬件占空比永远饱和。
7. 【SASS 级寄存器溢出排查】：通过 nvcc -Xptxas=-v 检查编译产物，确保引入局部缓存后 Spill stores 严格为 0，防止局部变量溢出至慢速 Local Memory。

8. Ringi 5 点核心速记口诀、自我检验清单与课后深度思考题

8.1 5 点押韵核心速记口诀

规约连续两端折，分支发散全消灭。
Warp 直通五步跃，寄存器里把和接。
未防溢出先减大，三遍扫描显存竭。
动态修正巧相加，一遍融合把速截。
流式缩放乘法借，Attention 革命写新页！

8.2 10 条白板自我检验清单

为什么在 FP16 精度下，注意力得分大于 11.1 时朴素 Softmax 会直接发生指数溢出？
简述 Mark Harris 规约树中，为什么把步长从“从小到大（K0）”改成“从大到小对折（K2）”就能彻底消除 Bank 冲突？
为什么在 Warp 内做规约时不需要显式调用 __syncthreads()？
__shfl_down_sync(0xffffffff, val, 16) 中，第一个参数掩码 0xffffffff 代表什么硬件含义？
传统 Safe Softmax 的“三遍扫描（3-Pass）”分别完成了什么计算？总显存流量是多少？
写出 Online Softmax 单元素在线递推更新公式，并指出动态缩放因子（Rescale Factor）的具体形式。
证明 Online Softmax 的两路状态合并公式满足结合律。
什么是 1-Pass Fused Softmax？它依靠什么硬件介质消除了最后一遍全局显存读取？
在大模型推理的长文本场景下，为什么说 Attention Softmax 是典型的 Memory-Bound 算子？
FlashAttention 是如何利用 Online Softmax 的原理，在不显式保存 $S = QK^T$ 矩阵的前提下完成注意力计算的？

8.3 3 道高阶开放式课后思考题（含极限 Corner Case）

思考题 1：超长行宽的极限规约（ $N > 65536$ ）

当 Softmax 作用在超大词表维度（例如某些多模态模型的词表大小

V = 131072

）时，单 Block 内部的寄存器和共享内存根本无法容纳整行数据。此时 1-Pass Fused Softmax 无法直接生效。请问在系统架构上，应如何设计跨 Block 的两阶段分布式 Online Softmax 算子？如何利用原子操作（Atomic）或跨 Block 协作网格完成全局归约？

思考题 2：数值精度边界 —— 修正因子的极端下溢

在 Online Softmax 中，动态修正因子为

\alpha = e^{m_{\text{old}} - m_{\text{new}}}

。如果新加入的元素极其巨大，使得

m_{\text{old}} - m_{\text{new}} = -100

，在 FP16 下

\alpha

将直接下溢为

0.0

。请从浮点分析角度推导：此时历史累加和

d

被乘成 0，算法在数学上是正确的还是会引发精度灾难？

思考题 3：算子融合前沿 —— FlashAttention-3 的 WGMMA 与 TMA 协同

在最新的 Hopper 架构中，FlashAttention-3 引入了硬件级异步拷贝 TMA 和 Warpgroup GEMM（WGMMA）。请分析：当 TMA 硬件直接把张量从全局显存搬入共享内存时，Online Softmax 的规约与修正逻辑应该由哪个 Warp 组负责？如何排布软流水线（Software Pipelining）以实现乘加计算与 Softmax 缩放的完美重叠（Overlap）？

9. 📚 参考资料与核心源码/经典论文指引

本讲内容与推导过程严格对照并依据本地知识库 AI_BOOK 中的权威一手文献与源码：

并行规约权威奠基专著：
- Mark Harris: Optimizing Parallel Reduction in CUDA, NVIDIA Developer Technology, 2007.
- 本地核心代码解析与压测：参见 10_reduction.md 与 3.1-CUDA Reduce算子优化.md。
Online Softmax 奠基论文与工程实现：
- Maxim Milakov, Natalia Gimelshein: Online normalizer calculation for softmax, NVIDIA Corporation, 2018 (arXiv:1805.02867).
- 本地工程级代码实现与递推详解：参见 5.2-CUDA Online Softmax实现.md 与 LeetCUDA/kernels/interview/base.cuh。
FlashAttention 核心前沿论文：
- Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022.
- 本地 FlashAttention 原理解析：参见 6.1-FlashAttention V1详解.md。

附录：Appendix A — 大厂硬核高频面试题与白板推导（Interview Drill）

题目 1：请在白板上推导 Online Softmax 的单元素增量递推公式与两块合并公式，并证明其满足结合律。

【面试官考察维度】

是否真正理解 Safe Softmax 的数值稳定性物理直觉；
能否独立推导 Online 修正因子 $e^{m_{\text{old}} - m_{\text{new}}}$ ；
能否运用代数证明两块状态合并满足结合律，从而论证 GPU 并行规约的正确性。

【白板标准解答与推导路径】

第一步：写出单元素增量递推公式 设已见前 $k$ 个元素最大值为 $m_k$ ，分母为 $d_k = \sum_{i=1}^k e^{x_i - m_k}$ 。新加入元素 $x_{k+1}$ 时：

m_{k+1} = \max(m_k, x_{k+1})

d_{k+1} = \sum_{i=1}^{k+1} e^{x_i - m_{k+1}} = \left(\sum_{i=1}^k e^{x_i - m_k}\right) e^{m_k - m_{k+1}} + e^{x_{k+1} - m_{k+1}} = d_k \cdot e^{m_k - m_{k+1}} + e^{x_{k+1} - m_{k+1}}

第二步：写出两块独立状态合并公式 设两块数据状态分别为 $(m_1, d_1)$ 和 $(m_2, d_2)$ ：

m_{\text{merged}} = \max(m_1, m_2)

d_{\text{merged}} = d_1 \cdot e^{m_1 - m_{\text{merged}}} + d_2 \cdot e^{m_2 - m_{\text{merged}}}

第三步：证明结合律 $[ (A \oplus B) \oplus C = A \oplus (B \oplus C) ]$ 定义状态合并算子 $\oplus$ ： $(m_1, d_1) \oplus (m_2, d_2) = (m_{12}, d_{12})$ 。易知 $m_{(12)3} = \max(\max(m_1, m_2), m_3) = \max(m_1, m_2, m_3) = M$ 显然满足结合律。再考察分母：

d_{(12)3} = d_{12} \cdot e^{m_{12} - M} + d_3 \cdot e^{m_3 - M} = \left( d_1 e^{m_1 - m_{12}} + d_2 e^{m_2 - m_{12}} \right) e^{m_{12} - M} + d_3 e^{m_3 - M}

指数展开相乘：

d_{(12)3} = d_1 e^{m_1 - M} + d_2 e^{m_2 - M} + d_3 e^{m_3 - M}

同理计算

d_{1(23)}

，展开后完全一致。结合律获证！ 这意味着无论 GPU 的线程树如何分叉折叠，最终结果严格恒等。

题目 2：为什么 `__shfl_down_sync` 可以在 Warp 内 5 次迭代完成 32 线程规约？画出数据流动拓扑图并解释掩码 `0xffffffff` 的含义。

【面试官考察维度】

考查对 GPU 底层 SIMT 指令、Lane 概念以及 Warp 级并行硬件原语的微架构理解。

【白板推导路径】

解释掩码 0xffffffff：
- 掩码是一个 32-bit 无符号整数，每一位对应 Warp 中的一个 Lane（线程 0~31）；
- 0xffffffff（二进制全 1）表示当前 Warp 内的全部 32 个线程都必须参与此条同步 Shuffle 指令；如果某位为 0，代表该线程不参与交换。
推导 5 步折叠拓扑：
- 32 是 2 的 5 次方（ $2^5 = 32$ ）；
- 迭代 1（offset = 16）：线程 $0 \sim 15$ 分别读取线程 $16 \sim 31$ 的寄存器并累加，此时前 16 个线程保存了 16 对和；
- 迭代 2（offset = 8）：线程 $0 \sim 7$ 分别读取线程 $8 \sim 15$ 的数据并累加；
- 迭代 3（offset = 4）：线程 $0 \sim 3$ 累加；
- 迭代 4（offset = 2）：线程 $0 \sim 1$ 累加；
- 迭代 5（offset = 1）：线程 0 读取线程 1 的数据累加。此时线程 0 的寄存器内保存了整个 Warp 32 线程的全部和。
硬件优势：数据全程在 SM 的通用寄存器物理交叉开关（Register Crossbar）上流动，无需访存指令，零延迟，无需显式 __syncthreads()。

题目 3：在行宽 $N = 4096$ 的矩阵 Softmax 中，如何设计 Block 与 Thread 的映射？为什么 1 个 Block 处理 1 行比 1 个 Thread 处理 1 行好？

【面试官考察维度】

考查将数学算法映射到 GPU 网格网格架构时的系统级权衡能力（访存合并度 vs 并行粒度）。

【白板推导路径】

方案 A（1 个 Thread 处理 1 行）：
- 优点：线程内部天然串行，不需要线程间规约和同步；
- 致命缺陷：矩阵在内存中是行优先排布的，相邻线程（处理相邻行）在访问第 $c$ 列时，物理地址相差整整一整行（ $N \times 4$ 字节）！导致全局内存读写完全是非合并访问（Strided Access），有效带宽暴跌 90% 以上；
方案 B（1 个 Block 处理 1 行，黄金方案）：
- 配置 dim3 block(256)，网格 dim3 grid(M)；
- Block 内连续的 256 个线程（变化最快的是 threadIdx.x）同时读取同一行的连续列；
- 天然完美触发全局内存合并访问（Coalesced Access），打满 HBM 物理总线；
- 行内 4096 个元素由 256 个线程分摊（每人处理 16 个），通过片上 Warp Shuffle 极速完成规约；
- 结论：方案 B 胜出，性能高出方案 A 一个数量级以上。

题目 4：FlashAttention 是如何将 Online Softmax 应用到分块注意力计算中的？请写出输出张量 $O$ 的动态缩放更新公式并解释为什么不需要保存注意力权重矩阵 $S$ 。

【面试官考察维度】

大模型 Infra 面试王牌大题，考察从基础算子到前沿系统的知识贯通能力。

【白板推导路径】

分块机制：将长序列按列切分，分批加载 $K_j, V_j$ 进 SRAM；
推导局部与全局更新：设第 $j$ 块计算出的局部得分为 $S_j = Q K_j^T$ ，局部最大值为 $m_j$ ，局部指数为 $P_j = \exp(S_j - m_j)$ ，局部总和为 $l_j = \text{rowsum}(P_j)$ ；维护全局状态：

m_{\text{new}} = \max(m_{\text{old}}, m_j)

d_{\text{new}} = d_{\text{old}} \cdot e^{m_{\text{old}} - m_{\text{new}}} + l_j \cdot e^{m_j - m_{\text{new}}}

输出矩阵 $O$ 的流式更新：在尚未做全局除法前，输出累加量 $O$ 维护的是 $\sum P_i V_i$ 的分子部分：

O_{\text{new}} = O_{\text{old}} \cdot e^{m_{\text{old}} - m_{\text{new}}} + P_j V_j \cdot e^{m_j - m_{\text{new}}}

彻底省去 $S$ 的物理原因：因为 $P_j$ 在片上 SRAM 算出来后，立即与 $V_j$ 相乘并累加进了 $O$ 中！完成累加后，局部矩阵 $S_j$ 和 $P_j$ 的使命彻底终结，可以直接丢弃覆写，完全无需向全局显存写回哪怕一个元素，显存复杂度从 $O(N^2)$ 骤降至 $O(N)$ 。

⌘I

​🏛️ 第23讲：为什么Attention必须先算Max？——从 Naive Reduce 到 Fused Online Softmax 深度解构

​📑 目录导航

​0. Ringi 开场：生产真实现场与痛点冲突

​0.1 真实工程矛盾：为什么长文本一开，Softmax 成了显存吞噬兽？

​0.2 线上事故复盘：某多模态大模型 FP16 数值溢出（NaN 灾难）与三遍扫描往返墙

​0.3 AI Infra 规约与 Softmax 各演进阶段速查表

​1. 规约（Reduction）体系结构：从交错寻址到 Warp Shuffle 的七级跳跃

​1.1 规约操作在 AI 算子中的统治地位

​1.2 Mark Harris 经典七级优化脉络全景解构

​Kernel 0：朴素交错寻址（Interleaved Addressing）

​Kernel 1 & 2：消除发散与连续寻址（Sequential Addressing）

​Kernel 3 & 4：展开最后一个 Warp 与完全循环展开（Loop Unrolling）

​1.3 Warp Shuffle 硬件微架构原语：跨 Lane 寄存器直通网络

核心原语家族：

5 步完成 32 线程 Warp 级规约的极简魔法：

​1.4 Ringi 工程师五问：规约计算视角下的 Shape 与 Cost

​2. 数值稳定性第一性原理：为什么 Softmax 必须做“Safe”保护？

​2.1 IEEE-754 浮点数的物理边界：FP16 与 FP32 的溢出悬崖

​2.2 No Naked Formula 2.0：平移不变性与 Safe Softmax 模型

① 为什么需要算它？

② Mental Model（物理直觉比喻）

③ Tiny Calculator（极简数字手算）

④ Formal Model（标准公式）

⑤ Sanity Check（代数恒等证明）

​2.3 传统 Safe Softmax 的三遍扫描之殇（3-Pass Memory Wall）

​3. Online Softmax 算法推导：如何在一遍扫描中同时求 Max 和 Sum？

​3.1 核心洞察：动态修正因子（Rescale Factor）的代数美学

​3.2 No Naked Formula 2.0：单元素增量递推模型

① 为什么需要算它？

② Mental Model（物理直觉）

③ Tiny Calculator（手算 3 个数字）

④ Formal Model（递推公式）

⑤ Sanity Check（数值安全性）

​3.3 树形规约结合律代数推导（Two-Block Merge）

严格代数推导：

​4. 工业级工程实现：从 2-Pass Online Softmax 到 1-Pass Fused Softmax

​4.1 核心架构：Warp Shuffle 双值规约 + Warp 间共享内存交换

完整的两级规约架构：

​4.2 终极一跃：1-Pass Fused Softmax（寄存器缓存消除最后一次全局读）

​4.3 向量化与多行网格跨步调度（Grid-Stride Multi-Row Parallelism）

​5. FlashAttention 核心前置：Online Softmax 是如何成就大模型注意力革新的？

​5.1 为什么标准 Attention 必须保存庞大的 S=QKTS = QK^TS=QKT 矩阵到显存？

​5.2 FlashAttention-1/2 的核心数学解构：输出张量的动态缩放递推

​5.3 体系结构级洞见：用 SRAM 乘法换取 HBM 流量清零

​6. 动手实战与代码实验室（Minimal Runnable Code）

​实验 1：树形规约演进基准（reduction_evolution_benchmark.cu）

​实验 2：经典 3-Pass Safe Softmax 基准（safe_softmax_3pass_benchmark.cu）

​实验 3：高性能 2-Pass / 1-Pass Online Softmax 算子（online_softmax_benchmark.cu）

​实验 4：端到端吞吐压测与显存流量对比（softmax_benchmark_harness.cu）

​7. Ringi 避坑指南与生产黄金准则

​7.1 避坑表格：7 大常见小白错误理解 vs 大厂 AI Infra 正确认知

​7.2 生产性能工程黄金 Checklist

​8. Ringi 5 点核心速记口诀、自我检验清单与课后深度思考题

​8.1 5 点押韵核心速记口诀

​8.2 10 条白板自我检验清单

​8.3 3 道高阶开放式课后思考题（含极限 Corner Case）

​思考题 1：超长行宽的极限规约（ N>65536N > 65536N>65536 ）

​思考题 2：数值精度边界 —— 修正因子的极端下溢

​思考题 3：算子融合前沿 —— FlashAttention-3 的 WGMMA 与 TMA 协同

​9. 📚 参考资料与核心源码/经典论文指引

​附录：Appendix A — 大厂硬核高频面试题与白板推导（Interview Drill）

​题目 1：请在白板上推导 Online Softmax 的单元素增量递推公式与两块合并公式，并证明其满足结合律。

【面试官考察维度】

【白板标准解答与推导路径】

​题目 2：为什么 __shfl_down_sync 可以在 Warp 内 5 次迭代完成 32 线程规约？画出数据流动拓扑图并解释掩码 0xffffffff 的含义。

【面试官考察维度】

【白板推导路径】

​题目 3：在行宽 N=4096N = 4096N=4096 的矩阵 Softmax 中，如何设计 Block 与 Thread 的映射？为什么 1 个 Block 处理 1 行比 1 个 Thread 处理 1 行好？

【面试官考察维度】

【白板推导路径】

​题目 4：FlashAttention 是如何将 Online Softmax 应用到分块注意力计算中的？请写出输出张量 OOO 的动态缩放更新公式并解释为什么不需要保存注意力权重矩阵 SSS。

【面试官考察维度】

【白板推导路径】

🏛️ 第23讲：为什么Attention必须先算Max？——从 Naive Reduce 到 Fused Online Softmax 深度解构

📑 目录导航

0. Ringi 开场：生产真实现场与痛点冲突

0.1 真实工程矛盾：为什么长文本一开，Softmax 成了显存吞噬兽？

0.2 线上事故复盘：某多模态大模型 FP16 数值溢出（NaN 灾难）与三遍扫描往返墙

0.3 AI Infra 规约与 Softmax 各演进阶段速查表

1. 规约（Reduction）体系结构：从交错寻址到 Warp Shuffle 的七级跳跃

1.1 规约操作在 AI 算子中的统治地位

1.2 Mark Harris 经典七级优化脉络全景解构

Kernel 0：朴素交错寻址（Interleaved Addressing）

Kernel 1 & 2：消除发散与连续寻址（Sequential Addressing）

Kernel 3 & 4：展开最后一个 Warp 与完全循环展开（Loop Unrolling）

1.3 Warp Shuffle 硬件微架构原语：跨 Lane 寄存器直通网络

1.4 Ringi 工程师五问：规约计算视角下的 Shape 与 Cost

2. 数值稳定性第一性原理：为什么 Softmax 必须做“Safe”保护？

2.1 IEEE-754 浮点数的物理边界：FP16 与 FP32 的溢出悬崖

2.2 No Naked Formula 2.0：平移不变性与 Safe Softmax 模型

2.3 传统 Safe Softmax 的三遍扫描之殇（3-Pass Memory Wall）

3. Online Softmax 算法推导：如何在一遍扫描中同时求 Max 和 Sum？

3.1 核心洞察：动态修正因子（Rescale Factor）的代数美学

3.2 No Naked Formula 2.0：单元素增量递推模型

3.3 树形规约结合律代数推导（Two-Block Merge）

4. 工业级工程实现：从 2-Pass Online Softmax 到 1-Pass Fused Softmax

4.1 核心架构：Warp Shuffle 双值规约 + Warp 间共享内存交换

4.2 终极一跃：1-Pass Fused Softmax（寄存器缓存消除最后一次全局读）

4.3 向量化与多行网格跨步调度（Grid-Stride Multi-Row Parallelism）

5. FlashAttention 核心前置：Online Softmax 是如何成就大模型注意力革新的？

5.1 为什么标准 Attention 必须保存庞大的 $S = QK^T$ 矩阵到显存？

5.2 FlashAttention-1/2 的核心数学解构：输出张量的动态缩放递推

5.3 体系结构级洞见：用 SRAM 乘法换取 HBM 流量清零

6. 动手实战与代码实验室（Minimal Runnable Code）

实验 1：树形规约演进基准（`reduction_evolution_benchmark.cu`）

实验 2：经典 3-Pass Safe Softmax 基准（`safe_softmax_3pass_benchmark.cu`）

实验 3：高性能 2-Pass / 1-Pass Online Softmax 算子（`online_softmax_benchmark.cu`）

实验 4：端到端吞吐压测与显存流量对比（`softmax_benchmark_harness.cu`）

7. Ringi 避坑指南与生产黄金准则

7.1 避坑表格：7 大常见小白错误理解 vs 大厂 AI Infra 正确认知

7.2 生产性能工程黄金 Checklist

8. Ringi 5 点核心速记口诀、自我检验清单与课后深度思考题

8.1 5 点押韵核心速记口诀

8.2 10 条白板自我检验清单

8.3 3 道高阶开放式课后思考题（含极限 Corner Case）

思考题 1：超长行宽的极限规约（ $N > 65536$ ）

思考题 2：数值精度边界 —— 修正因子的极端下溢

思考题 3：算子融合前沿 —— FlashAttention-3 的 WGMMA 与 TMA 协同

9. 📚 参考资料与核心源码/经典论文指引

附录：Appendix A — 大厂硬核高频面试题与白板推导（Interview Drill）

题目 1：请在白板上推导 Online Softmax 的单元素增量递推公式与两块合并公式，并证明其满足结合律。

题目 2：为什么 `__shfl_down_sync` 可以在 Warp 内 5 次迭代完成 32 线程规约？画出数据流动拓扑图并解释掩码 `0xffffffff` 的含义。

题目 3：在行宽 $N = 4096$ 的矩阵 Softmax 中，如何设计 Block 与 Thread 的映射？为什么 1 个 Block 处理 1 行比 1 个 Thread 处理 1 行好？

题目 4：FlashAttention 是如何将 Online Softmax 应用到分块注意力计算中的？请写出输出张量 $O$ 的动态缩放更新公式并解释为什么不需要保存注意力权重矩阵 $S$ 。