第 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）：

flowchart LR
    subgraph SafeSoftmax["数值稳定 Softmax 平移"]
        X["原始输入 x_i"] --> Max["找全局最大值 m = max(x)"]
        Max --> Shift["安全平移: x_i - m <= 0"]
        Shift --> Exp["指数计算: exp(x_i - m) ∈ (0, 1] (彻底杜绝浮点溢出)"]
        Exp --> Sum["求和分母: sum(exp)"]
        Sum --> Div["归一化输出: y_i = exp / sum"]
    end

图 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（一个线程负责一个输出点）：

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

sequenceDiagram
    autonumber
    participant GM as 全局显存 (HBM)
    participant Reg as 寄存器
    participant W as Warp 归约

    Note over GM,W: Pass 1: 扫描整行找全局最大值 m
    GM->>Reg: 读入 x_i
    Reg->>W: 折半找最大值 m
    Note over GM,W: Pass 2: 重新读取整行，计算 exp(x_i - m) 并求和 sum
    GM->>Reg: 重新读入 x_i
    Reg->>W: 折半求和 sum
    Note over GM,W: Pass 3: 第三次读取整行，除以 sum 并写回输出
    GM->>Reg: 第三次读入 x_i
    Reg->>GM: 写回 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. 正确性条件与常见误区

数学平移的严格一致性：
- 分子与分母必须减去同一个基准最大值 $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，直接破坏概率分布不变量；
Padding 与 Mask 边界处理：
- 在 Attention 中，被 Mask 掉的位置通常置为 $-\infty$ （在 FP32 中常设为 -1e9 或 -FLT_MAX）。求最大值时若整行全被 Mask，需要明确兜底机制；且在求分母 $\sum$ 时必须严格跳过 Padding 位置，否则多余的无效元素会分摊正常 Token 的概率权重。

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

计算复杂度与访存量爆炸 $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/L2 缓存彻底抖动击穿，带宽完全锁死在重复读取上；
从 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/______ 所在位置；不要删除断言或放宽误差。代码注释说明了每个边界条件。

Code

%%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 分；
一票否决：结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。

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

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

参考答案（仅 answer 分支）

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

Code

%%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 参考答案

输入阶段（Input & Thread Mapping）：
- 启动二维网格与二维线程块（如 dim3 block(16, 16)），每个激活的线程依据全局坐标 $(row, col)$ 对应输出矩阵的一个特定元素位置。输入为未归一化的连续实数矩阵；
状态变化阶段（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；
输出阶段（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)$ ，会引发严重的数学错误与数值崩溃：

打破同乘同除恒等变换：
- 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）；
无法解决分子的数值溢出：
- 减最大值的核心目的之一就是防止大数指数上溢。如果分子仍计算 $\exp(x_i)$ ，当输入得分 $x_i > 88.72$ 时，分子依然会先直接溢出为 +inf。哪怕分母被约束在有限值，+inf / d 仍然是 +inf，数值溢出缺陷完全没有被解决。

Q3 参考答案

Naive Kernel 在 $N$ 增大时性能恶化极其剧烈的本质原因在于计算与访存复杂度的维度级放大以及硬件缓存体系的彻底失效：

工作量与访存量复杂度差距： $O(N^2)$ vs. $O(N)$ ：
- 在 Naive Kernel 中，一行有 $N$ 个输出，由 $N$ 个独立线程分别计算。每个线程各自执行两次循环遍历整行（一次求 max，一次求 sum），整行被重复读取了 $2N \times N = 2N^2$ 次浮点数；
- 在“一 Block 负责一行”的协作模型中，Block 内全部线程通过协同跨步与共享内存归约处理该行，求 max 遍历一次（总计 $N$ 次读），求 sum 遍历一次（总计 $N$ 次读），整行全局访存复杂度严格为 $O(N)$ 。当 $N$ 从 1024 增加到 8192 时，协作模型的访存量增加 8 倍，而 Naive Kernel 的访存量暴增 64 倍！
Cache 抖动击穿与指令冗余：
- 当 $N$ 较小时，行数据可能暂时常驻在 SM 的 L1/L2 缓存中；但随着 $N$ 增大，一行的数据量超出 L1 缓存容量，同一行上的 $N$ 个线程互相竞争抢占缓存，导致 Cache Miss 率迅速飙升至接近 100%；
- 此外，Naive Kernel 中每个线程都重复计算了一模一样的 $m$ 与 $d$ ，浮点比较指令和 expf 指令被无谓地重复执行了 $N$ 次，造成极其昂贵的算力浪费，迅速将 GPU 显存总线和指令发射器彻底打瘫。

参考资料

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

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

--- title: "第 6 课：稳定 Softmax 与朴素复杂度" jupyter: python3 --- ::: {.callout-note appearance="simple"} **参考答案版**：包含完整代码实现与问答题深度参考解析。建议先独立完成练习题后再展开核对。 ::: ## 本课目标与完成标准学完后你应能：写出减最大值的稳定 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）**： ```{mermaid} flowchart LR subgraph SafeSoftmax["数值稳定 Softmax 平移"] X["原始输入 x_i"] --> Max["找全局最大值 m = max(x)"] Max --> Shift["安全平移: x_i - m <= 0"] Shift --> Exp["指数计算: exp(x_i - m) ∈ (0, 1] (彻底杜绝浮点溢出)"] Exp --> Sum["求和分母: sum(exp)"] Sum --> Div["归一化输出: y_i = exp / sum"] end ``` 图 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）**： ```{mermaid} sequenceDiagram autonumber participant GM as 全局显存 (HBM) participant Reg as 寄存器 participant W as Warp 归约 Note over GM,W: Pass 1: 扫描整行找全局最大值 m GM->>Reg: 读入 x_i Reg->>W: 折半找最大值 m Note over GM,W: Pass 2: 重新读取整行，计算 exp(x_i - m) 并求和 sum GM->>Reg: 重新读入 x_i Reg->>W: 折半求和 sum Note over GM,W: Pass 3: 第三次读取整行，除以 sum 并写回输出 GM->>Reg: 第三次读入 x_i Reg->>GM: 写回 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/L2 缓存彻底抖动击穿，带宽完全锁死在重复读取上； 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`/`______` 所在位置；不要删除断言或放宽误差。代码注释说明了每个边界条件。 ```{python} #| course_role: exercise %%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 分； - 一票否决：结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。需要提示时按四级机制请求：概念区域 → 具体方向 → 关键局部 → 完整答案。 ::: {.callout-tip collapse="true" title="💡 点击展开参考答案与详细代码实现" icon="false"} ## 参考答案（仅 answer 分支）先完成题目再核对。即使代码一致，也要能解释关键步骤，并尝试更换一个输入规模。 ```{python} #| course_role: solution %%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 参考答案 Naive Kernel 在 $N$ 增大时性能恶化极其剧烈的本质原因在于**计算与访存复杂度的维度级放大**以及**硬件缓存体系的彻底失效**： 1. **工作量与访存量复杂度差距：$O(N^2)$ vs. $O(N)$**： - 在 Naive Kernel 中，一行有 $N$ 个输出，由 $N$ 个独立线程分别计算。每个线程各自执行两次循环遍历整行（一次求 max，一次求 sum），整行被重复读取了 $2N \times N = 2N^2$ 次浮点数； - 在“一 Block 负责一行”的协作模型中，Block 内全部线程通过协同跨步与共享内存归约处理该行，求 max 遍历一次（总计 $N$ 次读），求 sum 遍历一次（总计 $N$ 次读），整行全局访存复杂度严格为 $O(N)$。当 $N$ 从 1024 增加到 8192 时，协作模型的访存量增加 8 倍，而 Naive Kernel 的访存量暴增 64 倍！ 2. **Cache 抖动击穿与指令冗余**： - 当 $N$ 较小时，行数据可能暂时常驻在 SM 的 L1/L2 缓存中；但随着 $N$ 增大，一行的数据量超出 L1 缓存容量，同一行上的 $N$ 个线程互相竞争抢占缓存，导致 Cache Miss 率迅速飙升至接近 100%； - 此外，Naive Kernel 中每个线程都重复计算了一模一样的 $m$ 与 $d$，浮点比较指令和 `expf` 指令被无谓地重复执行了 $N$ 次，造成极其昂贵的算力浪费，迅速将 GPU 显存总线和指令发射器彻底打瘫。 ## 参考资料 - [CUDA Programming Guide](https://docs.nvidia.com/cuda/cuda-programming-guide/) - [CUDA Best Practices Guide](https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/) 资料用于建立事实基线；面试回答仍需用自己的语言组织。 ::: ::: {.course-pagination} [**← 上一课** 第 5 课：LayerNorm](lesson05_layernorm.qmd){.pagination-link .pagination-prev} [**下一课 →** 第 7 课：Online Softmax 与...](lesson07_online_softmax.qmd){.pagination-link .pagination-next} :::