两个融合算子：RMSNorm 与 online softmax

一、第 10 章练的是归约，这里练的是「归约之后呢」↑

第 10 章的 R1–R6 只解决一个问题：把一个长数组变成一个标量。 归约就是终点，用 atomicAdd 写出结果就结束。

在真实算子里，归约几乎从来不是终点。 RMSNorm 先求一行的平方和，然后立刻用这个和回头缩放这一行的每个元素； softmax 先求一行的 max 和 exp 和，然后回头对每个元素做除法。 归约只是中间步骤，「回头」才是花钱的地方。

二、它们在 Qwen2.5-7B 里的真实位置↑

模型参数取自 Qwen2.5-7B 结构拆解： hidden 3584 · 28 层 · 28 Q 头 / 4 KV 头 · head_dim 128 · ffn 18944 · vocab 152064。

在 7B 的 decode 里，这两个算子都不是瓶颈。 decode 一个 token 时，57 次 RMSNorm 一共搬 1.6 MB， 按 1008 GB/s 算要 1.6 微秒。同一步要把 15 GB 权重读一遍，要 15 毫秒。 RMSNorm 只占万分之一。

28 头 × 4096 × 4096 × 4 字节 = 1.88 GB   // 每层，fp32

// 朴素三趟 softmax 要搬 16N/行：
16 × 4096 × (28 × 4096) = 7.52 GB  → 按 1008 GB/s 算 7.46 ms  // 每层
× 28 层 = 209 ms

// 就算用本文最好的 s3（8N/行）：
3.76 GB → 3.73 ms/层 × 28 = 104 ms

三、评价指标换成 MBU↑

MBU = 实测带宽 ÷ 理论峰值带宽 = (搬运字节数 ÷ 耗时) ÷ 1008 GB/s

// 4090 的峰值：内存时钟 10501 MHz × 2（DDR 双沿）× 384 bit ÷ 8 = 1008 GB/s
// 和 sgemm_all.cu 里算 roofline 用的是同一个数，口径一致

一、RMSNorm：12N 还是 8N↑

y = x · rsqrt(mean(x²) + eps) · w        // N = 3584

一句话：把一行（一个 token 的 3584 维向量）整体缩放到固定长度，再逐维乘一个学出来的增益。 公式里只有三个动作：求均方、取倒数平方根、逐维相乘。

今天的主流模型（Llama、Qwen、Mistral）全部用 RMSNorm。 砍掉「减均值」来自 Zhang & Sennrich 2019： 效果几乎不降，还省掉一个统计量和一遍减法。

权重 w 只有 4N = 14 KB，所有行共用， 一直待在 L1/L2 里，不计入显存流量。 N1 还要把长度为 M 的 rstd 数组写一次、读一次， 相对 12NM 可以忽略。

二、softmax：16N / 12N / 8N↑

softmax(x)_i = exp(x_i − m) / Σ_j exp(x_j − m)      // m = max(x)，N = 上下文长度

一句话：把一行任意实数变成一个概率分布：全为正、和为 1、大小顺序不变。 注意力里这一行是「这个 query 对 N 个 key 的打分」； softmax 后变成「分给每个 key 多少注意力」。

S1 → S2 省掉的那一趟靠数学改写，不靠代码技巧。 朴素做法要看完整行才知道 max，才能开始算 exp 和； online 递推让两者同时推进。 推导在 §5。

三、融合 ≠ 省带宽↑

真正省带宽的是 N2 → N3：把 x 留在寄存器里。 工程讨论里「融合」一词用得很松： 有人指合并 kernel，有人指数据留在片上。 这两件事的收益差一个数量级。 N1/N2/N3 三级并排跑一次，就能把它们分清。

一、down 还是 xor↑

// 第 10 章 R6：结果只汇到 0 号 lane，其余 lane 的值是垃圾
for (int off = 16; off > 0; off >>= 1)
    sum += __shfl_down_sync(0xffffffffu, sum, off);
if (lane == 0) atomicAdd(output, sum);        // 只有它要用

// 这里：256 个线程每一个都要拿归一化因子去缩放自己手上的元素
for (int off = 16; off > 0; off >>= 1)
    v += __shfl_xor_sync(0xffffffffu, v, off);
// 五轮之后每个 lane 手里都是全和 —— 少一个线程知道都不行

为什么 xor 一定是「两两互读」

因为异或是对合运算：(lane ^ off) ^ off == lane。 我算出伙伴是 p，p 按同一规则算出的伙伴正好是我。 所以「我读你、你读我」天然成对，一条指令完成一次双向交换。

// xor：对合，成对
lane 2 的伙伴 = 2 ^ 4 = 6
lane 6 的伙伴 = 6 ^ 4 = 2   // 正好读回来，互为伙伴

// down：加法不对合，单向
lane 2 读 lane 2 + 4 = 6
lane 6 读 lane 6 + 4 = 10  // 不是 2，而是往更远处要 —— 所以是漏斗，不是交换

两个版本用的是同一组数、同样三轮、同样三条 shuffle 指令。 区别在于：down 版里 lane 4–7 从第一轮起就向「warp 外面」要数 （越界时硬件返回自己的值），所以它们累加出的是没有意义的部分和。 只有 lane 0 那一格两边都是 31。这就是 R6 只敢让 lane 0 做 atomicAdd 的原因。

五轮为什么正好够：off 就是在逐位翻 lane 号

off 取 16、8、4、2、1，二进制正好是 10000、01000、 00100、00010、00001，每轮翻 lane 号的一个 bit。 lane 号有 5 个 bit。翻完这 5 个 bit，就能从任意一个 lane 走到任意另一个 lane。 以 lane 5（00101）为例：

xor 版不多花钱。 两版都是五轮、每轮一条 shuffle。 xor 版只把「漏斗」换成「全交换」，代价相同，还省掉了之后再广播一次。 所以凡是归约后所有线程都要用结果的场合，直接用 xor。

二、block 级两段式归约↑

BLOCK = 256 = 8 个 warp。 warp 内部用 shuffle 归约，跨 warp 只能走 shared memory。

__inline__ __device__ float blockReduceSum(float v, float *smem) {
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;

    v = warpReduceSum(v);                      // ① 每个 warp 各归各的
    if (lane == 0) smem[warp] = v;             // ② 8 个结果落进 shared
    __syncthreads();                           // ③ 等 8 个都写完

    v = (threadIdx.x < (BLOCK >> 5)) ? smem[lane] : 0.0f;
    if (warp == 0) v = warpReduceSum(v);       // ④ 0 号 warp 把 8 个归成 1 个

    if (threadIdx.x == 0) smem[0] = v;         // ⑤ 放回 shared
    __syncthreads();                           // ⑥ 等它写完
    return smem[0];                            // ⑦ 全 block 一起读走
}

第 ④ 步的 : 0.0f 是加法的幺元，不是随便填的。 8 个 warp 只用到前 8 个 lane，其余 24 个 lane 必须填一个 「参与运算但不改变结果」的值：求和用 0，求最大值用 -FLT_MAX， 求乘积用 1。填错幺元是这类代码里最隐蔽的错误。

三、联合归约：把 (m, s) 当成一个整体↑

online softmax 要同时归约两个量，而且两者不能各归各的： s 的合并方式依赖 m。

__inline__ __device__ void warpReduceMaxSum(float &m, float &s) {
    for (int off = 16; off > 0; off >>= 1) {
        float m_o = __shfl_xor_sync(0xffffffffu, m, off);   // 伙伴的 max
        float s_o = __shfl_xor_sync(0xffffffffu, s, off);   // 伙伴的 sum
        float m_n = fmaxf(m, m_o);                          // 新的公共基准
        s = s * __expf(m - m_n) + s_o * __expf(m_o - m_n);
        m = m_n;
    }
}

两边都先换算到新基准 m_n，再相加。 max 较小的一边，exp 因子小于 1，会把它的和缩小； 两边 max 相等时，两个 exp 都是 1，退化成普通加法。

一、循环体那一行算什么

每个 lane 手里的 s 不是裸和，而是「以我自己的 m 为基准」的和： s = Σ exp(x_j − m)。 我和伙伴的 m 通常不同， 所以两个 s 的量纲不同，不能直接相加。 必须先把两边换算到同一个新基准 m_n = max(m, m_o)，再相加。那一行做的就是这件事。

二、一轮只碰一个伙伴，为什么能覆盖全 warp

m = -INFINITY, s = 0        // 没分到元素的线程
m_n = max(-inf, -inf) = -inf
m - m_n = (-inf) - (-inf) = NaN
s = 0 × NaN + 0 × NaN = NaN   // 然后 NaN 顺着归约树污染整行

// 换成 -FLT_MAX：
m - m_n = (-FLT_MAX) - (-FLT_MAX) = 0
s = 0 × exp(0) + 0 × exp(0) = 0   // 安全退化

一、N1：两趟，两个 kernel↑

先把变量对上号：RMSNorm 只沿最后一维做，所以前两维可以直接拉平。 每个 token 用自己的 3584 个数独立归一化，与它在哪个 batch、哪个位置无关。

命令行 ./fused_ops n1 4096 3584 里的两个数字就是 M 和 N。 n1 1 3584 对应 decode batch=1： M=1，整张卡只有一个 block 在工作（本节末尾的 callout 讲这件事）。

rms_n1_pass1 拿到一行之后只做三件事：累加 → 归约 → 写回。 其中「归约」是两层嵌套的调用： blockReduceSum 内部调两次 warpReduceSum。 第一次 8 个 warp 并行各做各的，第二次只有 0 号 warp 做。

换个角度：同一份代码，不同的层面

CUDA 最别扭的一点：你写的是一个线程的代码，运行时却有 256 个线程同时跑这一份。 所以读这三个函数，不能只看「它算什么」，还要对每一句问 「这一句作用在哪个层面」：每个线程各算各的，一个 warp 的 32 个线程一起， 还是整个 block 的 256 个线程一起。

第二行最容易看错：这不叫「共享」，它们只是恰好相等。 row 也是普通寄存器变量，也有 256 份。 只是 blockIdx.x 对同一个 block 的所有线程都相同，所以算出来的值相同。

选几个代表？—— lane == 0 和 threadIdx.x == 0 不是一回事

这是这段代码最容易看错的地方。 threadIdx.x 是整个 block 内的编号（0…255）； lane 是每个 warp 内的编号（0…31，每个 warp 都从 0 重新数）。 两者的拆法如下：

const int lane = threadIdx.x & 31;   // 低 5 位 → 在本 warp 里排第几  0…31
const int warp = threadIdx.x >> 5;   // 高位   → 是第几个 warp      0…7

④ 那行有一个隐含前提：warp 数 ≤ 32。这个前提永远成立。 有了它，⑤ 用一个 warp 就能收完。BLOCK = 256 只有 8 个 warp； CUDA 也限制一个 block 最多 1024 线程（= 32 个 warp）。

下面把 §3 的两块归约积木、两个 kernel 和调用方式一次贴全，方便连续读完。 四级台阶共用同一份 blockReduceSum，后面 N2/N3/N4 不再重复贴。 完整文件见 §8 · 二。

// ═══ 积木：warp 内 + 跨 warp 的两段式求和（§3 · 一、二）══════════════════
static const int BLOCK = 256;                  // 8 个 warp

__inline__ __device__ float warpReduceSum(float v) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        v += __shfl_xor_sync(0xffffffffu, v, off);   // 蝶形，5 轮后全员持有全和
    return v;
}

__inline__ __device__ float blockReduceSum(float v, float *smem) {
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    v = warpReduceSum(v);                      // ① 每个 warp 各归各的
    if (lane == 0) smem[warp] = v;             // ② 8 个结果落进 shared
    __syncthreads();                           // ③ 等 8 个都写完
    v = (threadIdx.x < (BLOCK >> 5)) ? smem[lane] : 0.0f;  // 加法的幺元
    if (warp == 0) v = warpReduceSum(v);       // ④ 0 号 warp 把 8 个归成 1 个
    if (threadIdx.x == 0) smem[0] = v;         // ⑤ 放回 shared
    __syncthreads();                           // ⑥ 等它写完
    return smem[0];                            // ⑦ 全 block 一起读走
}

// ═══ N1 第一趟：只算每行的 rstd，写进一个长度 M 的小数组 ══════════════════
__global__ void rms_n1_pass1(const float *__restrict__ x, float *__restrict__ rstd,
                             int N, float eps) {
    __shared__ float smem[BLOCK / 32];
    const long long row = blockIdx.x;          // 一个 block 负责一行
    const float *xr = x + row * (long long)N;  // (long long) 防 M·N 溢出 int

    float acc = 0.0f;
    for (int i = threadIdx.x; i < N; i += BLOCK) {   // 跨步循环，天然合并
        float v = xr[i];
        acc += v * v;
    }
    acc = blockReduceSum(acc, smem);
    if (threadIdx.x == 0)
        rstd[row] = rsqrtf(acc / N + eps);     // rsqrtf 是硬件指令，一条搞定
}

// ═══ N1 第二趟：把 x 再读一遍，乘上 rstd 和权重，写出去 ═══════════════════
__global__ void rms_n1_pass2(const float *__restrict__ x, const float *__restrict__ w,
                             const float *__restrict__ rstd, float *__restrict__ y, int N) {
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    float       *yr = y + row * (long long)N;
    const float r = rstd[row];                 // 整个 block 读同一个地址 → 广播
    for (int i = threadIdx.x; i < N; i += BLOCK)
        yr[i] = xr[i] * r * w[i];               // ← 第二次读 x，12N 的根源
}

// ═══ 怎么调 ═════════════════════════════════════════════════════════════
// 显存四块：dx / dy 各 M×N 个 float；dw 长度 N（权重）；
//          dm 长度 M —— 就是 rstd，两趟之间靠它传值
CUDA_CHECK(cudaMalloc(&dm, M * sizeof(float)));

// grid = M（行数，一个 block 一行），block = BLOCK = 256
// 全文里这两行在 launch 那个 lambda 的 switch (kind) 的 case K_N1 分支下
// 两个 kernel 都不需要动态 shared，所以 <<<>>> 里只写两个参数
rms_n1_pass1<<<M, BLOCK>>>(dx, dm, N, eps);           // 读 x        → 写 rstd
rms_n1_pass2<<<M, BLOCK>>>(dx, dw, dm, dy, N);        // 读 x+rstd+w → 写 y

// 注意：两次发射之间不用写任何同步。同一个 stream 上，后发的 kernel
// 一定在前一个完全结束之后才开始 —— pass2 读到的 rstd 必然是全的。

代码里那两个关键字

__inline__ 是 C++ inline 的 GNU / CUDA 写法，nvcc 预定义了它。 它不是 CUDA 特有的加速开关，这里写它主要是习惯。 有三处容易误解：

要明确指定，用 __forceinline__（强制） 或 __noinline__（禁止）。

__restrict__：一个你对编译器的承诺

含义：在这个指针的作用域内，它指向的内存只通过它访问，不和别的指针重叠。 它是 C99 restrict 的 C++ 写法 （C++ 标准没有 restrict，各编译器用 __restrict__ 扩展）。

const 和 __restrict__ 要成对使用，缺一不可。 nvcc 只有在能证明「这块数据在整个 kernel 执行期间只读」时， 才生成走只读数据缓存的 load（LDG.E.CI，即 __ldg() 的路径，sm_35 起）。 const 说明「我不写它」，__restrict__ 说明「别人也不会通过别名写它」。 两条都有，编译器才用这条路径。

const float *__restrict__ x    // __restrict__ 在 * 右边 → 修饰指针本身
                               // const 在 * 左边       → 修饰指向的数据
float *const x                 // 这才是「指针本身不可改」，完全两码事

二、N2：融成一个 kernel，字节数一个都没省↑

__global__ void rms_n2(const float *__restrict__ x, const float *__restrict__ w,
                       float *__restrict__ y, int N, float eps) {
    __shared__ float smem[BLOCK / 32];
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    float       *yr = y + row * (long long)N;

    float acc = 0.0f;
    for (int i = threadIdx.x; i < N; i += BLOCK) {
        float v = xr[i];                       // ← 第一次读
        acc += v * v;
    }
    const float r = rsqrtf(blockReduceSum(acc, smem) / N + eps);
    for (int i = threadIdx.x; i < N; i += BLOCK)
        yr[i] = xr[i] * r * w[i];             // ← 第二次读，同样的地址
}

// ═══ 怎么调：和 N1 的唯一区别是只发一次 ═════════════════════════════════
rms_n2<<<M, BLOCK>>>(dx, dw, dy, N, eps);              // 读 x 两次 → 写 y

// 参数表里没有 dm 了 —— 不再需要那个长度 M 的中间数组，也省掉一次 kernel launch。
// 但 x 仍然从显存读了两遍，字节数一个都没省，还是 12N。

代码变短了，kernel 少了一个，但 xr[i] 仍出现两次。 显存流量还是 12N，与 N1 相同。

三、N3：把 x 留在寄存器 —— 这一级才动字节账↑

// 每线程最多存多少个 float。32 个 float = 32 个寄存器，离 255 的上限还远，不会 spill。
static const int MAXS = 32;                    // 单行最长 = BLOCK × MAXS = 8192

__global__ void rms_n3(const float *__restrict__ x, const float *__restrict__ w,
                       float *__restrict__ y, int N, float eps) {
    __shared__ float smem[BLOCK / 32];
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    float       *yr = y + row * (long long)N;

    float xs[MAXS];                            // 本线程负责的元素，全放寄存器
    float acc = 0.0f;
    #pragma unroll
    for (int c = 0; c < MAXS; ++c) {           // ← 固定次数，常量下标
        int i = threadIdx.x + c * BLOCK;
        float v = (i < N) ? xr[i] : 0.0f;      // 越界喂 0，不影响平方和
        xs[c] = v;
        acc += v * v;
    }
    const float r = rsqrtf(blockReduceSum(acc, smem) / N + eps);
    #pragma unroll
    for (int c = 0; c < MAXS; ++c) {
        int i = threadIdx.x + c * BLOCK;
        if (i < N) yr[i] = xs[c] * r * w[i];   // ← 直接用寄存器，不回显存
    }
}

// ═══ 怎么调：发射方式和 N2 一模一样 ═════════════════════════════════════
rms_n3<<<M, BLOCK>>>(dx, dw, dy, N, eps);              // 读 x 一次 → 写 y

// 但多了一个前提，必须在 host 侧先查 —— 整行得塞得进寄存器：
const bool needReg = (kind == K_N3 || kind == K_N4 || kind == K_S3 || kind == K_S4);
if (needReg && N > BLOCK * MAXS) {                 // 8192
    printf("✗ 这一级把整行存在寄存器里，单行最长 %d，而 N=%d。\n",
           BLOCK * MAXS, N);
    return 3;                                  // 不是实现偷懒，是硬件约束（§7·三）
}

// 错的写法（看起来更自然，但会毁掉这一级）
int cnt = 0;
for (int i = threadIdx.x; i < N; i += BLOCK) {
    xs[cnt++] = xr[i];                     // 下标是运行期算出来的
    ...
}

// 对的写法
#pragma unroll
for (int c = 0; c < MAXS; ++c) {              // 展开后 xs[0] xs[1] … 全是常量下标
    int i = threadIdx.x + c * BLOCK;
    ...
}

代价是单行长度有上限：N ≤ BLOCK × MAXS = 256 × 32 = 8192。 超过就放不下。程序会直接拒绝并说明原因，不会悄悄 spill。 这个限制不是实现偷懒， 它和 FlashAttention 必须分块是同一个物理限制，§5 · 五 细说。

四、N4：float4，字节不变，只动条数↑

// float4 版：8 个 float4 也是 32 个 float。两条路径的寄存器压力刻意做成一样，
// 这样 N3→N4 就是干净的单变量对照 —— 只动指令条数，不动寄存器。
static const int MAXV = 8;                     // MAXV × 4 = MAXS = 32

__global__ void rms_n4(const float *__restrict__ x, const float *__restrict__ w,
                       float *__restrict__ y, int N, float eps) {
    __shared__ float smem[BLOCK / 32];
    const long long row = blockIdx.x;
    const int n4 = N >> 2;                        // 一行有多少个 float4

    // reinterpret_cast 把 float* 重新解释成 float4*，地址必须 16 字节对齐 ——
    // cudaMalloc 返回的指针是 256 字节对齐的，而 row*N 只要 N%4==0 就仍然对齐。
    const float4 *xr = reinterpret_cast<const float4 *>(x + row * (long long)N);
    float4       *yr = reinterpret_cast<float4 *>(y + row * (long long)N);
    const float4 *wr = reinterpret_cast<const float4 *>(w);

    float4 xs[MAXV];                           // 8 个 float4 = 32 个寄存器，和 N3 一样
    float acc = 0.0f;
    #pragma unroll
    for (int c = 0; c < MAXV; ++c) {
        int i = threadIdx.x + c * BLOCK;
        float4 v = (i < n4) ? xr[i] : make_float4(0.f, 0.f, 0.f, 0.f);
        xs[c] = v;
        acc += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
    }
    const float r = rsqrtf(blockReduceSum(acc, smem) / N + eps);
    #pragma unroll
    for (int c = 0; c < MAXV; ++c) {
        int i = threadIdx.x + c * BLOCK;
        if (i >= n4) continue;
        float4 v = xs[c], g = wr[i], o;        // 权重也按 float4 读
        o.x = v.x * r * g.x;  o.y = v.y * r * g.y;
        o.z = v.z * r * g.z;  o.w = v.w * r * g.w;
        yr[i] = o;                             // 一条 STG.128 写回 4 个 float
    }
}

// ═══ 怎么调：和 N3 一字不差 ═════════════════════════════════════════════
rms_n4<<<M, BLOCK>>>(dx, dw, dy, N, eps);              // 读 x 一次 → 写 y

// 除了 N3 那条寄存器长度的守卫，还多一条对齐守卫：
const bool needVec4 = (kind == K_N4 || kind == K_S4);
if (needVec4 && (N & 3)) {                          // N 不是 4 的倍数
    printf("✗ float4 版要求 N 是 4 的倍数（当前 N=%d）\n", N);
    return 3;
}

两个读入循环到底差在哪

两边的 int i = threadIdx.x + c * BLOCK; 一字不差 ，但 N3 里 i 是「第几个 float」， N4 里是「第几个 float4」。 这是全部差异的来源，其余几处差异都由它引出。

MAXV = 8 是故意选的：8 个 float4 = 32 个 float，与 N3 的 MAXS = 32 相同。 两条路径的寄存器压力相同， N3→N4 才是干净的单变量对照：只动指令宽度，不动寄存器数。

预测：N4 比 N3 快约 1.0×–1.15×。 字节数没变，理论上不该变快。如果快一点， 原因是访存指令更少、内存请求更少（4 条 32 位并成 1 条 128 位）， 即 SGEMM 那边的「条数账」和「请求账」。 但这两本账只在带宽没跑满时才有空间： 如果 N3 的 MBU 已经很高，N4 应该几乎不变快。这一条本身就是判据。

一、为什么要减 max↑

softmax(x)_i = exp(x_i) / Σ_j exp(x_j)       // 教科书定义

直接照着算会溢出。 fp32 能表示的最大值约 3.4×10³⁸，而 exp(89) 就超了。 注意力打分经过 1/√d 缩放后通常在 ±20 以内，但训练不稳或上下文很长时可能超过 90。 一旦溢出，结果就是 inf / inf = NaN。

softmax(x)_i = exp(x_i − m) / Σ_j exp(x_j − m)    // m = max(x)，恒等变形

// 为什么相等：分子分母同乘 exp(−m)，约掉了
// 为什么安全：x_i − m ≤ 0，所以 exp(·) ∈ (0, 1]，永远不会溢出

减 max 的代价是多扫一趟。 要减 max，就得先知道 max； 要知道 max，就得把整行看完 —— 于是朴素实现必须扫三趟。

二、递推式怎么来的↑

换个做法：不等看完整行，边扫边维护一对 (m, s)。 约定 m_k 是前 k 个元素的最大值， s_k = Σ_{j≤k} exp(x_j − m_k) 是以当前最大值为基准的部分和。

现在来了第 k+1 个元素 x。新的最大值是 m' = max(m_k, x)。 但 s_k 的基准是旧的 m_k。基准变了，s_k 就不能直接加。

先把两个记号分清楚 —— 同一批数 x_1 … x_k，用不同基准记账，得到的是不同的数：

s_k      = Σ_{j≤k} exp(x_j − m_k)      // 手里已经有的，旧基准 m_k
s_k(m')  = Σ_{j≤k} exp(x_j − m')       // 想要的，新基准 m'

两者只差一个常数倍 —— 下面五步就是把这个倍数找出来。

// ① 在指数里加个零（−m' + m'），什么都没改变
    x_j − m_k  =  x_j − m' + m' − m_k

// ② 重新分组：一半随 j 变，一半是常数
               =  (x_j − m')  +  (m' − m_k)
                  └ 随 j 变 ┘     └ 与 j 无关 ┘

// ③ 用 exp(a + b) = exp(a) · exp(b) 拆开
    exp(x_j − m_k)  =  exp(x_j − m') · exp(m' − m_k)

// ④ 对 j 求和 —— 常数可以提到求和号外面
    s_k  =  Σ_{j≤k} exp(x_j − m_k)
         =  Σ_{j≤k} [ exp(x_j − m') · exp(m' − m_k) ]
         =  exp(m' − m_k) · Σ_{j≤k} exp(x_j − m')
            └ 提出来的公因子 ┘     └──── 这就是 s_k(m') ────┘
         =  exp(m' − m_k) · s_k(m')

// ⑤ 反解出我们真正想要的 s_k(m')：两边同除 exp(m' − m_k)
    s_k(m')  =  s_k / exp(m' − m_k)
             =  s_k · exp(m_k − m')       // ← 指数反号了！
                      ↑ 这才是代码里那个缩放因子

把换算好的旧账，和本来就以 m′ 记账的新项加起来，就是那两行递推：

m' = max(m_k, x)
s' = s_k · exp(m_k − m') + exp(x − m')
     └─── 旧账换算到新基准 ───┘   └─ 新项，本来就以 m' 记账 ─┘

递推只有这两行。 直观含义：遇到更大的 max，就把已经累计的和按比例"缩"回去。 缩放因子 exp(m_k − m') 一定 ≤ 1，因为 m' ≥ m_k。 没遇到更大的 max 时，m' = m_k，因子是 exp(0) = 1，退化成普通累加。

递推一路只维护两个标量，却和「先扫完整行拿到 max、再从头累加」得到同一个数。 最后一行的校验：exp(1−5) + exp(2−5) + exp(5−5) + exp(3−5) = 0.0183 + 0.0498 + 1 + 0.1353 = 1.2034。 「加 3」这一步 max 没变，因子退化成 1，就是普通累加。

三、结合律：为什么能做树形归约↑

先把合并运算写成一般形式（不再限定第二个操作数是单个元素）：

(m₁, s₁) ⊕ (m₂, s₂) = ( max(m₁, m₂),
                         s₁·exp(m₁ − max) + s₂·exp(m₂ − max) )

式子的两个分量，各用一句话就能说清。 式子里 max 出现三次，指的是同一个数，先给它起个名：

令 m = max(m₁, m₂)

(m₁,s₁) ⊕ (m₂,s₂) = (  m  ,  s₁·exp(m₁−m) + s₂·exp(m₂−m)  )
                       ↑              ↑                  ↑
                  合并后的最大值   A 的账换到新基准    B 的账换到新基准

// 第一个分量：两堆数合起来的最大值，就是两个最大值里更大的那个 —— 显然
// 第二个分量：s₁ 以 m₁ 记账、s₂ 以 m₂ 记账，量纲不同不能直接加，
//             所以两边都先换算到新基准 m 再加（§5·二 那个换基准公式，这次两边都要换）

// 代进上面的例子：
exp(m₁ − m) = exp(2−5) = 0.0498      // m 小的那边，被缩下去
exp(m₂ − m) = exp(5−5) = 1.0000      // m 大的那边，因子恒为 1

两个因子里必有一个是 exp(0) = 1，另一个 ≤ 1 —— 所以这一步永远不会溢出。

(m_k, s_k) ⊕ (x, 1)
  = ( max(m_k, x),  s_k·exp(m_k−m′) + 1·exp(x−m′) )
  = ( m′,           s_k·exp(m_k−m′) + exp(x−m′)   )   // 和 §5·二 一字不差

证法一：合出来的，就是并集的摘要

图 23 已经证完了结合律的一大半。 图 23 对应一条引理： 两段各自的摘要合并起来，正好等于「这两段并在一起」那一整段的摘要。 把引理当作已知，剩下的证明只有三行：

// 引理（就是图 23，直接代定义验算即可）
    摘要(A) ⊕ 摘要(B)  =  摘要(A ∪ B)

// 于是两种括号法分别是
    ( 摘要A ⊕ 摘要B ) ⊕ 摘要C  =  摘要(A∪B) ⊕ 摘要C  =  摘要( (A∪B) ∪ C )
    摘要A ⊕ ( 摘要B ⊕ 摘要C )  =  摘要A ⊕ 摘要(B∪C)  =  摘要( A ∪ (B∪C) )

// 而集合的并当然可结合
    (A∪B) ∪ C  =  A ∪ (B∪C)

// 两边摘要的输入是同一个集合，输出自然相同
    (摘要A ⊕ 摘要B) ⊕ 摘要C  =  摘要A ⊕ (摘要B ⊕ 摘要C)        // ∎

整个证明归结到「集合的并可结合」，这一点无可怀疑。 不管你怎么加括号，合出来的都是「这几段并在一起」 那一整段的摘要；而并在一起，跟你怎么分组无关。

T(A) = s · exp(m) = Σ_{j∈A} exp(x_j)      // 真正的"无偏指数和"

// 换坐标的映射
φ(m, s) = ( m , s·exp(m) ) = ( M , T )

// 新坐标里的运算：两个分量各管各的
(M₁, T₁) ⊞ (M₂, T₂) = ( max(M₁, M₂) , T₁ + T₂ )

// 要证：⊕ 可结合。路线四步 ——
①  φ 是单射                     // 能原路还回来
②  φ(A ⊕ B) = φ(A) ⊞ φ(B)       // φ 把 ⊕ 翻译成 ⊞
③  ⊞ 可结合                      // 显然，但要说
④  ①+②+③  ⟹  ⊕ 可结合

给定新坐标 (M, T)，可以唯一还原回去：
    m = M
    s = T · exp(−M)

所以   φ(m₁,s₁) = φ(m₂,s₂)  ⟹  (m₁,s₁) = (m₂,s₂)

令 m = max(m₁, m₂)。按 ⊕ 的定义
    A ⊕ B = ( m , s₁·exp(m₁−m) + s₂·exp(m₂−m) )

// 第一个分量
    M(A ⊕ B) = m = max(m₁, m₂) = max( M(A), M(B) )        // ⊞ 的第一分量 ✓

// 第二个分量 —— 关键是 m 会被完全消掉
    T(A ⊕ B) = [ s₁·exp(m₁−m) + s₂·exp(m₂−m) ] · exp(m)
             = s₁·exp(m₁−m)·exp(m) + s₂·exp(m₂−m)·exp(m)   // 乘进括号
             = s₁·exp(m₁−m+m)      + s₂·exp(m₂−m+m)        // exp(a)exp(b)=exp(a+b)
             = s₁·exp(m₁)          + s₂·exp(m₂)
             = T(A) + T(B)                                // ⊞ 的第二分量 ✓

⟹  φ(A ⊕ B) = φ(A) ⊞ φ(B)

φ( (A⊕B) ⊕ C )  =  φ(A⊕B) ⊞ φ(C)          // 由 ②
                 =  ( φA ⊞ φB ) ⊞ φC        // 由 ②
                 =  φA ⊞ ( φB ⊞ φC )        // 由 ③ ← 唯一用到结合律的地方
                 =  φA ⊞ φ(B⊕C)             // 由 ②
                 =  φ( A ⊕ (B⊕C) )          // 由 ②

两端 φ 的值相同，而 φ 是单射（①），所以

    (A⊕B) ⊕ C  =  A ⊕ (B⊕C)                        // ∎

幺元：为什么是 (−∞, 0)

取 e = (−∞, 0)，对应 T(e) = 0 —— 也就是空集。

验证 e ⊕ (m, s)：令 M = max(−∞, m) = m
    s′ = 0·exp(−∞ − m) + s·exp(m − m)
       = 0·0            + s·1
       = s
⟹  e ⊕ (m, s) = (m, s)   // 另一侧同理，交换律已证

数学上 −∞ 没有问题，但在浮点里 (−∞) − (−∞) = NaN。 所以代码里用 -FLT_MAX 顶替 —— 那个坑和修法见 §3 · 三。

再核对 T：exp(1)+exp(2)+exp(5) = 158.520497， 而 1.068102707 × exp(5) = 158.520497 —— 两种括号法算出来的 (m, s) 还原出同一个 T，正是 ② 说的那件事。

（T 只是证明用的中间量，代码里绝不能真的去算它 —— exp(m) 恰恰就是我们一开始要避免的那个溢出。 保持 (m, s) 的形式，是为了让每一步的中间值都留在安全区间里。）

四、数值上到底安不安全↑

唯一真正的风险是 s = 0 导致除零。 但 s 至少包含 exp(m − m) = 1 这一项 （最大的那个元素），所以 s ≥ 1，永远不会是 0。 前提是这一行至少有一个元素 —— 空行是调用方的错，kernel 不负责。

五、和 FlashAttention 的关系↑

把 §5 · 三 证明的结合律和 §4 · 三 那个"单行最长 8192"的约束放在一起看， FlashAttention 就是必然结果。

朴素三趟 S1     7.52 GB/层  →  7.46 ms  →  ×28 层 = 209 ms
本文最优 S3     3.76 GB/层  →  3.73 ms  →  ×28 层 = 104 ms
FlashAttention  0 GB        →  打分矩阵根本不存在

一、S1：三个 kernel，16N↑

S1 和 N1 是同一个套路：趟与趟之间只能靠 kernel 边界同步， 中间结果存在长度 M 的小数组里。 softmax 要传两个中间量（max 和 sum）， 所以 S1 有三个 kernel、两个小数组。

// ═══ 积木一：两个 warp 级归约，只差一个运算符 ═════════════════
__inline__ __device__ float warpReduceMax(float v) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, off));   // ← 取大
    return v;
}
__inline__ __device__ float warpReduceSum(float v) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        v += __shfl_xor_sync(0xffffffffu, v, off);            // ← 求和，其余一模一样
    return v;
}

// ═══ 积木二：block 级求和 —— 结尾那三行是广播 ══════════════
__inline__ __device__ float blockReduceSum(float v, float *smem) {
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    v = warpReduceSum(v);                      // ① 8 个 warp 并行各归各的
    if (lane == 0) smem[warp] = v;             // ② 每 warp 交一个
    __syncthreads();                           // ③ 等 8 个都写完
    v = (threadIdx.x < (BLOCK >> 5)) ? smem[lane] : 0.0f;  // 加法的幺元
    if (warp == 0) v = warpReduceSum(v);       // ④ warp 0 收尾
    if (threadIdx.x == 0) smem[0] = v;         // ⑤ ┐
    __syncthreads();                           // ⑥ ├ 广播给全 block
    return smem[0];                            // ⑦ ┘
}

// 注意：这份代码里没有 blockReduceMax —— 求 max 那边是直接写在 sm_s1_max 里的，
//       而且故意省掉了 ⑤⑥⑦，因为只有 thread 0 要用。下面对照着看。

// ═══ ① 求 max：整行扫一遍，每行吐出一个 float ═════════════════════════
__global__ void sm_s1_max(const float *__restrict__ x, float *__restrict__ mx, int N) {
    __shared__ float smem[BLOCK / 32];
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;

    float m = -FLT_MAX;                            // 求 max 的幺元（不能用 -INFINITY，§3·三）
    for (int i = threadIdx.x; i < N; i += BLOCK) m = fmaxf(m, xr[i]);

    // 两段式归约，但这里手写、没有广播尾巴 —— 只有 thread 0 要用结果
    const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
    m = warpReduceMax(m);                      // warp 内，8 个 warp 并行
    if (lane == 0) smem[warp] = m;             // 每 warp 交一个
    __syncthreads();
    m = (threadIdx.x < BLOCK / 32) ? smem[lane] : -FLT_MAX;   // 其余 lane 喂幺元
    if (warp == 0) m = warpReduceMax(m);       // warp 0 收尾
    if (threadIdx.x == 0) mx[row] = m;         // 全 block 只落笔一次
}

// ═══ ② 求 exp 和：要用 ① 的 m，所以必须等 ① 整个 grid 跑完 ═══════
__global__ void sm_s1_sum(const float *__restrict__ x, const float *__restrict__ mx,
                          float *__restrict__ sm, int N) {
    __shared__ float smem[BLOCK / 32];
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    const float m = mx[row];                   // 256 个线程读同一个地址 → 广播

    float s = 0.0f;
    for (int i = threadIdx.x; i < N; i += BLOCK)
        s += __expf(xr[i] - m);                  // ← 第二遍读 x
    s = blockReduceSum(s, smem);               // 复用 §3 的 helper（带广播尾巴）
    if (threadIdx.x == 0) sm[row] = s;
}

// ═══ ③ 归一化：要用 ② 的 s ═══════════════════════════════════════════
__global__ void sm_s1_div(const float *__restrict__ x, const float *__restrict__ mx,
                          const float *__restrict__ sm, float *__restrict__ y, int N) {
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    float       *yr = y + row * (long long)N;
    const float m = mx[row], inv = 1.0f / sm[row];   // 先取倒数，循环里就不用除法

    for (int i = threadIdx.x; i < N; i += BLOCK)
        yr[i] = __expf(xr[i] - m) * inv;           // ← 第三遍读 x；exp 又算了一遍
}

// ═══ 怎么调 ═════════════════════════════════════════════════════════════
// 两个长度 M 的小数组接力：dm 存每行的 max，ds 存每行的 Σexp
CUDA_CHECK(cudaMalloc(&dm, M * sizeof(float)));
CUDA_CHECK(cudaMalloc(&ds, M * sizeof(float)));

sm_s1_max<<<M, BLOCK>>>(dx, dm, N);                     // 读 x        → 写 dm
sm_s1_sum<<<M, BLOCK>>>(dx, dm, ds, N);                  // 读 x + dm   → 写 ds
sm_s1_div<<<M, BLOCK>>>(dx, dm, ds, dy, N);              // 读 x + dm+ds→ 写 y

// 三次发射之间不用写同步 —— 同一个 stream 上后发的 kernel
// 一定在前一个整个 grid 结束之后才开始。这正是趟与趟之间唯一的同步手段。

__expf 是硬件近似指令（对应 SASS 的 MUFU.EX2）， 比 expf 快一个量级。误差随 |x| 增大（NVIDIA 给的上界是 2 + floor(|1.16x|) ulp），在打分常见的 ±20 区间内约 1e-5 量级。 推理里普遍用它：softmax 的结果马上要做加权平均， 这点误差远小于 fp16 本身的表示误差。校验时容差要放宽到 2e-3。

二、S2：online 递推，省掉一趟↑

S2 省掉的是 S1 的第 ① 趟，即求 max 那一趟。 靠的是 §5 · 二 的数学改写，不是代码技巧： online 递推让 max 和 sum 在同一趟里同时推进。 S1 的第三个 kernel 原样复用，所以这一级只新写了一个 kernel。

// ═══ 积木一：联合归约 —— warp 内把 (m, s) 当一个整体 ═════════════════
__inline__ __device__ void warpReduceMaxSum(float &m, float &s) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        float m_o = __shfl_xor_sync(0xffffffffu, m, off);   // 伙伴的 max
        float s_o = __shfl_xor_sync(0xffffffffu, s, off);   // 伙伴的 sum（每轮 2 条 shuffle）
        float m_n = fmaxf(m, m_o);                          // 新的公共基准
        s = s * __expf(m - m_n) + s_o * __expf(m_o - m_n);  // 两边都换算到 m_n 再加
        m = m_n;
    }
}

// ═══ 积木二：block 级联合归约 —— 两个分量一起收 ═══════════════════════
__inline__ __device__ void blockReduceMaxSum(float &m, float &s, float *smem) {
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int nwarp = BLOCK >> 5;                           // = 8
    warpReduceMaxSum(m, s);                                 // ① 8 个 warp 并行各归各的
    if (lane == 0) { smem[warp] = m; smem[nwarp + warp] = s; }  // ② m 放前 8 格，s 放后 8 格
    __syncthreads();                                        // ③
    if (warp == 0) {
        // 幺元是 (-inf, 0)：最大值取 -inf 不影响 max，和取 0 不影响求和
        m = (lane < nwarp) ? smem[lane]         : -FLT_MAX;
        s = (lane < nwarp) ? smem[nwarp + lane] : 0.0f;
        warpReduceMaxSum(m, s);                             // ④ warp 0 收尾
        if (lane == 0) { smem[2 * nwarp] = m; smem[2 * nwarp + 1] = s; }   // ⑤ 存进最后 2 格
    }
    __syncthreads();                                        // ⑥ 注意：屏障在 if 外面
    m = smem[2 * nwarp];                                      // ⑦ ┐ 广播给全 block
    s = smem[2 * nwarp + 1];                                  //   ┘
}

// ═══ 一趟同时拿到 m 和 s ═════════════════════════════════════════════
__global__ void sm_s2_ms(const float *__restrict__ x, float *__restrict__ mx,
                         float *__restrict__ sm, int N) {
    __shared__ float smem[BLOCK / 32 * 2 + 2];              // 18 个 float
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;

    float m = -FLT_MAX, s = 0.0f;              // 幺元：空集的摘要
    for (int i = threadIdx.x; i < N; i += BLOCK) {
        float v = xr[i];                       // ← 整个 S2 只在这里读 x
        // 递推三行：见到更大的 max，就把已经累好的 s 按比例缩回去
        float m_n = fmaxf(m, v);
        s = s * __expf(m - m_n) + __expf(v - m_n);
        m = m_n;
    }
    blockReduceMaxSum(m, s, smem);             // 256 对 (m,s) 用 ⊕ 归成一对
    if (threadIdx.x == 0) { mx[row] = m; sm[row] = s; }     // 一次写两个小数组
}

// ═══ 第二个 kernel：原样复用 S1 的 sm_s1_div（再贴一遍，省得翻回去）═══
__global__ void sm_s1_div(const float *__restrict__ x, const float *__restrict__ mx,
                          const float *__restrict__ sm, float *__restrict__ y, int N) {
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    float       *yr = y + row * (long long)N;
    const float m = mx[row], inv = 1.0f / sm[row];   // 先取倒数，循环里不做除法

    for (int i = threadIdx.x; i < N; i += BLOCK)
        yr[i] = __expf(xr[i] - m) * inv;           // ← 第二遍读 x
}

// ═══ 怎么调：只发两次 ═══════════════════════════════════════════════════
sm_s2_ms <<<M, BLOCK>>>(dx, dm, ds, N);                  // 读 x        → 写 dm + ds
sm_s1_div<<<M, BLOCK>>>(dx, dm, ds, dy, N);              // ← S1 的第三个 kernel，原样复用

// dm / ds 的分配和 S1 一模一样；少的只是一次 kernel 发射和一整趟读 x。

三、S3：一趟，8N↑

S3 把最后一趟也省掉了：整行只经过显存一次，(m, s) 全程留在片上。 代价是整行要放进寄存器。这和 N3 是同一个套路， 连「固定次数 + 常量下标」那条规矩也一样。 于是 dm / ds 两个中间数组消失，只剩一个 kernel。

// 每线程最多存多少个 float —— 和 N3 共用同一个常量
static const int MAXS = 32;                    // 单行最长 = BLOCK × MAXS = 8192

// 积木仍是 §6·二 那两个：warpReduceMaxSum / blockReduceMaxSum，原样复用
// 但这一级真的需要它结尾那几行广播 —— 256 个线程每人都要拿 (m,s) 去归一化自己手上的元素

// ═══ 一个 kernel 走完全程 ═══════════════════════════════════════════════
__global__ void sm_s3(const float *__restrict__ x, float *__restrict__ y, int N) {
    __shared__ float smem[BLOCK / 32 * 2 + 2];
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    float       *yr = y + row * (long long)N;

    // 常量下标 + #pragma unroll，理由同 rms_n3（否则 xs 会被塞进 local memory）
    float xs[MAXS];                            // 整行摊在寄存器上
    float m = -FLT_MAX, s = 0.0f;
    #pragma unroll
    for (int c = 0; c < MAXS; ++c) {
        int i = threadIdx.x + c * BLOCK;
        if (i >= N) { xs[c] = -FLT_MAX; continue; }  // 越界塞幺元，见下面的 callout
        float v = xr[i];                       // ← 整个 S3 只在这里读一次 x
        xs[c] = v;                             //   读完就留着，不再回显存
        float m_n = fmaxf(m, v);               // ┐
        s = s * __expf(m - m_n) + __expf(v - m_n);  // ├ 递推三行，同 S2
        m = m_n;                               // ┘
    }
    blockReduceMaxSum(m, s, smem);             // 256 对 (m,s) 归成一对，并广播给全员

    const float inv = 1.0f / s;                // 先取倒数，循环里不做除法
    #pragma unroll
    for (int c = 0; c < MAXS; ++c) {
        int i = threadIdx.x + c * BLOCK;
        if (i < N) yr[i] = __expf(xs[c] - m) * inv;   // ← 直接用寄存器，不回显存
    }
}

// ═══ 怎么调：只发一次，而且参数表里没有中间数组 ═════════════════════════
sm_s3<<<M, BLOCK>>>(dx, dy, N);                          // 读 x 一次 → 写 y

// 和 N3 共用同一条守卫 —— 整行得塞得进寄存器：
const bool needReg = (kind == K_N3 || kind == K_N4 || kind == K_S3 || kind == K_S4);
if (needReg && N > BLOCK * MAXS) {                 // 8192
    printf("✗ 这一级把整行存在寄存器里，单行最长 %d，而 N=%d。\n",
           BLOCK * MAXS, N);
    return 3;
}

四、S4：float4，顺便演示了分块合并↑

S4 = S3 + float4：字节数不变（仍是 8N），减少的是访存指令条数。 它还多做了一件事：一次读进 4 个元素， 所以必须先把这 4 个合成一个小块摘要，再并进累计值。 这用的是 §5 · 三 那个一般形式的 ⊕， 而不是 S3 那个「并进一个元素」的特例。

// float4 版：8 个 float4 也是 32 个 float，寄存器压力和 S3 刻意做成一样
static const int MAXV = 8;                     // MAXV × 4 = MAXS = 32

// 积木仍是 §6·二 那两个，原样复用；广播这一级同样真的需要

// ═══ 一个 kernel，一次搬 4 个 ═══════════════════════════════════════════
__global__ void sm_s4(const float *__restrict__ x, float *__restrict__ y, int N) {
    __shared__ float smem[BLOCK / 32 * 2 + 2];
    const long long row = blockIdx.x;
    const int n4 = N >> 2;                        // 一行有多少个 float4
    const float4 *xr = reinterpret_cast<const float4 *>(x + row * (long long)N);
    float4       *yr = reinterpret_cast<float4 *>(y + row * (long long)N);

    float4 xs[MAXV];                           // 8 个 float4 = 32 个寄存器
    float m = -FLT_MAX, s = 0.0f;
    #pragma unroll
    for (int c = 0; c < MAXV; ++c) {
        int i = threadIdx.x + c * BLOCK;
        if (i >= n4) { xs[c] = make_float4(-FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX); continue; }
        float4 v = xr[i];                      // ← 一条 LDG.E.128，16 字节
        xs[c] = v;

        // 第 ① 层：先在寄存器里把这 4 个合成一个小块摘要 (mv, sv)
        float mv = fmaxf(fmaxf(v.x, v.y), fmaxf(v.z, v.w));
        float sv = __expf(v.x - mv) + __expf(v.y - mv) + __expf(v.z - mv) + __expf(v.w - mv);

        // 第 ② 层：再把这个小块用 ⊕ 并进累计值 —— 第二个操作数不再是单个元素
        float m_n = fmaxf(m, mv);
        s = s * __expf(m - m_n) + sv * __expf(mv - m_n);   // ← 注意 sv 前面这个系数
        m = m_n;
    }
    blockReduceMaxSum(m, s, smem);

    const float inv = 1.0f / s;
    #pragma unroll
    for (int c = 0; c < MAXV; ++c) {
        int i = threadIdx.x + c * BLOCK;
        if (i >= n4) continue;
        float4 v = xs[c], o;                   // 直接用寄存器里的 xs
        o.x = __expf(v.x - m) * inv;  o.y = __expf(v.y - m) * inv;
        o.z = __expf(v.z - m) * inv;  o.w = __expf(v.w - m) * inv;
        yr[i] = o;                             // 一条 STG.128 写回 4 个 float
    }
}

// ═══ 怎么调：和 S3 一字不差 ═════════════════════════════════════════════
sm_s4<<<M, BLOCK>>>(dx, dy, N);                          // 读 x 一次 → 写 y

// 除了 S3 那条寄存器长度守卫，还多一条对齐守卫：
const bool needVec4 = (kind == K_N4 || kind == K_S4);
if (needVec4 && (N & 3)) {                          // N 不是 4 的倍数
    printf("✗ float4 版要求 N 是 4 的倍数（当前 N=%d）\n", N);
    return 3;
}

sv = … 这一行算的是什么

这一行直接照抄「摘要」的定义，一个字都没多。

// 摘要的定义（§5·三）
摘要(A) = ( max(A) , Σ_{j∈A} exp(x_j − max(A)) )

// 代进 A = { v.x, v.y, v.z, v.w }
mv = max(A)                                                // ← 第一个分量
sv = exp(v.x−mv) + exp(v.y−mv) + exp(v.z−mv) + exp(v.w−mv)   // ← 第二个分量

// 所以这两行合起来 = 摘要({v.x, v.y, v.z, v.w})，(mv, sv) 是一对，拆开都没意义

块内用的是「朴素两趟」，不是递推

这两行的结构是先求 max，再求和，正是 S1 那套朴素做法， 不是 online 递推。

// 块内（4 个元素，全在寄存器里）：朴素两趟
mv = max(...)                    // 第 1 趟：求 max，0 次 exp
sv = Σ exp(v_i − mv)             // 第 2 趟：求和，4 次 exp        → 共 4 次

// 如果改成递推、逐个并进来：
for v_i in {x, y, z, w}:
    m_n = fmaxf(m, v_i);
    s = s * exp(m − m_n) + exp(v_i − m_n);   // 每元素 2 次 exp   → 共 8 次

4 次对 8 次：在块内，朴素法反而少算一半 exp。 S1 的问题不在算得多， 而在「两趟」意味着两趟读显存。 而这四个数就在寄存器里，两趟不花钱， 于是朴素法「少算一半 exp」的优势就显出来了。

B + 2 怎么数出来

先定义「块」：块 = 一次循环迭代里拿到的那几个元素。 两个 kernel 覆盖的元素总数相同，差别只在每轮拿几个：

S3：每元素 2 次（fused_ops.cu:378）

float m_n = fmaxf(m, v);                            // 0 次
s = s * __expf(m - m_n) + __expf(v - m_n);          // ← 2 次
        └ 缩放因子：把已累计的 s ┘   └ 新项 ┘
          从旧基准 m 换到 m_n

关键：每处理一个元素，都要重新换一次基准， 因为 m_n 每轮都可能变。所以 B 个元素要 2B 次。

S4：B + 2 次（fused_ops.cu:409 与 411）

float mv = fmaxf(fmaxf(v.x, v.y), fmaxf(v.z, v.w));                     // 0 次
float sv = __expf(v.x-mv)+__expf(v.y-mv)+__expf(v.z-mv)+__expf(v.w-mv);   // ← B = 4 次
           ↑ 基准 mv 是这 4 个自己的 max，一次定好、中途不变 → 不需要任何缩放因子

float m_n = fmaxf(m, mv);                                               // 0 次
s = s * __expf(m - m_n) + sv * __expf(mv - m_n);                     // ← 2 次
        └ 换累计值的基准 ┘   └ 换整块 sv 的基准 ┘
                             ↑ 和 S3 的唯一结构差别：sv 这个系数

// 整个块只换一次基准 → B + 2 = 6

五、长行：s3/s4 会拒绝跑，这是设计↑

if (needReg && N > BLOCK * MAXS) {
    printf("✗ 这一级把整行存在寄存器里，单行最长 %d，而 N=%d。\n", BLOCK * MAXS, N);
    printf("  这不是实现偷懒，是硬件约束 —— 也正是 FlashAttention 必须分块的原因。\n");
    return 3;
}

Qwen2.5-7B 的 vocab 是 152064。 采样时要对这么长的一行做 softmax，而 S3 的上限是 8192。 差 18 倍，无论如何都放不进寄存器。

一、主阶梯：八条↑

二、形状扫描：四条↑

第 ⑫ 条和 SGEMM 形状扫描的结论相同： 在 decode 一侧，辛苦优化的那本账会失效， 因为瓶颈变了。 SGEMM 那边是"砍指令条数救不了带宽墙"， 这边是"砍字节救不了空转的 SM"。

三、长行那一条↑

一、验收清单↑

# 测试机
nvidia-smi                          # 确认卡上没有别人
bash run_fused.sh                   # 约 20 分钟

# 想快速过一遍
ROUNDS=1 bash run_fused.sh

二、fused_ops.cu 全文↑

十一个 kernel + 完整脚手架（设备查询、occupancy、字节账、CPU 参照校验、中位数计时）。 计时和校验的写法和 sgemm_all.cu 相同： 预热次数按 ITERS 分档；样本数为偶数时，取中间两个的平均。 这两处写法来自 SGEMM 那一轮的教训，来历见本阶段 README「怎么复现」的四条规矩。

// fused_ops.cu:533-539
const size_t nElem = (size_t)M * N;
const size_t bytes = nElem * sizeof(float);
std::vector<float> hx(nElem), hw(N), hy(nElem);
std::mt19937 rng(12345);
std::uniform_real_distribution<float> dist(-3.0f, 3.0f);
for (size_t i = 0; i < nElem; ++i) hx[i] = dist(rng);
for (int i = 0; i < N; ++i)        hw[i] = 0.5f + 0.5f * dist(rng) / 3.0f;

// U(−L, L) 的密度是 1/(2L)
E[x²] = ∫_{−L}^{L} x² · 1/(2L) dx = (1/(2L))·[x³/3]_{−L}^{L}
      = (1/(2L))·(L³/3 + L³/3)
      = L²/3

// 于是缩放因子
r = 1/√(E[x²]) = √3 / L

// 对任意 j：x_j ≥ −L 而 max < L
x_j − max ∈ (−2L, 0]        →        exp(x_j − max) ∈ (e^(−2L), 1]

// 把 max 近似当常数 L（N 大时 E[max] = L − 2L/(N+1)，近似极好）
E[exp(x − L)] = ∫_{−L}^{L} e^(x−L) · 1/(2L) dx = (1/(2L))·[e^(x−L)]_{−L}^{L}
              = (1 − e^(−2L)) / (2L)

// L = 3 时
E[s] = N · (1 − e^(−6)) / 6 = 0.166254 · N

// fused_ops.cu:461 —— verifyRows 的判定
if (d > 2e-3 * fabs(ref[i]) + 1e-7) { ... }

// 常数输入 x ≡ c 时，走一遍递推：
// 第 1 个元素：m = −FLT_MAX → m_n = c
正确：  s = s * exp(m − m_n) + exp(c − c) = 0 * exp(−FLT_MAX − c) + 1 = 1
漏掉：  s = s               + exp(c − c) = 0                        + 1 = 1   ← 一样

// 第 2 个元素起：m 再也不变，m_n == m
正确：  s = s * exp(0) + 1 = s + 1
漏掉：  s =              s + 1                                            ← 还是一样

// 结论：rescale 这一步【一次都没被真正触发过】

// dist(rng)              ∈ [-3, 3)      ← 复用上一行那个 ±3 的分布
// 0.5f * dist(rng)/3.0f  ∈ [-0.5, 0.5)  ← 实际就是 d/6
// 0.5f + ...             ∈ [0, 1)       ← 均匀，均值 0.5
hw[i] = 0.5f + 0.5f * dist(rng) / 3.0f;

// fused_ops.cu:570-593
// ---- occupancy：这一级理论上每个 SM 能同时跑几个 block ----
int blocksPerSM = 0;
const void *rep = (kind == K_N1) ? (const void *)rms_n1_pass1
                : (kind == K_N2) ? (const void *)rms_n2
                : …（八级各挑一个代表）
                                 : (const void *)sm_s4;
CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocksPerSM, rep, BLOCK, 0));

const double smFill = (double)M / (prop.multiProcessorCount * blocksPerSM);

printf("footprint = x + y = %.1f MB（L2 %.1f MB）%s\n",
       2.0 * bytes / 1048576.0, prop.l2CacheSize / 1048576.0,
       2.0 * bytes > prop.l2CacheSize ? "  → 装不下，必须走显存"
                                      : "  ⚠ 装得下 → 数字会虚高");

// RTX 4090（sm_89），BLOCK = 256
线程限    = 1536 / 256                       = 6
寄存器限  = 65536 / 每 block 占用的寄存器数
shared 限 = 102400 / 每 block 的 shared 字节

// 关键：寄存器按 warp 分配，而且要向上取整到 256 的倍数
每 warp 占 = 256 × ceil(寄存器数 × 32 / 256) = 256 × ceil(寄存器数 / 8)
每 block 占 = 8 个 warp × 每 warp 占

// 行号      kernel          ptxas info : Used …
 17   sm_s4          Used 48 registers, used 1 barriers, 72 bytes smem, 372 bytes cmem[0]
 22   sm_s3          Used 56 registers, used 1 barriers, 72 bytes smem, 372 bytes cmem[0]
 27   sm_s2_ms       Used 22 registers, used 1 barriers, 72 bytes smem, 380 bytes cmem[0]
 32   sm_s1_div      Used 24 registers, used 0 barriers,                388 bytes cmem[0]   ← 没有 smem
 37   sm_s1_sum      Used 40 registers, used 1 barriers, 32 bytes smem, 380 bytes cmem[0]
 42   sm_s1_max      Used 39 registers, used 1 barriers, 32 bytes smem, 372 bytes cmem[0]
 47   rms_n4         Used 55 registers, used 1 barriers, 32 bytes smem, 384 bytes cmem[0]
 52   rms_n3         Used 52 registers, used 1 barriers, 32 bytes smem, 384 bytes cmem[0]
 57   rms_n2         Used 39 registers, used 1 barriers, 32 bytes smem, 384 bytes cmem[0]
 62   rms_n1_pass2   Used 26 registers, used 0 barriers,                388 bytes cmem[0]   ← 没有 smem
 67   rms_n1_pass1   Used 40 registers, used 1 barriers, 32 bytes smem, 376 bytes cmem[0]

// BLOCK = 256 固定，把上面的式子化简
blocks_reg = floor( 32 / ceil(寄存器数 / 8) )

// fused_ops.cu —— 两个访存密集算子的融合阶梯：RMSNorm 与 online softmax
//
//   这份文件是第 10 章（Reduction）的实际应用。第 10 章练的是"把一个长数组归约成
//   一个标量"，这里练的是"把每一行归约成一个标量，然后立刻拿它去缩放这一行" ——
//   归约不再是目的，而是中间步骤，这才是它在真实推理栈里的样子。
//
//   ---- RMSNorm 阶梯（Qwen2.5 用的就是它，作用在 hidden=3584 上）----
//   n1   两趟两个 kernel：先算 sum(x²)，再归一化          —— 基线   12N 字节
//   n2   融成一个 kernel，但 x 仍然从 global 读两次        —— 只省 launch，不省字节  12N
//   n3   x 留在寄存器，读一次                              —— 这一步才动字节账   8N
//   n4   n3 + float4 向量化                                —— 只动条数账，字节不变  8N
//
//   ---- online softmax 阶梯（attention 打分和 logits 采样都要用）----
//   s1   三趟三个 kernel：max → sum(exp) → 归一化          —— 基线   16N 字节
//   s2   online 递推，两趟两个 kernel：(m,s) → 归一化      —— 省一趟  12N
//   s3   一趟，x 留寄存器                                  —— 8N
//   s4   s3 + float4                                       —— 8N
//
//   ⚠ n2 和 n1 的字节数【完全一样】。这是故意安排的一级台阶：
//     它用来证明"融合"本身不等于省带宽 —— 省带宽的是"让数据留在片上"，
//     而不是"把两个 kernel 写成一个"。这两件事经常被混为一谈。
//
//   评价指标是 MBU（Memory Bandwidth Utilization）= 实测 GB/s ÷ 理论峰值 GB/s。
//   访存密集算子的唯一目标就是把 MBU 顶上去 —— 和 SGEMM 看"占峰值算力"正好是
//   roofline 的两端。
//
//   编译:   nvcc -O3 -arch=sm_89 -lineinfo -Xptxas -v fused_ops.cu -o fused_ops
//   运行:   ./fused_ops <kernel> [M] [N] [ITERS]
//             M = 行数（token 数 / 打分矩阵的行数），N = 每行长度
//             RMSNorm 默认 M=4096  N=3584   （Qwen2.5-7B 的 hidden）
//             softmax 默认 M=4096  N=2048   （一段 2048 上下文的打分行）
//   查竞态: compute-sanitizer --tool racecheck ./fused_ops n3 8 3584 1
//
//   语法速查见 sgemm_all.cu 开头那一段，这里只标注新出现的写法。

#include <cuda_runtime.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <vector>
#include <algorithm>
#include <random>
#include <cfloat>

// ---------------------------------------------------------------------------
// 出错就地报，不要等到后面某个莫名其妙的地方才崩
// ---------------------------------------------------------------------------
#define CUDA_CHECK(call) do {                                                  \
    cudaError_t e_ = (call);                                                   \
    if (e_ != cudaSuccess) {                                                   \
        printf("[CUDA 错误] %s:%d  %s\n  ↳ %s\n",                              \
               __FILE__, __LINE__, cudaGetErrorString(e_), #call);             \
        exit(1);                                                               \
    }                                                                          \
} while (0)

// 每个 block 固定 256 线程。一个 block 负责一行。
static const int BLOCK = 256;
// 寄存器驻留版每线程最多存多少个 float。32 个 float = 32 个寄存器，
// 加上其余开销还远没到 255 的上限，不会 spill。
// 于是单行最长 = BLOCK × MAXS = 256 × 32 = 8192。
static const int MAXS = 32;
// float4 版：8 个 float4 也是 32 个 float，两条路径的寄存器压力刻意做成一样，
// 这样 n3→n4 / s3→s4 就是干净的单变量对照（只动指令条数，不动寄存器）。
static const int MAXV = 8;

// ===========================================================================
//                            归约积木
// ===========================================================================
//
// 和第 10 章 R6 的区别，就一处，但很关键：
//
//   R6 用的是 __shfl_down_sync —— 归约结果只汇聚到 0 号 lane，因为那里只需要
//   一个线程把结果 atomicAdd 出去，其余线程的值是垃圾无所谓。
//
//   这里用 __shfl_xor_sync（蝶形/butterfly）—— 归约完【每个 lane 手里都是全和】。
//   因为接下来 32 个线程都要拿这个和去缩放自己负责的那几个元素，
//   少一个线程知道都不行。
//
//   代价是一样的：同样 5 轮，同样每轮一条 shuffle 指令。xor 版不多花钱，
//   只是把"漏斗"换成了"全交换"。
// ---------------------------------------------------------------------------
__inline__ __device__ float warpReduceSum(float v) {
    // off 依次取 16 8 4 2 1；lane ^ off 是"异或伙伴"，
    // 第 1 轮 0↔16、1↔17……第 5 轮 0↔1、2↔3……五轮之后所有 lane 都持有全和。
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        v += __shfl_xor_sync(0xffffffffu, v, off);
    return v;
}

__inline__ __device__ float warpReduceMax(float v) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, off));
    return v;
}

// ---------------------------------------------------------------------------
// online softmax 的联合归约：把两个 (m, s) 合成一个
//
//   合并规则：
//       m = max(m1, m2)
//       s = s1·exp(m1 − m) + s2·exp(m2 − m)
//
//   为什么可以拿它做树形归约？因为这个运算【满足结合律】。
//   不满足的话，warp 内的蝶形归约就是错的 —— 它的合并顺序和串行扫描完全不同。
//   结合律的证明见导学 §5 · 三，一句话版本：
//   两个 (m,s) 合并后等价于"把两段数据拼起来重新算一遍 max 和 sum"，
//   而"拼起来"这个操作本身显然是可结合的。
// ---------------------------------------------------------------------------
// 注意这里【不能】用 -INFINITY 当初值。
//   没分到元素的线程 m 保持初值、s=0；两个这样的线程合并时
//   m - m_n = (-inf) - (-inf) = NaN，s = 0×NaN + 0×NaN = NaN，然后污染整行。
//   换成 -FLT_MAX：相减得 0，exp(0)=1，s = 0×1 + 0×1 = 0，安全退化。
__inline__ __device__ void warpReduceMaxSum(float &m, float &s) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        float m_o = __shfl_xor_sync(0xffffffffu, m, off);
        float s_o = __shfl_xor_sync(0xffffffffu, s, off);
        float m_n = fmaxf(m, m_o);
        // 两边都先缩放到新的 m_n 再相加。谁的 m 小，谁的 exp 就是个小于 1 的数，
        // 不会溢出；相等时两个 exp 都是 1，退化成普通求和。
        s = s * __expf(m - m_n) + s_o * __expf(m_o - m_n);
        m = m_n;
    }
}

// ---------------------------------------------------------------------------
// block 级归约：先 warp 内，再跨 warp。
//   BLOCK=256 → 8 个 warp。第一轮每个 warp 归出一个值放进 shared，
//   第二轮让 0 号 warp 把这 8 个值再归一次，最后广播给全 block。
// ---------------------------------------------------------------------------
__inline__ __device__ float blockReduceSum(float v, float *smem) {
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    v = warpReduceSum(v);
    if (lane == 0) smem[warp] = v;
    __syncthreads();
    // 只有前 8 个 lane 有真值，其余喂 0（加法的幺元）
    v = (threadIdx.x < (BLOCK >> 5)) ? smem[lane] : 0.0f;
    if (warp == 0) v = warpReduceSum(v);
    // 0 号 warp 算完了，但全 block 都要用，所以再过一次 shared 广播
    if (threadIdx.x == 0) smem[0] = v;
    __syncthreads();
    return smem[0];
}

__inline__ __device__ void blockReduceMaxSum(float &m, float &s, float *smem) {
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int nwarp = BLOCK >> 5;
    warpReduceMaxSum(m, s);
    if (lane == 0) { smem[warp] = m; smem[nwarp + warp] = s; }
    __syncthreads();
    if (warp == 0) {
        // 幺元是 (-inf, 0)：最大值取 -inf 不影响 max，和取 0 不影响求和
        m = (lane < nwarp) ? smem[lane]         : -FLT_MAX;
        s = (lane < nwarp) ? smem[nwarp + lane] : 0.0f;
        warpReduceMaxSum(m, s);
        if (lane == 0) { smem[2 * nwarp] = m; smem[2 * nwarp + 1] = s; }
    }
    __syncthreads();
    m = smem[2 * nwarp];
    s = smem[2 * nwarp + 1];
}

// ===========================================================================
//                            RMSNorm
//                y = x · rsqrt(mean(x²) + eps) · w
// ===========================================================================

// ---- N1：两趟，两个 kernel ------------------------------------------------
// 第一趟只算每行的 rstd，写进一个 M 长的小数组。
// 字节：读 x 一次 = 4N/行
__global__ void rms_n1_pass1(const float *__restrict__ x, float *__restrict__ rstd,
                             int N, float eps) {
    __shared__ float smem[BLOCK / 32];
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;

    float acc = 0.0f;
    for (int i = threadIdx.x; i < N; i += BLOCK) {
        float v = xr[i];
        acc += v * v;
    }
    acc = blockReduceSum(acc, smem);
    if (threadIdx.x == 0)
        rstd[row] = rsqrtf(acc / N + eps);
}

// 第二趟再把 x 读一遍，乘上第一趟算好的 rstd 和权重，写出去。
// 字节：读 x 一次 + 写 y 一次 = 8N/行   （合计 12N）
__global__ void rms_n1_pass2(const float *__restrict__ x, const float *__restrict__ w,
                             const float *__restrict__ rstd, float *__restrict__ y, int N) {
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    float       *yr = y + row * (long long)N;
    const float r = rstd[row];
    for (int i = threadIdx.x; i < N; i += BLOCK)
        yr[i] = xr[i] * r * w[i];
}

// ---- N2：融成一个 kernel，但 x 还是读两次 ---------------------------------
// 和 N1 的字节数【一模一样】。省掉的只有：一次 kernel launch、
// rstd 数组的一写一读（4M 字节，相对 12NM 可以忽略）。
// 留着这一级，就是为了让"融合 ≠ 省带宽"这件事有个可量化的对照。
__global__ void rms_n2(const float *__restrict__ x, const float *__restrict__ w,
                       float *__restrict__ y, int N, float eps) {
    __shared__ float smem[BLOCK / 32];
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    float       *yr = y + row * (long long)N;

    float acc = 0.0f;
    for (int i = threadIdx.x; i < N; i += BLOCK) {
        float v = xr[i];                       // ← 第一次读
        acc += v * v;
    }
    const float r = rsqrtf(blockReduceSum(acc, smem) / N + eps);
    for (int i = threadIdx.x; i < N; i += BLOCK)
        yr[i] = xr[i] * r * w[i];              // ← 第二次读，同样的地址
}

// ---- N3：x 留在寄存器，只读一次 -------------------------------------------
// 字节：读 x 一次 + 写 y 一次 = 8N/行。相对 N1/N2 的 12N，理论加速 1.5×。
// 代价是每线程多占 MAXS 个寄存器，且单行长度被 BLOCK×MAXS 卡死。
__global__ void rms_n3(const float *__restrict__ x, const float *__restrict__ w,
                       float *__restrict__ y, int N, float eps) {
    __shared__ float smem[BLOCK / 32];
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    float       *yr = y + row * (long long)N;

    // 循环必须写成"固定 MAXS 次 + 边界判断"，不能写成 for(i=tid; i<N; i+=BLOCK)
    // 配一个 xs[cnt++]。后者的下标是运行期算出来的，而【寄存器不能动态寻址】——
    // 编译器只好把 xs 放进 local memory（也就是显存），ptxas 会报 spill stores，
    // 这一级"数据留在片上"的前提当场作废。
    // 写成常量下标 xs[c] 并 #pragma unroll，编译器才会把它真正分配到寄存器里。
    float xs[MAXS];
    float acc = 0.0f;
    #pragma unroll
    for (int c = 0; c < MAXS; ++c) {
        int i = threadIdx.x + c * BLOCK;
        float v = (i < N) ? xr[i] : 0.0f;      // 越界的线程喂 0，不影响平方和
        xs[c] = v;
        acc += v * v;
    }
    const float r = rsqrtf(blockReduceSum(acc, smem) / N + eps);
    #pragma unroll
    for (int c = 0; c < MAXS; ++c) {
        int i = threadIdx.x + c * BLOCK;
        if (i < N) yr[i] = xs[c] * r * w[i];
    }
}

// ---- N4：N3 + float4 -------------------------------------------------------
// 字节数和 N3 完全相同，变的只有指令条数和内存请求数：
// 4 个 32 位 load 并成 1 个 128 位 load。
__global__ void rms_n4(const float *__restrict__ x, const float *__restrict__ w,
                       float *__restrict__ y, int N, float eps) {
    __shared__ float smem[BLOCK / 32];
    const long long row = blockIdx.x;
    const int n4 = N >> 2;                     // 一行有多少个 float4

    // reinterpret_cast 把 float* 重新解释成 float4*，地址必须 16 字节对齐 ——
    // cudaMalloc 返回的指针是 256 字节对齐的，而 row*N 只要 N%4==0 就仍然对齐。
    const float4 *xr = reinterpret_cast<const float4 *>(x + row * (long long)N);
    float4       *yr = reinterpret_cast<float4 *>(y + row * (long long)N);
    const float4 *wr = reinterpret_cast<const float4 *>(w);

    float4 xs[MAXV];
    float acc = 0.0f;
    #pragma unroll
    for (int c = 0; c < MAXV; ++c) {
        int i = threadIdx.x + c * BLOCK;
        float4 v = (i < n4) ? xr[i] : make_float4(0.f, 0.f, 0.f, 0.f);
        xs[c] = v;
        acc += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
    }
    const float r = rsqrtf(blockReduceSum(acc, smem) / N + eps);
    #pragma unroll
    for (int c = 0; c < MAXV; ++c) {
        int i = threadIdx.x + c * BLOCK;
        if (i >= n4) continue;
        float4 v = xs[c], g = wr[i], o;
        o.x = v.x * r * g.x;  o.y = v.y * r * g.y;
        o.z = v.z * r * g.z;  o.w = v.w * r * g.w;
        yr[i] = o;
    }
}

// ===========================================================================
//                          online softmax
//                  y_i = exp(x_i − max) / Σ exp(x_j − max)
// ===========================================================================

// ---- S1：三趟，三个 kernel -------------------------------------------------
// 字节：4N + 4N + (4N+4N) = 16N/行
__global__ void sm_s1_max(const float *__restrict__ x, float *__restrict__ mx, int N) {
    __shared__ float smem[BLOCK / 32];
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    float m = -FLT_MAX;
    for (int i = threadIdx.x; i < N; i += BLOCK) m = fmaxf(m, xr[i]);

    const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
    m = warpReduceMax(m);
    if (lane == 0) smem[warp] = m;
    __syncthreads();
    m = (threadIdx.x < BLOCK / 32) ? smem[lane] : -FLT_MAX;
    if (warp == 0) m = warpReduceMax(m);
    if (threadIdx.x == 0) mx[row] = m;
}

__global__ void sm_s1_sum(const float *__restrict__ x, const float *__restrict__ mx,
                          float *__restrict__ sm, int N) {
    __shared__ float smem[BLOCK / 32];
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    const float m = mx[row];
    float s = 0.0f;
    for (int i = threadIdx.x; i < N; i += BLOCK) s += __expf(xr[i] - m);
    s = blockReduceSum(s, smem);
    if (threadIdx.x == 0) sm[row] = s;
}

__global__ void sm_s1_div(const float *__restrict__ x, const float *__restrict__ mx,
                          const float *__restrict__ sm, float *__restrict__ y, int N) {
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    float       *yr = y + row * (long long)N;
    const float m = mx[row], inv = 1.0f / sm[row];
    for (int i = threadIdx.x; i < N; i += BLOCK)
        yr[i] = __expf(xr[i] - m) * inv;
}

// ---- S2：online 递推，两趟 -------------------------------------------------
// 第一趟【一次扫描同时得到 max 和 sum】—— 这就是 online softmax 的全部内容。
// 字节：4N + (4N+4N) = 12N/行
__global__ void sm_s2_ms(const float *__restrict__ x, float *__restrict__ mx,
                         float *__restrict__ sm, int N) {
    __shared__ float smem[BLOCK / 32 * 2 + 2];
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;

    float m = -FLT_MAX, s = 0.0f;
    for (int i = threadIdx.x; i < N; i += BLOCK) {
        float v = xr[i];
        // 递推三行：见到更大的 max，就把已经累好的 s 按比例缩回去
        float m_n = fmaxf(m, v);
        s = s * __expf(m - m_n) + __expf(v - m_n);
        m = m_n;
    }
    blockReduceMaxSum(m, s, smem);
    if (threadIdx.x == 0) { mx[row] = m; sm[row] = s; }
}

// ---- S3：一趟，x 留寄存器 ---------------------------------------------------
// 字节：4N + 4N = 8N/行。相对 S1 的 16N，理论加速 2.0×。
__global__ void sm_s3(const float *__restrict__ x, float *__restrict__ y, int N) {
    __shared__ float smem[BLOCK / 32 * 2 + 2];
    const long long row = blockIdx.x;
    const float *xr = x + row * (long long)N;
    float       *yr = y + row * (long long)N;

    // 常量下标 + #pragma unroll，理由同 rms_n3
    float xs[MAXS];
    float m = -FLT_MAX, s = 0.0f;
    #pragma unroll
    for (int c = 0; c < MAXS; ++c) {
        int i = threadIdx.x + c * BLOCK;
        if (i >= N) { xs[c] = -FLT_MAX; continue; }  // 越界位置塞 -FLT_MAX：
        float v = xr[i];                             // 既不会成为 max，
        xs[c] = v;                                   // exp 出来也是 0
        float m_n = fmaxf(m, v);
        s = s * __expf(m - m_n) + __expf(v - m_n);
        m = m_n;
    }
    blockReduceMaxSum(m, s, smem);
    const float inv = 1.0f / s;
    #pragma unroll
    for (int c = 0; c < MAXS; ++c) {
        int i = threadIdx.x + c * BLOCK;
        if (i < N) yr[i] = __expf(xs[c] - m) * inv;
    }
}

// ---- S4：S3 + float4 -------------------------------------------------------
__global__ void sm_s4(const float *__restrict__ x, float *__restrict__ y, int N) {
    __shared__ float smem[BLOCK / 32 * 2 + 2];
    const long long row = blockIdx.x;
    const int n4 = N >> 2;
    const float4 *xr = reinterpret_cast<const float4 *>(x + row * (long long)N);
    float4       *yr = reinterpret_cast<float4 *>(y + row * (long long)N);

    float4 xs[MAXV];
    float m = -FLT_MAX, s = 0.0f;
    #pragma unroll
    for (int c = 0; c < MAXV; ++c) {
        int i = threadIdx.x + c * BLOCK;
        if (i >= n4) { xs[c] = make_float4(-FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX); continue; }
        float4 v = xr[i];
        xs[c] = v;
        // 一次进来 4 个元素，先在寄存器里把这 4 个合成一个 (m,s)，再并进累计值。
        // 这已经是 FlashAttention 分块合并的最小形态了 —— 只不过块大小是 4。
        float mv = fmaxf(fmaxf(v.x, v.y), fmaxf(v.z, v.w));
        float sv = __expf(v.x - mv) + __expf(v.y - mv) + __expf(v.z - mv) + __expf(v.w - mv);
        float m_n = fmaxf(m, mv);
        s = s * __expf(m - m_n) + sv * __expf(mv - m_n);
        m = m_n;
    }
    blockReduceMaxSum(m, s, smem);
    const float inv = 1.0f / s;
    #pragma unroll
    for (int c = 0; c < MAXV; ++c) {
        int i = threadIdx.x + c * BLOCK;
        if (i >= n4) continue;
        float4 v = xs[c], o;
        o.x = __expf(v.x - m) * inv;  o.y = __expf(v.y - m) * inv;
        o.z = __expf(v.z - m) * inv;  o.w = __expf(v.w - m) * inv;
        yr[i] = o;
    }
}

// ===========================================================================
//                         CPU 参照与抽样校验
// ===========================================================================
enum Kind { K_N1, K_N2, K_N3, K_N4, K_S1, K_S2, K_S3, K_S4 };
static bool isNorm(Kind k) { return k <= K_N4; }

// 只抽查若干行。全查的话 M=4096、N=3584 要在 CPU 上算 1400 万次 exp，太慢。
static int verifyRows(Kind kind, int M, int N, const std::vector<float> &hx,
                      const std::vector<float> &hw, const std::vector<float> &hy,
                      int nrows) {
    std::mt19937 rng(20260925);
    int bad = 0;
    for (int t = 0; t < nrows; ++t) {
        long long row = rng() % (unsigned)M;
        const float *xr = hx.data() + row * (long long)N;
        const float *yr = hy.data() + row * (long long)N;

        std::vector<double> ref(N);
        if (isNorm(kind)) {
            double acc = 0.0;
            for (int i = 0; i < N; ++i) acc += (double)xr[i] * xr[i];
            double r = 1.0 / std::sqrt(acc / N + 1e-6);
            for (int i = 0; i < N; ++i) ref[i] = xr[i] * r * hw[i];
        } else {
            double m = -1e300;
            for (int i = 0; i < N; ++i) m = std::max(m, (double)xr[i]);
            double s = 0.0;
            for (int i = 0; i < N; ++i) s += std::exp(xr[i] - m);
            for (int i = 0; i < N; ++i) ref[i] = std::exp(xr[i] - m) / s;
        }
        // __expf / rsqrtf 是快速近似指令，容差要比普通浮点宽一点。
        // 这里用相对误差 2e-3，softmax 的绝对值很小，再加一个 1e-7 的绝对兜底。
        for (int i = 0; i < N; ++i) {
            double d = std::fabs(ref[i] - yr[i]);
            if (d > 2e-3 * std::fabs(ref[i]) + 1e-7) {
                if (bad < 3)
                    printf("  行 %lld 列 %d: GPU %.6g  CPU %.6g\n", row, i, yr[i], ref[i]);
                ++bad;
            }
        }
    }
    return bad;
}

// ===========================================================================
//                                main
// ===========================================================================
int main(int argc, char **argv) {
    if (argc < 2) {
        printf("用法: %s <kernel> [M] [N] [ITERS]\n"
               "  kernel = n1|n2|n3|n4   RMSNorm     默认 M=4096 N=3584\n"
               "           s1|s2|s3|s4   softmax     默认 M=4096 N=2048\n"
               "  M = 行数（token 数 / 打分矩阵的行数），N = 每行长度\n"
               "  例： %s n4 4096 3584 11      Qwen2.5-7B 的一次 RMSNorm\n"
               "       %s n4 1 3584 11         decode batch=1（只有 1 个 block！）\n"
               "       %s s4 114688 2048 11    28 头 × 4096 个 query 的打分\n",
               argv[0], argv[0], argv[0], argv[0]);
        return 2;
    }
    const char *name = argv[1];
    Kind kind;
    if      (!strcmp(name, "n1")) kind = K_N1;
    else if (!strcmp(name, "n2")) kind = K_N2;
    else if (!strcmp(name, "n3")) kind = K_N3;
    else if (!strcmp(name, "n4")) kind = K_N4;
    else if (!strcmp(name, "s1")) kind = K_S1;
    else if (!strcmp(name, "s2")) kind = K_S2;
    else if (!strcmp(name, "s3")) kind = K_S3;
    else if (!strcmp(name, "s4")) kind = K_S4;
    else { printf("不认识的 kernel: %s\n", name); return 2; }

    int M     = (argc > 2) ? atoi(argv[2]) : 4096;
    int N     = (argc > 3) ? atoi(argv[3]) : (isNorm(kind) ? 3584 : 2048);
    int ITERS = (argc > 4) ? atoi(argv[4]) : 11;
    if (M <= 0 || N <= 0) { printf("M/N 必须为正\n"); return 2; }

    cudaDeviceProp prop;
    CUDA_CHECK(cudaGetDeviceProperties(&prop, 0));
    int memClkKHz = 0, gpuClkKHz = 0;
    cudaDeviceGetAttribute(&memClkKHz, cudaDevAttrMemoryClockRate, 0);
    cudaDeviceGetAttribute(&gpuClkKHz, cudaDevAttrClockRate, 0);
    // 峰值带宽 = 内存时钟 × 2（DDR 双沿）× 位宽/8。和 sgemm_all.cu 用同一个口径。
    const double peakBW = memClkKHz * 2.0 * (prop.memoryBusWidth / 8) / 1.0e6;

    printf("GPU: %s  sm_%d%d  %d SM  L2 %.1f MB   显存带宽峰值 %.0f GB/s\n",
           prop.name, prop.major, prop.minor, prop.multiProcessorCount,
           prop.l2CacheSize / 1048576.0, peakBW);
    printf("[%s] %s   M=%d 行 × N=%d 列   block %d 线程   grid %d 个 block\n",
           name, isNorm(kind) ? "RMSNorm" : "softmax", M, N, BLOCK, M);

    // ---- 两个会让结果作废的前提，先查 ----
    const bool needReg = (kind == K_N3 || kind == K_N4 || kind == K_S3 || kind == K_S4);
    if (needReg && N > BLOCK * MAXS) {
        printf("✗ 这一级把整行存在寄存器里，单行最长 %d，而 N=%d。\n",
               BLOCK * MAXS, N);
        printf("  这不是实现偷懒，是硬件约束 —— 也正是 FlashAttention 必须分块的原因，\n"
               "  详见导学 §7 · 三。想跑长行请用 n1/n2 或 s1/s2。\n");
        return 3;
    }
    const bool needVec4 = (kind == K_N4 || kind == K_S4);
    if (needVec4 && (N & 3)) {
        printf("✗ float4 版要求 N 是 4 的倍数（当前 N=%d）\n", N);
        return 3;
    }

    // ---- 数据 ----
    const size_t nElem = (size_t)M * N;
    const size_t bytes = nElem * sizeof(float);
    std::vector<float> hx(nElem), hw(N), hy(nElem);
    std::mt19937 rng(12345);
    std::uniform_real_distribution<float> dist(-3.0f, 3.0f);
    for (size_t i = 0; i < nElem; ++i) hx[i] = dist(rng);
    for (int i = 0; i < N; ++i)        hw[i] = 0.5f + 0.5f * dist(rng) / 3.0f;

    float *dx = nullptr, *dw = nullptr, *dy = nullptr, *dm = nullptr, *ds = nullptr;
    CUDA_CHECK(cudaMalloc(&dx, bytes));
    CUDA_CHECK(cudaMalloc(&dy, bytes));
    CUDA_CHECK(cudaMalloc(&dw, N * sizeof(float)));
    CUDA_CHECK(cudaMalloc(&dm, M * sizeof(float)));
    CUDA_CHECK(cudaMalloc(&ds, M * sizeof(float)));
    CUDA_CHECK(cudaMemcpy(dx, hx.data(), bytes, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(dw, hw.data(), N * sizeof(float), cudaMemcpyHostToDevice));

    const float eps = 1e-6f;
    // lambda：把"发射这一级的全部 kernel"打包成一个可重复调用的东西。
    // [&] 表示按引用捕获外面所有用到的变量。
    auto launch = [&]() {
        switch (kind) {
            case K_N1: rms_n1_pass1<<<M, BLOCK>>>(dx, dm, N, eps);
                       rms_n1_pass2<<<M, BLOCK>>>(dx, dw, dm, dy, N);          break;
            case K_N2: rms_n2<<<M, BLOCK>>>(dx, dw, dy, N, eps);               break;
            case K_N3: rms_n3<<<M, BLOCK>>>(dx, dw, dy, N, eps);               break;
            case K_N4: rms_n4<<<M, BLOCK>>>(dx, dw, dy, N, eps);               break;
            case K_S1: sm_s1_max<<<M, BLOCK>>>(dx, dm, N);
                       sm_s1_sum<<<M, BLOCK>>>(dx, dm, ds, N);
                       sm_s1_div<<<M, BLOCK>>>(dx, dm, ds, dy, N);             break;
            case K_S2: sm_s2_ms <<<M, BLOCK>>>(dx, dm, ds, N);
                       sm_s1_div<<<M, BLOCK>>>(dx, dm, ds, dy, N);             break;
            case K_S3: sm_s3<<<M, BLOCK>>>(dx, dy, N);                         break;
            case K_S4: sm_s4<<<M, BLOCK>>>(dx, dy, N);                         break;
        }
    };

    // ---- occupancy：这一级理论上每个 SM 能同时跑几个 block ----
    int blocksPerSM = 0;
    const void *rep = (kind == K_N1) ? (const void *)rms_n1_pass1
                    : (kind == K_N2) ? (const void *)rms_n2
                    : (kind == K_N3) ? (const void *)rms_n3
                    : (kind == K_N4) ? (const void *)rms_n4
                    : (kind == K_S1) ? (const void *)sm_s1_max
                    : (kind == K_S2) ? (const void *)sm_s2_ms
                    : (kind == K_S3) ? (const void *)sm_s3
                                     : (const void *)sm_s4;
    CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocksPerSM, rep, BLOCK, 0));
    printf("occupancy: %d block × %d 线程 = %d / %d = %.1f%%",
           blocksPerSM, BLOCK, blocksPerSM * BLOCK, prop.maxThreadsPerMultiProcessor,
           100.0 * blocksPerSM * BLOCK / prop.maxThreadsPerMultiProcessor);
    // grid 太小的话，占用率再高也没用 —— 卡根本没被填满。这是 decode 的常态。
    const double smFill = (double)M / (prop.multiProcessorCount * blocksPerSM);
    printf("   grid 填充度 %d / (%d SM × %d) = %.2f%s\n",
           M, prop.multiProcessorCount, blocksPerSM, smFill,
           smFill < 1.0 ? "   ⚠ 填不满，卡在空转" : "");

    printf("footprint = x + y = %.1f MB（L2 %.1f MB）%s\n",
           2.0 * bytes / 1048576.0, prop.l2CacheSize / 1048576.0,
           2.0 * bytes > prop.l2CacheSize ? "  → 装不下，必须走显存"
                                          : "  ⚠ 装得下 → 数字会虚高");

    // ---- 计时：和 sgemm_all.cu 同一套（预热 + 中位数）----
    const int WARMUP = (ITERS <= 1) ? 0 : (ITERS <= 2 ? 1 : 3);
    for (int i = 0; i < WARMUP; ++i) launch();
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaDeviceSynchronize());

    cudaEvent_t ev0, ev1;
    CUDA_CHECK(cudaEventCreate(&ev0));
    CUDA_CHECK(cudaEventCreate(&ev1));
    std::vector<float> samples;
    for (int it = 0; it < ITERS; ++it) {
        CUDA_CHECK(cudaEventRecord(ev0));
        launch();
        CUDA_CHECK(cudaEventRecord(ev1));
        CUDA_CHECK(cudaEventSynchronize(ev1));
        float t; CUDA_CHECK(cudaEventElapsedTime(&t, ev0, ev1));
        samples.push_back(t);
    }
    std::sort(samples.begin(), samples.end());
    const size_t ns = samples.size();
    const double ms = (ns % 2) ? samples[ns / 2]
                               : 0.5 * (samples[ns / 2 - 1] + samples[ns / 2]);

    CUDA_CHECK(cudaMemcpy(hy.data(), dy, bytes, cudaMemcpyDeviceToHost));
    const int bad = verifyRows(kind, M, N, hx, hw, hy, 8);
    printf("%s（抽查 8 行）\n", bad == 0 ? "✓ 抽样校验通过" : "✗ 抽样校验失败");

    // ---- 字节账：这是这两个算子唯一重要的一本账 ----
    // 每行搬多少字节，按"必须经过显存的量"算。权重 w 只有 4N 字节，
    // 所有行共用、稳稳待在 L2 里，不计入。
    double bytesPerRow;
    switch (kind) {
        case K_N1: case K_N2: bytesPerRow = 12.0 * N; break;   // 读两次 + 写一次
        case K_N3: case K_N4: bytesPerRow =  8.0 * N; break;   // 读一次 + 写一次
        case K_S1:            bytesPerRow = 16.0 * N; break;   // 读三次 + 写一次
        case K_S2:            bytesPerRow = 12.0 * N; break;
        default:              bytesPerRow =  8.0 * N; break;
    }
    const double moved = bytesPerRow * M;
    const double gbps  = moved / (ms / 1000.0) / 1.0e9;
    const double mbu   = 100.0 * gbps / peakBW;

    printf("中位耗时 %.4f ms   搬运 %.1f MB   %.1f GB/s   MBU %.1f%%\n",
           ms, moved / 1048576.0, gbps, mbu);
    printf("字节账: %.0fN 字节/行 × %d 行 = %.1f MB；"
           "理论下限 %.4f ms（按 %.0f GB/s 跑满）\n",
           bytesPerRow / N, M, moved / 1048576.0,
           moved / (peakBW * 1e9) * 1000.0, peakBW);

    CUDA_CHECK(cudaFree(dx)); CUDA_CHECK(cudaFree(dy)); CUDA_CHECK(cudaFree(dw));
    CUDA_CHECK(cudaFree(dm)); CUDA_CHECK(cudaFree(ds));
    return bad == 0 ? 0 : 1;
}

三、compute-sanitizer 在查什么：读懂 racecheck 日志↑

compute-sanitizer 是 CUDA 自带的运行时检查工具。 它照常运行你的程序，同时监控每一条访存指令。

这里用的是它的 racecheck 工具，只查一类错误：shared memory 上的竞态。 同时满足下面三个条件，就是一处竞态：

• 同一个 block 里的两个线程，访问同一个 shared memory 地址；
• 两次访问中，至少有一次是写；
• 两次访问之间，没有 __syncthreads()。

一、为什么需要它：抽样校验抓不住竞态

漏写 __syncthreads() 不会让程序崩溃，也不会报错。 它只会让结果偏一点点。偏多少取决于 warp 的执行顺序。 换一个形状或换一张卡，结果可能就不同。 所以「✓ 抽样校验通过」不能证明没有竞态。 这一次结果对，可能只是执行顺序碰巧合适。

racecheck 不看结果对不对，只看访问顺序。 只要访问顺序满足上面三个条件，它就报告一处竞态，哪怕这次读到的值碰巧正确。 racecheck 把每一处竞态记为一个 hazard，按读写的先后分三类：

二、对照源码：racecheck 检查哪几行

下面是 fused_ops.cu 里的 blockReduceSum，n3 用的就是它。 注释说明 racecheck 怎样看每一行。 蓝色是 shared memory 访问； 红色是 __syncthreads()，删掉任何一个，racecheck 都会报告竞态。

__inline__ __device__ float blockReduceSum(float v, float *smem) {
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    v = warpReduceSum(v);                // warp 内用 shuffle 交换，数据在寄存器里，racecheck 不检查
    if (lane == 0) smem[warp] = v;       // ② 写：8 个线程各写一个槽 smem[0..7]，地址不重叠，没有 WAW
    __syncthreads();                     // ③ 同步：删掉它，smem[1..7] 报 RAW
                                         //    （0 号 warp 读的时候，7 号 warp 可能还没写）
    v = (threadIdx.x < (BLOCK >> 5)) ? smem[lane] : 0.0f;
                                         // ④ 读：只有 0 号 warp 的前 8 个线程读
    if (warp == 0) v = warpReduceSum(v);
    if (threadIdx.x == 0) smem[0] = v;   // ⑤ 写回 smem[0]：④ 里读 smem[0] 的也是 0 号线程
                                         //    同一个线程先读后写，不算竞态，所以这里不需要同步
    __syncthreads();                     // ⑥ 同步：删掉它，smem[0] 报 RAW
                                         //    （其余 255 个线程可能读到 ② 写入的部分和）
    return smem[0];                      // ⑦ 读：256 个线程读同一个地址，只读不写，没有竞态
}

三、逐行读 fused_20261001_005225.txt 的 racecheck 段

注意：只看最后一行 RACECHECK SUMMARY。 racecheck 运行时，程序打印的耗时、GB/s 和 MBU 都没有意义。

--- racecheck · n3 ---
$ compute-sanitizer --tool racecheck ./fused_ops n3 8 1000 1
========= COMPUTE-SANITIZER                // 以 ========= 开头的行由 sanitizer 打印；其余行是程序自己的输出
GPU: NVIDIA GeForce RTX 4090  sm_89  128 SM  L2 72.0 MB   显存带宽峰值 1008 GB/s
[n3] RMSNorm   M=8 行 × N=1000 列   block 256 线程   grid 8 个 block
occupancy: 4 block × 256 线程 = 1024 / 1536 = 66.7%   grid 填充度 8 / (128 SM × 4) = 0.02   ⚠ 填不满，卡在空转
                                           // 这里的 ⚠ 是预期的：这一步故意用很小的形状，不测性能
footprint = x + y = 0.1 MB（L2 72.0 MB）  ⚠ 装得下 → 数字会虚高
✓ 抽样校验通过（抽查 8 行）                // 程序自己的正确性校验，和 racecheck 无关
中位耗时 2.3409 ms   搬运 0.1 MB   0.0 GB/s   MBU 0.0%
                                           // 正常应是微秒级；插桩让程序慢了几百倍，这个数字没有意义
字节账: 8N 字节/行 × 8 行 = 0.1 MB；理论下限 0.0001 ms（按 1008 GB/s 跑满）
========= RACECHECK SUMMARY: 0 hazards displayed (0 errors, 0 warnings)
                                           // 只看这一行：0 hazards 表示没有发现竞态

s3 和 s4 两段的格式相同，结尾也都是 0 hazards。 这满足验收清单里「三段 racecheck 都要 0 hazards」的要求。

如果有人删掉一个 __syncthreads()，SUMMARY 前面会多出几条报告。 每条报告写明竞态类型、shared memory 地址，以及读写两侧的源码位置。 编译时要加 -lineinfo，报告才能定位到行号。SUMMARY 里的 hazard 数也不再是 0。

四、为什么只测三个 kernel，为什么形状这么小

小结：三项检查各回答一个问题。

• 计时：kernel 快不快。
• 抽样校验：这一次结果对不对。
• racecheck：结果对，是不是只因为执行顺序碰巧合适。

三段都是 0 hazards，说明 blockReduceSum 和 blockReduceMaxSum 里的 __syncthreads() 位置正确。

四、run_fused.sh 全文↑

一条命令跑完八个 kernel × 五种形状 × 三轮，外加边界、长行拒绝和 racecheck。 这次实跑的产物就是 §9 那三张表。

• 脚本开头有一段 shell 语法速查。 它讲脚本用到的写法：${VAR:-默认}、$( )、"$@" 与 "$*"、 2>&1、set -u、set --、command -v 等。 正文每段开头也标了该段做什么。不熟悉 shell 的话，先读那段，再读正文。
• 全脚本只有一处故意不加引号： 形状扫描那段的 set -- $shp。目的是让 "prefill 2048" 按空格拆成两个词。注释里也标了。

#!/usr/bin/env bash
# =============================================================================
# run_fused.sh —— RMSNorm / online softmax 两条融合阶梯的完整测试
#
#   八个 kernel：n1 n2 n3 n4（RMSNorm）+ s1 s2 s3 s4（softmax）
#
#   和 SGEMM 那套的区别：这两个算子是【访存密集】的，
#   所以看的不是"占峰值算力"，是 MBU = 实测 GB/s ÷ 1008。
#
#   阶段：
#     0 编译 + spill 检查   ← 这一步比 SGEMM 那边更要命，见下面
#     1 正确性（含边界）
#     2 主阶梯（固定形状，看 1.5× 和 2.0× 兑现了多少）
#     3 形状扫描（decode b1 / b32 / prefill / 大批）
#     4 长行：寄存器放不下会怎样
#     5 竞态检查
#
#   ⚠ 第 0 步的 spill 检查是这份测试的命门：
#     n3/n4/s3/s4 的全部立论就是"整行留在寄存器里，只读一次 global"。
#     一旦 ptxas 报 spill stores，那行数据其实被挪到了 local memory（也就是显存），
#     于是它们退化成"读两次"，1.5× 和 2.0× 都不会出现，而且你从耗时上看不出原因。
#
#   用法：bash run_fused.sh
#         ROUNDS=1 bash run_fused.sh      # 快速过一遍
# =============================================================================
#
# ---- 给不熟悉 shell 的读者：这份脚本用到的语法，一次讲清 --------------------
#
#   #!/usr/bin/env bash  第一行的 #! 叫 shebang，告诉系统拿哪个解释器跑这个文件。
#                      写 env bash 而不是写死 /bin/bash，是为了用 PATH 里找到的那个。
#   #                  行内从 # 到行尾都是注释（shebang 那一行除外）。
#   $VAR / ${VAR}      取变量的值。${} 是为了和紧跟的字符分开，
#                      比如 "fused_${STAMP}.txt" —— 不加花括号会被当成变量名 STAMP.txt。
#   "$VAR"             几乎永远要加双引号 —— 不加的话，值里有空格就会被拆成多个参数。
#                      下面只有一处故意不加（set -- $shp），那里就是要它拆开。
#   ${VAR:-默认值}      VAR 没设置或为空时取「默认值」。ROUNDS=${ROUNDS:-3} 的意思是：
#                      允许调用者用 ROUNDS=1 bash run_fused.sh 从环境覆盖它。
#   $(命令)             命令替换：先跑命令，把它打印的内容替换到这里。
#   名字() { ...; }     定义函数。调用时直接写名字，第 1、2 个参数是 $1 $2 …
#                      函数体写在一行时，最后一条命令后面也要有分号，} 前要有空格。
#   "$@"  与  "$*"     都表示"全部参数"，但用途相反：
#                      "$@" 保持每个参数各自独立 → 拿来【执行】；
#                      "$*" 拼成一整串           → 拿来【打印】。下面 run() 里两个都用到了。
#   cmd1 | cmd2        管道：把 cmd1 的输出喂给 cmd2 的输入。
#   >  文件            把输出写进文件（覆盖）；  >>  是追加。
#   2>/dev/null        把「错误输出」丢掉（1 是正常输出，2 是错误输出）。
#   2>&1               把错误输出并进正常输出，这样管道和重定向才收得到它。
#   \$                 反斜杠转义，输出一个字面的 $ 符号，而不是去取变量。
#   for x in 一串词     按空格逐个取出。词可以来自字面量、变量展开或 $( )。
#   do … done          for / while 的循环体；if 用 then … fi 收尾。
#   A || { B; }        A 失败（返回非 0）时才执行 B。A && B 则相反。
#                      { } 把多条命令打包成一组（注意最后一条也要有分号）。
#   set -u             用到没定义过的变量就立刻报错退出，防手滑打错变量名。
#   set -- 词 词       重设位置参数：执行后 $1 $2 就变成这些词。
#                      下面拿它把 "prefill 2048" 这种「名字+数值」的配对拆成两个变量。
#   command -v xxx     查「系统里有没有 xxx 这个命令」，输出丢掉只看成败。
#   exit 1             中止脚本并返回非 0（约定：0 成功，非 0 失败）。
#
#   ---- 被测程序的调用约定 ----
#   ./fused_ops <kind> <M> <N> <iters>
#       kind   n1|n2|n3|n4|s1|s2|s3|s4
#       M      行数（= batch × seq，每行一个 block）
#       N      列数（RMSNorm 用 hidden=3584，softmax 用 seq=2048）
#       iters  计时重复次数，程序取中位数
#
# =============================================================================
set -u

# ---- 可调参数：都用 ${VAR:-默认} 写，于是能从命令行环境覆盖 ----
ROUNDS=${ROUNDS:-3}
HID=${HID:-3584}                          # Qwen2.5-7B 的 hidden_size
SEQ=${SEQ:-2048}                          # 打分行长度（上下文长度）
# date +格式 生成时间戳，$( ) 把它取出来拼进文件名 → 每次运行各自一份输出
STAMP=$(date +%Y%m%d_%H%M%S)
OUT="fused_${STAMP}.txt"                  # ${STAMP} 的花括号是为了和后面的 .txt 分开

# ---- 三个小函数 ----
# log：  "$@" = 调用时传进来的全部参数（保留原有的分词）；
#        tee 同时写屏幕和文件，-a 表示追加（没有 -a 会每次清空）
# rule： 固定打一条横线，纯粹为了日志好读
# run：  先把命令行原样记进日志（方便复现），再真正执行它。函数体是两条命令，用 ; 隔开：
#          echo "\$ $*" >>"$OUT"     ← 记日志的是这半句。\$ 输出字面的 $ 符号（模仿终端提示符），
#                                      $* 把全部参数拼成一行；>> 只写文件、不上屏幕
#          "$@" 2>&1 | tee -a "$OUT" ← 这半句才是真执行。"$@" 保持每个参数独立
#                                      （$* 会拼成一整串，不能拿来执行）；2>&1 让报错也进日志
log()  { echo "$@" | tee -a "$OUT"; }
rule() { log "-------------------------------------------------------------"; }
run()  { echo "\$ $*" >>"$OUT"; "$@" 2>&1 | tee -a "$OUT"; }

# ---------------------------------------------------------------- 环境快照
# date -Is 打 ISO 格式时间；hostname 机器名；uname -sr 内核名和版本
log "============================================================="
log " 融合算子基准 · RMSNorm + online softmax   $(date -Is)"
log " 主机 $(hostname)   $(uname -sr)"
log " hidden=$HID  seq=$SEQ  （Qwen2.5-7B 的真实维度）"
log "============================================================="
# 记下卡型号和功耗墙 —— 换机器重跑时，这几行决定数字能不能对比
run nvidia-smi --query-gpu=name,driver_version,clocks.max.sm,clocks.max.mem,power.limit --format=csv

# ---------------------------------------------------------------- 0 编译
rule
log "编译 + spill 检查"
rule
# -Xptxas -v 让 ptxas 把每个 kernel 的寄存器数 / 栈帧 / spill 打出来 —— 这一步的全部目的
# ||  前面失败才执行后面的 { }；{ } 是把多条命令打包成一组（注意最后一条也要有分号）
run nvcc -O3 -arch=sm_89 -lineinfo -Xptxas -v fused_ops.cu -o fused_ops || {
  log "!! 编译失败，停在这里"; exit 1; }
log ""
# log "" 打一个空行；下面这几行 >>> 是写给读日志的人看的验收提示，不参与计算
log ">>> 十一个 kernel 全都要是 0 bytes stack frame / 0 bytes spill stores。"
log ">>> 尤其盯住 rms_n3 / rms_n4 / sm_s3 / sm_s4 —— 它们每线程要拿 32 个寄存器"
log ">>> 存整行数据。一旦 spill，那行数据就跑到显存里去了，"
log ">>> '只读一次' 的前提作废，1.5× / 2.0× 一个都不会出现。"
log ">>> 顺带记下寄存器数：n1/n2 应该很少（十几个），n3/n4/s3/s4 应该明显更多。"

# ---------------------------------------------------------------- 1 正确性
rule
log "正确性：八个 kernel 各跑一次小规模"
rule
# for k in n1 n2 n3 n4 —— 直接给字面量列表，$k 依次取到每一个
# M 只取 64：这一步看的是校验行，不是速度，跑快点好
for k in n1 n2 n3 n4; do
  log ""; log "--- $k · M=64 N=$HID ---"
  run ./fused_ops "$k" 64 "$HID" 3
done
# softmax 这四个换成 N=$SEQ —— 行长不同，两条阶梯不能混在同一个 N 上比
for k in s1 s2 s3 s4; do
  log ""; log "--- $k · M=64 N=$SEQ ---"
  run ./fused_ops "$k" 64 "$SEQ" 3
done
log ""
log ">>> 八行都要看到 ✓ 抽样校验通过。"
log ">>> softmax 那四个尤其要看 —— online 递推写错的典型症状不是崩溃，"
log ">>> 是结果偏一点点，只有对照 CPU 才看得出来。"

log ""
log "--- 边界：N 不是 BLOCK 的整数倍（最后一轮只有部分线程有活）---"
# 1000 不是 256 的整数倍：1000 = 3×256 + 232，最后一轮有 24 个线程闲着
# 只测 n3/s3 —— 寄存器驻留那两级才有"幺元 + 越界判断"的组合风险
for k in n3 s3; do
  log ""; log "--- $k · N=1000 ---"
  run ./fused_ops "$k" 64 1000 3
done
log ">>> 这两个必须也通过。写越界判断时最容易在这里翻车。"

# ---------------------------------------------------------------- 2 主阶梯
log ""
log "============================================================="
log " 主阶梯：${ROUNDS} 轮"
log " RMSNorm  M=8192 行 × N=$HID    理论 n1→n3 = 12N/8N = 1.50×"
log " softmax  M=8192 行 × N=$SEQ    理论 s1→s3 = 16N/8N = 2.00×"
log "============================================================="
# seq 1 3 打印 1 2 3，$( ) 接住当成 for 的列表 —— 跑 ROUNDS 轮取中位数，防偶发抖动
for r in $(seq 1 "$ROUNDS"); do
  log ""
  log "##################### 第 $r 轮 #####################"
  # M=8192 是主基准形状；iters 给到 11，程序取中位数
  for k in n1 n2 n3 n4; do
    log ""; log "--- round $r · rmsnorm · $k ---"
    run ./fused_ops "$k" 8192 "$HID" 11
  done
  for k in s1 s2 s3 s4; do
    log ""; log "--- round $r · softmax · $k ---"
    run ./fused_ops "$k" 8192 "$SEQ" 11
  done
done

# ---------------------------------------------------------------- 3 形状扫描
log ""
log "============================================================="
log " 形状扫描 —— 行数变了，结论会不会变"
log "   decode b1   M=1      一个 block，128 个 SM 里 127 个在看戏"
log "   decode b32  M=32     还是填不满"
log "   prefill     M=2048   一次前向 2048 个 token"
log "   大批        M=16384  把卡喂饱，MBU 的上限在这里"
log "============================================================="
# 这里要同时带「名字」和「M 值」两个信息。shell 没有字典，常见做法是把一对
# 塞进一个字符串，循环里再拆开：
#   set -- $shp   把 $shp 按空格拆成词，重设位置参数 → $1 是名字，$2 是 M
#   注意 $shp 这里【故意不加引号】—— 加了就成一个词，拆不开了。
#   这是本脚本唯一一处不加引号的变量展开。
for shp in "decode-b1 1" "decode-b32 32" "prefill 2048" "大批 16384"; do
  set -- $shp
  nm=$1; m=$2
  log ""
  log "========== $nm : M=$m =========="
  # 每种形状只跑 n1/n3/n4 和 s1/s3/s4 —— n2/s2 是中间台阶，形状扫描里不需要
  for k in n1 n3 n4; do
    log ""; log "--- $nm · $k ---"
    run ./fused_ops "$k" "$m" "$HID" 11
  done
  for k in s1 s3 s4; do
    log ""; log "--- $nm · $k ---"
    run ./fused_ops "$k" "$m" "$SEQ" 11
  done
done

# ---------------------------------------------------------------- 4 长行
log ""
log "============================================================="
log " 长行：寄存器放不下的时候"
log "   Qwen2.5-7B 的 vocab 是 152064 —— 采样时要对这么长的一行做 softmax。"
log "   而 s3/s4 的单行上限是 BLOCK × MAXS = 256 × 32 = 8192。"
log "   下面三条命令里，前两条能跑，第三条会被程序拒绝并说明原因。"
log ""
log "   这不是实现偷懒，而是【FlashAttention 为什么必须分块】的同一个约束："
log "   行放不进片上存储 → 必须分块 → 分块就要能增量合并 softmax"
log "   → 于是 online 递推从'优化技巧'变成'必需品'。"
log "============================================================="
# 这三条不进循环，因为每条要配一句不同的说明 —— 日志是给人读的
log ""; log "--- logits softmax · s1（三趟，不吃寄存器）---"
run ./fused_ops s1 32 152064 11
log ""; log "--- logits softmax · s2（online 两趟，也不吃寄存器）---"
run ./fused_ops s2 32 152064 11
log ""; log "--- logits softmax · s3（寄存器驻留，预期被拒绝）---"
run ./fused_ops s3 32 152064 11
log ""
# 注意：这一条【预期失败】。脚本没开 set -e，加上 run() 用了管道（退出码取自末尾的 tee），
# 所以这次非 0 退出不会打断后面的步骤 —— 失败与否要靠读日志里的 ✗ 判断
log ">>> s3 应该打印 '✗ 这一级把整行存在寄存器里' 并退出（返回码 3），这是预期行为。"
log ">>> 同时注意 s1 和 s2 在这个形状上的 MBU —— 行长 152064、只有 32 行，"
log ">>> 两个数字放在一起看，能说明'长行 + 少行'是最难喂饱卡的一种形状。"

# ---------------------------------------------------------------- 5 竞态
# command -v 查这台机器上装没装 compute-sanitizer；没装就整段跳过，不报错
# >/dev/null 2>&1 把它的正常输出和错误输出都丢掉 —— 只要成败，不要内容
if command -v compute-sanitizer >/dev/null 2>&1; then
  log ""
  log "============================================================="
  log " 竞态检查 —— 放在计时之后"
  log " 这两个算子的 shared 用法比第 10 章复杂：blockReduceMaxSum 里"
  log " 同一块 shared 先写 warp 部分和、再被 0 号 warp 读、再广播，"
  log " 三个阶段之间少一个 __syncthreads 就是竞态。"
  log "============================================================="
  # 规模必须小（M=8 N=1000 跑 1 次）：racecheck 会逐条指令检查 shared 访问，
  # 慢几百倍，拿主基准的形状去跑等于死机
  for k in n3 s3 s4; do
    log ""; log "--- racecheck · $k ---"
    run compute-sanitizer --tool racecheck ./fused_ops "$k" 8 1000 1
  done
  log ""
  log ">>> 每段结尾都要看到 RACECHECK SUMMARY: 0 hazards displayed"
fi

# ---------------------------------------------------------------- 收尾
log ""
log "============================================================="
log " 完成 → $OUT"
log ""
log " 本地取回后："
log "   python parse_fused.py $OUT      # 解析成 CSV"
log "============================================================="

一、主阶梯↑

二、形状扫描↑

三、十四条预测的判决↑

§7 三张表的「实测」「判决」两列已经填好。 成立 9 条、落空 4 条、口径失效 1 条。 没成立的 5 条都栽在同一件事上：L2。

五条的共同原因：重用距离

字节账算的是 kernel 发出了多少访存请求；耗时取决于显存实际读写了多少。 两者之间隔着 72 MiB 的 L2。 同一行数据被读两次，两次读之间隔了多少数据，叫重用距离。 重用距离小于 L2 容量，第二次读就不花钱。

① N2 为什么快 1.49×

把 N3 实测的 917.5 GB/s 当作这张卡显存带宽的实际上限 （峰值 1008 GB/s 的 91.0%，八级里六级都贴着它跑）。 用「耗时 × 917.5 GB/s」反推每一级实际搬了多少，再和账面比：

N2 账面 12N，实际只搬了 8.1N：第二次读 xr[i] 没有走到显存。 原因见代码：

// rms_n2：两次读之间，这个 block 只处理了自己这一行
for (int i = threadIdx.x; i < N; i += BLOCK) {
    float v = xr[i];                       // 第一次读，3584 个 float = 14 KB
    acc += v * v;
}
const float r = rsqrtf(blockReduceSum(acc, smem) / N + eps);
for (int i = threadIdx.x; i < N; i += BLOCK)
    yr[i] = xr[i] * r * w[i];              // 第二次读，几微秒后，同样 14 KB

// rms_n1：pass1 读完【全部 8192 行 = 112 MiB】，kernel 结束，
//         pass2 才从头再读一遍 —— 那时 72 MiB 的 L2 早被后面的行挤干净了

预测 ① 错在默认「访存指令发出的字节」等于「显存搬运的字节」。 实测 N2（0.2591 ms）和 N3（0.2560 ms）只差 1.2% ，因为对显存来说两者都是 8N。 N2 只是多发了一遍 L2 读请求。

④⑤ S2 / S3 为什么没达到理论值

问题不在 S2 和 S3，而是基线 S1 被 L2 抬快了。 softmax 在 M=8192 × N=2048 时， x 整张是 64 MiB，L2 是 72 MiB。 刚好装得下。 S1 的 max kernel 扫完 x 后， sum kernel 再读同一份 x，其中相当一部分还在 L2 里。 所以 S1 账面 16N，实际只付了 13.7N。基线虚快，比值就缩水。

⑪⑭ MBU 什么时候失效

MBU = 账面字节 ÷ 耗时 ÷ 1008 GB/s。只要数据整个装进 L2，分子就不再是显存流量，这个数也就失去意义。

⑪ 的预测本意（2048 个 block 够把卡填满）是对的 ：grid 填充度从 2.67（N1）到 4.00（N3），确实填满了。 错在用 MBU 去验证它： 数据进了 L2，这把尺子就不再量显存。

⑫ 唯一一条在五种形状下都成立的

一个例外：M=1 和 M=32 时，N4 比 N3 快 1.49× （0.0041 ms 对 0.0061 ms）。原因不是省字节（字节一个没省）， 而是 float4 把访存指令条数减到四分之一，缩短了延迟受限时的关键路径。 喂饱卡时它没有用（1.00×）；喂不饱时，它是唯一还起作用的一级。

一、三条线索↑

二、这份报告没做、但值得做的↑

三、读完能不能回答这 6 个问题↑

• ① 为什么这里用 __shfl_xor_sync，而第 10 章 R6 用 __shfl_down_sync？ 换过来会怎样？
• ② N1 和 N2 发出的访存指令都是 12N，为什么 N2 实测快 1.49×？「指令发出的字节」和「显存搬运的字节」差在哪一层、差多少？
• ③ 为什么 float xs[MAXS] 必须配 #pragma unroll 和常量下标？ 写成 xs[cnt++] 会发生什么？从哪里能看出来？
• ④ 不加证明就把 online softmax 的递推用在蝶形归约上，错在哪？ 那半页结合律证明证的是什么？
• ⑤ 幺元为什么不能写 -INFINITY？什么输入会触发这个 bug？
• ⑥ S3 已经做到「每个元素只过一次显存」， 为什么 FlashAttention 还能比它快一个数量级？