第 6 课：稳定 Softmax 与朴素复杂度

本课目标与完成标准

学完后你应能：写出减最大值的稳定 softmax，并准确分析朴素 kernel 的 O(N²) 行扫描。

通过必须同时满足：

• 独立补齐本课唯一的代码填空题，并通过给定检查；
• 三个问答题均说明因果链，而不是只报术语；
• 能指出至少一个正确性边界和一个性能取舍；
• 总分不低于 8/10，且没有一票否决级概念错误。

前置关系

• 课程前置：C/C++、线性代数、基本并行编程
• 本课在路线中的作用：Softmax 把 logits 变成正且和为 1 的分布；指数使它对溢出和下溢敏感。

核心心智模型

1. 它是什么，解决什么问题

Softmax 算子 是分类任务和自注意力机制（Self-Attention）的核心基础。它将未归一化的实数得分向量 $\mathbf{x} = [x_1, x_2, \dots, x_N]$ 映射为非负且归一化的概率分布向量 $\mathbf{y}$： $$y_i = \frac{e^{x_i}}{\sum_{j=1}^N e^{x_j}}, \quad \text{其中 } y_i > 0, \ \sum_{i=1}^N y_i = 1$$

面对的数值危险（溢出与下溢）： - 上溢（Overflow）：在 IEEE 754 单精度浮点数（FP32）中，最大可表示正数为 $\approx 3.4 \times 10^{38}$。由于指数函数的爆炸性增长，当输入分量 $x_i > \ln(3.4 \times 10^{38}) \approx 88.72$ 时，expf(x_i) 就会直接溢出为正无穷 +inf。若分母中包含 +inf，或者分子分母均为 +inf，浮点运算将直接生成不可逆的 NaN； - 下溢（Underflow）：当输入分量极小（如 $x_i < -87.33$）时，expf(x_i) 下溢归零。若所有分量皆下溢，分母归零，触发除以零异常； - 解决方案（Safe Softmax，图 6-1）：

数值稳定 Softmax 平移

原始输入 x_i

找全局最大值 m = max(x)

安全平移: x_i - m <= 0

指数计算: exp(x_i - m)
∈ (0, 1] (杜绝浮点溢出)

求和分母: sum(exp)

归一化输出: y_i = exp / sum

图 6-1：数值安全 Softmax 减最大值平移与归一化机制

利用恒等式 $\frac{e^{x_i}}{\sum_j e^{x_j}} = \frac{e^{x_i - m}}{\sum_j e^{x_j - m}}$，其中 $m = \max_{1 \le k \le N} x_k$。平移后所有项的指数 $x_i - m \le 0$，保证了 $e^{x_i - m} \in (0, 1]$，杜绝上溢；且至少有一项 $e^{m - m} = e^0 = 1$，分母恒满足 $\sum_j e^{x_j - m} \ge 1$，杜绝除零。

2. 它如何工作（结合图 6-1 与图 6-2）

本课的 朴素实现（Naive Softmax） 采用直觉最简单但性能最差的映射策略——One Thread per Output Element（一个线程负责一个输出点）：

1. 线程映射与寻址：
   • 构造 2D Grid 和 2D Block，线程坐标映射为二维索引 $(row, col)$；
   • 只要索引有效（$row < M$ 且 $col < N$），该线程独立负责计算矩阵中 $(row, col)$ 位置的 Softmax 输出。
2. 三步独立行扫描（图 6-2）：

全局显存 (HBM)

寄存器

Warp 归约

Pass 1: 扫描整行找全局最大值 m

1读入 x_i

2折半找最大值 m

Pass 2: 重新读取整行，计算 exp(x_i - m) 并求和 sum

3重新读入 x_i

4折半求和 sum

Pass 3: 第三次读取整行，除以 sum 并写回输出

5第三次读入 x_i

6写回 y_i = exp(x_i - m) / sum (3 次全行往返，带宽严重浪费)

图 6-2：朴素 Softmax 三遍访存流水线与内存开销瓶颈

• 第一遍扫描（找最大值）：线程沿着自己的行，在 Global Memory 中串行遍历 $i = 0 \dots N-1$，求得该行的全局最大值 $m = \max_i x_i$；
• 第二遍扫描（求和）：线程再次沿自己的行串行遍历 $i = 0 \dots N-1$，在寄存器中累加 $d = \sum_i \exp(x_i - m)$；
• 第三步（计算写回）：线程计算自己专属列号 $col$ 的分子 $\exp(x_{col} - m)$，除以分母 $d$，并将概率值写回输出矩阵 output[row, col]。

3. 正确性条件与常见误区

1. 数学平移的严格一致性：
   • 分子与分母必须减去同一个基准最大值 $m$。如果分子减去了局部最大值，或者只在分母减去了 $m$ 而分子保持 $x_i$，则打破了恒等变换条件： $$\frac{e^{x_i}}{\sum_j e^{x_j - m}} = e^m \cdot \frac{e^{x_i}}{\sum_j e^{x_j}} \neq S(x_i)$$ 计算结果将整体被放大 $e^m$ 倍，概率和不再为 1，直接破坏概率分布不变量；
2. Padding 与 Mask 边界处理：
   • 在 Attention 中，被 Mask 掉的位置通常置为 $-\infty$（在 FP32 中常设为 -1e9 或 -FLT_MAX）。求最大值时若整行全被 Mask，需要明确兜底机制；且在求分母 $\sum$ 时必须严格跳过 Padding 位置，否则多余的无效元素会分摊正常 Token 的概率权重。

4. 性能与工程取舍（为什么 Naive 实现是算力灾难）

1. 计算复杂度与访存量爆炸 $O(N^2)$：
   • 设单行长度为 $N$。一行共有 $N$ 个输出元素，由同一行上的 $N$ 个线程分别负责；
   • 每个线程各自遍历该行 2 遍，每次读取 $N$ 个浮点数。因此，整行被全体线程累计重复读取了 $2N \times N = 2N^2$ 次。
   • 对于 $N = 4096$ 的典型序列长度，一行仅有 16 KB 数据，但该朴素 Kernel 会对这一行发起 $2 \times 4096^2 \times 4 \text{ B} \approx 134 \text{ MB}$ 的重复读取请求。这些请求大多命中 L1（同一行的 16 个 lane 读同一地址时还会广播），所以主要浪费的是 load 指令、L1 带宽和发射槽，而不是 DRAM 带宽；
2. 从 Naive 到 Block 协作：
   • 朴素 Kernel 极度简单，无需任何线程间同步或 Shared Memory，适合建立基线概念；
   • 但在工程实践中，必须将调度粒度重构为 One Block per Row（一个 ThreadBlock 协作处理一行）：Block 内所有线程协同遍历一行，通过 Warp Shuffle 和 Shared Memory 归约求出 $m$ 和 $d$。整行数据仅被全局扫描 2~3 次，总访存复杂度从 $O(N^2)$ 骤降至 $O(N)$。

具体演示：数值安全对比实测

考虑输入张量包含两个非常大的得分：$x = [1000.0, 1001.0]$：

• 非安全直接计算：
  • $\exp(1000.0) \to +\infty$（由于超出了 FP32 最大浮点表示范围 $3.4 \times 10^{38}$）；
  • $\exp(1001.0) \to +\infty$；
  • 分母 $d = +\infty + (+\infty) = +\infty$；
  • 输出 $y_1 = \frac{+\infty}{+\infty} = \text{NaN}$。整个网络训练当场报废。
• 减最大值安全计算（Safe Softmax）：
  • 最大值 $m = \max(1000.0, 1001.0) = 1001.0$；
  • 平移后指数差分向量：$x - m = [-1.0, 0.0]$；
  • 计算指数项：$\exp(-1.0) \approx 0.367879, \ \exp(0.0) = 1.0$；
  • 分母 $d = 0.367879 + 1.0 = 1.367879$；
  • 最终概率输出：$y = [0.26894, 0.73106]$。数值丝毫不溢出，且概率总和精确等于 1.0。

实践任务：唯一代码填空题

补齐最终稳定输出。

规则：只能修改 TODO/______ 所在位置；不要删除断言或放宽误差。代码注释说明了每个边界条件。

%%writefile /tmp/06_softmax_naive.cu
#include <float.h>
#include <cuda_runtime.h>

// Naive Softmax: 对 [M, N] 的每一行做 softmax。
//
// 这个版本保留“一个线程负责一个输出元素”的写法，
// 适合面试里先讲清楚 softmax 的数学公式和它为什么低效。
//
// output[row, col] = exp(input[row, col] - max(row)) / sum(exp(input[row, i) - max(row)))
//
// 低效点：
// 对同一行来说，max(row) 和 sum(row) 是所有 col 共享的，
// 但这里每个输出元素线程都会重新扫完整行求 max 和 sum。
// 所以每行复杂度接近 O(N^2)，优化版应改成一个 block 合作处理一行。
__global__ void softmax_naive_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int M,
    int N
) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int row = blockIdx.y * blockDim.y + threadIdx.y;

    if (row < M && col < N) {
        float max_val = -FLT_MAX;

        // 第一次扫整行，求最大值，避免 exp 溢出。
        for (int i = 0; i < N; ++i) {
            max_val = fmaxf(max_val, input[row * N + i]);
        }

        float sum_exp = 0.0f;

        // 第二次扫整行，求 denominator。
        for (int i = 0; i < N; ++i) {
            sum_exp += expf(input[row * N + i] - max_val);
        }

        // 当前线程只写自己负责的一个 output[row, col]。
        output[row * N + col] = ______;  // TODO: 稳定归一化
    }
}

void launch_softmax_naive(
    const float* input,
    float* output,
    int M,
    int N,
    cudaStream_t stream
) {
    dim3 block(16, 16);
    dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);
    softmax_naive_kernel<<<grid, block, 0, stream>>>(input, output, M, N);
}

检查方法

有 CUDA 环境时执行 nvcc -std=c++17 -c /tmp/06_softmax_naive.cu -o /tmp/06_softmax_naive.cu.o；无 CUDA 环境时只做静态审查并登记待验证。

提交时请给出：补齐后的代码、实际运行输出（环境不可用时注明“仅静态审查”）以及对失败用例的解释。

Q1

不要背定义：请从输入、状态变化和输出三个阶段解释“稳定 Softmax 与朴素复杂度”的工作机制。

你的答案：

Q2

只在分母减 max、分子不减会怎样？

你的答案：

Q3

为什么这个 naive kernel 在 N 增大时比“一 block 一行”恶化得更快？

你的答案：

评分与通过规则

• 代码 4 分：正常输入 2 分，边界输入 1 分，解释实现 1 分；
• Q1～Q3 各 2 分；
• 一票否决：结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。

需要提示时按四级机制请求：概念区域 → 具体方向 → 关键局部 → 完整答案。

Tip💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

先完成题目再核对。即使代码一致，也要能解释关键步骤，并尝试更换一个输入规模。

%%writefile /tmp/06_softmax_naive.cu
#include <float.h>
#include <cuda_runtime.h>

// Naive Softmax: 对 [M, N] 的每一行做 softmax。
//
// 这个版本保留“一个线程负责一个输出元素”的写法，
// 适合面试里先讲清楚 softmax 的数学公式和它为什么低效。
//
// output[row, col] = exp(input[row, col] - max(row)) / sum(exp(input[row, i) - max(row)))
//
// 低效点：
// 对同一行来说，max(row) 和 sum(row) 是所有 col 共享的，
// 但这里每个输出元素线程都会重新扫完整行求 max 和 sum。
// 所以每行复杂度接近 O(N^2)，优化版应改成一个 block 合作处理一行。
__global__ void softmax_naive_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int M,
    int N
) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int row = blockIdx.y * blockDim.y + threadIdx.y;

    if (row < M && col < N) {
        float max_val = -FLT_MAX;

        // 第一次扫整行，求最大值，避免 exp 溢出。
        for (int i = 0; i < N; ++i) {
            max_val = fmaxf(max_val, input[row * N + i]);
        }

        float sum_exp = 0.0f;

        // 第二次扫整行，求 denominator。
        for (int i = 0; i < N; ++i) {
            sum_exp += expf(input[row * N + i] - max_val);
        }

        // 当前线程只写自己负责的一个 output[row, col]。
        output[row * N + col] = expf(input[row * N + col] - max_val) / sum_exp;
    }
}

void launch_softmax_naive(
    const float* input,
    float* output,
    int M,
    int N,
    cudaStream_t stream
) {
    dim3 block(16, 16);
    dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);
    softmax_naive_kernel<<<grid, block, 0, stream>>>(input, output, M, N);
}

Q1 参考答案

1. 输入阶段（Input & Thread Mapping）：
   • 启动二维网格与二维线程块（如 dim3 block(16, 16)），每个激活的线程依据全局坐标 $(row, col)$ 对应输出矩阵的一个特定元素位置。输入为未归一化的连续实数矩阵；
2. 状态变化阶段（Two-Pass Row Traversal & Safe Exponentiation）：
   • 求最大值（Pass 1）：线程串行遍历 Global Memory 中对应行的全部 $N$ 个输入，维护寄存器局部变量 max_val，提取该行全局最大得分 $m = \max_i x_i$；
   • 指数累加（Pass 2）：线程再次串行遍历该行全部 $N$ 个输入，对每个元素做负偏置指数变换 $\exp(x_i - m)$，在寄存器中累加归一化分母 $d = \sum_{i=0}^{N-1} \exp(x_i - m)$。这一步保证了所有指数项输入非正，避免了单精度浮点数溢出为 +inf；
3. 输出阶段（Pointwise Normalization & Store）：
   • 线程读取自己负责的目标位置原始输入 $x_{row, col}$，计算分子 $\exp(x_{row, col} - m)$ 并除以分母 $d$，得到合法概率值 $y_{row, col} \in (0, 1]$，单次写回显存 output[row * N + col]。

Q2 参考答案

若只在分母减去最大值 $m$，而分子保持原始指数 $\exp(x_i)$，会引发严重的数学错误与数值崩溃：

1. 打破同乘同除恒等变换：
   • Softmax 的数学本质是比例归一化：$\frac{e^{x_i}}{\sum_j e^{x_j}} = \frac{e^{x_i} \cdot e^{-m}}{\sum_j e^{x_j} \cdot e^{-m}} = \frac{e^{x_i - m}}{\sum_j e^{x_j - m}}$；
   • 若分子不乘 $e^{-m}$，则实际计算的表达式变为： $$\hat{y}_i = \frac{e^{x_i}}{\sum_j e^{x_j - m}} = \frac{e^{x_i}}{e^{-m} \sum_j e^{x_j}} = e^m \cdot \frac{e^{x_i}}{\sum_j e^{x_j}} = e^m \cdot y_i$$ 整个输出概率分布被无故放大了 $e^m$ 倍。此时输出元素总和不再是 1，而是 $\sum_i \hat{y}_i = e^m$。这破坏了概率分布的基本公理（和为 1）；
2. 无法解决分子的数值溢出：
   • 减最大值的核心目的之一就是防止大数指数上溢。如果分子仍计算 $\exp(x_i)$，当输入得分 $x_i > 88.72$ 时，分子依然会先直接溢出为 +inf。哪怕分母被约束在有限值，+inf / d 仍然是 +inf，数值溢出缺陷完全没有被解决。

Q3 参考答案

• 工作量 $O(N^2)$ vs. $O(N)$：Naive Kernel 中一行的 $N$ 个线程各自遍历整行两遍，整行被读取 $2N^2$ 次，expf 与比较也重复 $N$ 次；“一 Block 一行”只读两遍，共 $2N$ 次。$N$ 从 1024 增到 8192 时，前者的工作量增加 64 倍，后者只增加 8 倍。
• 瓶颈在重复的 L1/LSU 流量，不在 DRAM：block(16,16) 下一个 Warp 覆盖 2 行，循环中同一行的 16 个 lane 读同一地址（广播），一行 16 KB（$N=4096$）可以留在 L1 中，所以 Cache Miss 率并不高；代价是 $O(N^2)$ 条 load 指令和对应的 L1 带宽、发射槽，以及重复的 expf 计算。

参考资料

• CUDA Programming Guide
• CUDA Best Practices Guide

资料用于建立事实基线；面试回答仍需用自己的语言组织。

← 上一课
第 5 课：LayerNorm 下一课 →
第 7 课：Online Softmax 与…