K1_bwd “prepare” kernel 逐行精讲（小白可推导版）

对应源码：csrc/smxx/bwd_kernel1.cuh 的 _flash_kda_bwd_prepare

目标读者：没接触过自动微分 / CUDA kernel 的人。读完应能自己推出 L2 归一化反向、数值稳定的、以及门控的逆向 cumsum，并理解代码为何这样写。

约定：d<量> 表示 loss 对该量的梯度。exp(x) 实际是（exp2）。公式用 LaTeX，伪代码用代码块。

0. 它在整条反向链里的位置

1	Backward: K2_bwd (反向递推, 后->前) -> K1_bwd (本文, 还原输入梯度)

K2_bwd 算出了一堆"派生量的梯度" 。 K1_bwd 是前向 K1 的逆过程：把这些梯度还原回原始输入的梯度—— （门）、，以及参数、。

0.1 前向 K1 做了什么（理解反向的前提）

设每行（一个 token）的原始，门控相关原始量、参数、。前向：

归 一 化 块 内 前 缀 和 整 块 总 和

所以反向要做三件事：把收拢成和；把过 L2 反向得；把（+）过 cumsum 和门函数得。

1. 线程模型（line 10-62）

1 2	constexpr int ELEMS_PER_THREAD = 8; constexpr int THREADS_PER_ROW = D / 8; // = 16

Grid = (total_tiles, H)：blockIdx.x = global_tile_idx，blockIdx.y = head_idx。每个 CTA 处理一个 chunk 的一个 head（line 7）。
每行 16 个线程，每线程负责 8 个元素：

这个布局和前向 K1 完全一致（line 60），保证归约顺序相同 → 数值可复现。

2. tile → 序列定位（line 64-99）

把 global_tile_idx 映射到 seq_idx / local_t，并算出：

越界 tile（local_t >= t_tiles_this_seq）直接 return（line 95）。 a_log_exp = exp(A_log[head])（line 101）。

shared memory 只有一块：dgc_smem[C*D]（line 105）。

3. Phase 1+2：融合算 dgc 与 dq/dk（line 129-218）

这是全 kernel 最核心、最讲数值稳定的部分。

3.1 重算 L2 范数（line 136-157）

每线程读自己负责的 8 个元素，累加平方和；尾块越界行读 0；再用 16 线程的 __shfl_xor 归约（）把整行平方和汇总：

归约方式必须与前向逐位一致。

3.2 重算门控指数、算 dqn/dkn 与 dgc（line 163-193）

对每个元素（列）：

（代码里 gt_tile[col] 存的就是，所以 exp_gt_gc = gt_tile[col]*exp_neg_gc。）

(a) dqn / dkn（供 L2 反向）

前向，所以对：

同时进了三个分支，梯度相加：

(b) dgc —— 数值稳定写法（关键 trick）

下面先给从零详解（b.0–b.4），已熟悉的读者可直接看 b.3 的结论框。

(b.0) 先搞清是什么、要求什么

是"门控累加和"，一个标量（对每行每维）。前向它出现在四个派生量的指数里：

上游已经给了这四个量的梯度。一个变量被用到多处（扇出 / fan-out），它的总梯度 = 各处回传之和（多元链式法则）。所以：

(b.1) 求每个偏导（指数函数求导）

对指数函数。这里底数，所以。本 kernel 约定把公共因子推迟到 Phase 4 统一乘（因为对所有项线性，提一个常数出来最后乘即可），故这里先省略。导数符号跟着指数的号走：

（指数是，链式带出一个负号。）代回得天真公式（数学正确）：

天 真

(b.2) 为什么天真公式在 fp32 里会崩（灾难性抵消）

，。当较大时这两者量级天差地别。举个具体例子，设、：

两者相差约。fp32 只有约 7 位十进制有效数字（相对精度）。一旦把（可能上千万量级）和（极小）放进同一个加减式：

大 + 小：小项落在大项的有效位之外，被直接舍掉 → 信息丢失；
大 − 大：若两个大项接近，相减后高位全抵消、只剩低位舍入噪声 → 灾难性抵消（catastrophic cancellation）。

结果基本是垃圾。注释 line 134 说的 “kd, ki can differ by ~1e20” 就是这个意思。

(b.3) 稳定写法：把提到括号外

注意，对同理。把公共的（和）提出来：

其中。对应源码 line 186-187。

(b.4) 为什么这样就稳了

关键观察：括号里的乘积 本身就是 （正常大小）。原因是上游梯度是按的比例传下来的（前向，反向链式正好带一个），二者相乘把抵消掉。于是：

括号内是"几个正常大小的数相减" → 良态，没有大数吞小数；
再乘一个的（单位向量的分量）→ 仍然正常。

全程不出现这种怪物。对比一句话：

天 真 ： 稳 定 ：

两者数学恒等，只是运算顺序不同，浮点结果天差地别。这是典型的"重新结合以避免中间量爆炸"的数值技巧。

注意区分两个量（来源相同、符号不同）：

全 对 的 梯 度 ， 给 反 向

带 符 号 ， 对 的 梯 度

进不带的符号（前向里是线性因子，故全加）；而进指数带号（故带符号）。两者只差中间两项的正负号，别混用。

3.3 L2 归一化的反向 → dq/dk（line 195-216）

前向。要求已知求。

下面 3.3.0 是写给线代基础薄弱读者的从零详解；已熟悉的读者可直接看 3.3.1 的结论。

3.3.0 从零推导

为什么不简单：若这种逐元素映射，求导就完事。但归一化里 装着所有分量 ，所以改动任一会通过影响每一个 。这就是最终会冒出耦合项的根源。

第一块：对的导数。 设，。只有第项含，故，再用的链式：

第二块：对的导数（核心）。 ，用乘积法则，（克罗内克符号，为 1 否则 0），（代入 ★）：

☆

两部分含义：是"自己对自己"的直接项（仅）；是通过产生的"所有分量互相耦合"项（任意都有）。

第三块：链式法则汇总。 影响所有，把贡献全加起来，并代入 (☆)：

◇

第一项里使求和塌缩成单项。

第四块：用化简。 关键代换：

代回 (◇)：

3.3.1 结论

标 量

对应代码：

点积用 16 线程 shfl 归约出整行（line 195-205），仅 my_row < actual_len 才写出（line 207）。这个归约就是在算第二块说的"耦合标量"——所以归一化反向必须先做一次行内归约，无法纯逐元素。

直觉：是单位向量，沿方向拉伸不改变（只改长度），所以梯度里"沿的径向分量"无效。是在上的投影长度，是该径向分量，减掉它只留切向分量；再乘（越长，同样的对影响越小）。一句话：把梯度投影到方向，再按长度缩放。

3.3.2 小数字验证（2 维）

取，则，。设：

自检：恒垂直于（沿推不改变）：。一般地

可作为写单测时的快速断言。

4. Phase 3：dv 直接搬运（line 220-224）

在 K2_bwd 已算好（bf16），这里只是从 workspace 的 tile 布局拷到最终布局，仅拷 actual_len*D 个元素。

5. Phase 4：门控反向 = dgc 的逆向 cumsum + dgT（line 226-261）

下面 5.0 是写给基础薄弱读者的从零详解（核心就一句：前缀和的反向是后缀和）；已熟悉的读者可直接看 5.1。

5.0 从零理解"cumsum 的反向"

(a) 什么是 cumsum（前缀和）

前向把一串数累加成"到当前为止的和"：

展开看就是（以为例）：

(b) 关键观察：每个输入"扇出"到哪些输出

竖着看上面这张表：出现在第行及其下面所有行里。比如出现在（不在）。用偏导写：

© 反向：把扇出的梯度收回来 = 后缀和

多元链式法则：被用到多处，它的总梯度 = 这些输出回传的梯度之和，每个偏导都是 1，所以就是直接相加：

即"从到末尾"的求和——这叫后缀和（suffix sum）。还是：

一句话规律：前向是前缀和（从前往后累加），反向就是后缀和（从后往前累加）。 方向正好反过来，所以叫"逆向 cumsum"。

(d) 怎么算出来：一个滚动变量

不用对每个都重新求和（那是）。从最后一行往前走，维护一个累加器 rev_sum，每步把当前行的加进去：

rev_sum = 0
for row = C-1 down to 0:          // 从后往前
    rev_sum += dgc[row]           // 此刻 rev_sum == sum_{r>=row} dgc[r]
    dg_total[row] = rev_sum

走到 row 时，rev_sum 恰好就是，正是我们要的后缀和。对应源码 line 237-238。

(e) 数字小例（）

设（下标 0…3）。从后往前滚：

步 row	加入	rev_sum	即
3	3	3	3
2	1	4	4
1	5	9	9
0	2	11	11

可手动验证 ✓。

(f) 再叠加那条路

其实还流进了整块总和（见 5.1）。对所有，所以每行都再加一份相同的：

来 自 （ 后 缀 和 ） 来 自 （ 常 数 ）

小结：cumsum 反向 = 后缀和（一个滚动累加器搞定）；再贡献一个对每行都相同的常数。下面 5.1 把这两件事正式写出来。

5.1 g_total 扇出到两个去处

前向同时进入两处（计算图里是并列两条边）：

多元链式法则：一个量被用到多处（扇出），其总梯度 = 各下游路径回传之和。

来 自 来 自

gc 这条：因仅当，所以是"后缀和"。代码从 row=C-1 往下走，维护 rev_sum += dgc[row]，到时 rev_sum 正是（这就是"逆向 cumsum"）。
gT 这条：因对所有，所以每行都加同一个常数 dgT_val（line 232 循环外读一次，line 239 每行加）。

（此处仍是"未乘 "约定值。）

5.2 过门函数反传到 dg_raw 与参数（line 241-260）

把（=dg_nat）沿往回传，并补上延后的。

下面 5.2.0 是从零详解（单变量链式法则 + sigmoid 求导 + 一个量被多处使用怎么办）；已熟悉的读者可直接看 (a)(b)©。

5.2.0 从零理解这条链

(i) 前向是一串"嵌套函数"。 从最里到最外：

原始输入和两个参数都在这条链上。反向就是从右端的已知梯度一站一站往左乘"每站的局部导数"。

(ii) 单变量链式法则（唯一要会的）。 若，则。也就是"上游梯度 × 本站导数"。一站一站接力即可。

(iii) sigmoid 的导数为什么是。 。求导：

（用了。）这就是代码里的 dsig = sig*(1-sig)。

(iv) 从哪冒出来。 回忆 3.2(b)：算时我们故意省略了 求导该带的公因子，约定"最后统一补"。是的来源（是的前缀和），所以传到这里的 dg_nat 也是"少乘了一个 "的版本。现在就把它补回来：真正的梯度 = dg_nat 。因为整条链对 dg_nat 是线性的，补在哪一步都行，代码补在算这一步。

(v) 一个量被用在多处怎么办（的情形）。 同时乘进了 ——它只出现一次在里，简单。但反过来看对三个东西都有依赖：、、。求对它们各自的偏导就是把另外的当常数：

(a) 过 ，：

（代码 line 247 把写在最后，等价。）

(b) 过 ：

© 过 ，，：

代码 line 252-253 先累加 ddt_bias_acc += dz*a、dA_log_acc += dz*(g_raw+dt)，循环后：

为什么用 atomicAdd（line 259-260）：是每 head 一个标量、是每 head 每维，同一 head 的很多 chunk-CTA 会并发写同一地址，必须原子加。而是每 token 独立位置，直接写。

尾块掩码：row >= actual_len 的行 dg_out=0（line 255），且不累加参数梯度。

6. Phase 5：dbeta 从 chunk 空间转回 logit 空间（line 263-272）

K2_bwd 给的是对"已过 sigmoid 的 "的梯度。前向，反传过 sigmoid 回 logit：

仅前 actual_len 个线程各写一个 token 行。

7. 全部梯度公式速查表

重 算 收 拢 反 向 门 控 直 接 从 拷 贝

8. 设计要点总结

主题	做法
数值稳定 dgc	把量级悬殊（）的改写成 $项相减$ ，避免 fp32 抵消
16 线程/行 + shfl	严格对齐前向 K1 的归约顺序，保证 bit 级可复现
dkn vs dkn_signed	前者全加（对，给 L2 反向）；后者带符号（对）
门控逆向 cumsum	是前缀和 → 反向是后缀和（rev_sum），再叠加常数
ln2 延后	求导的统一在 Phase 4 乘入
atomicAdd	跨 chunk-CTA 累加；其余每 token 独立直接写
尾块掩码	`actual_len` 控制越界行读 0 / 写 0 / 不累加参数

配套阅读：上游 K2_bwd 的逐行精讲见 k2_bwd_recurrence_deep_dive.md。本 kernel 的输入全部来自那里。

k1_bwd_prepare_deep_dive