k2_bwd_recurrence_deep_dive

K2_bwd 反向递推 kernel 逐行精讲（小白可推导版）

对应源码：csrc/smxx/bwd_kernel2.cuh 的 _flash_kda_bwd_recurrence

目标读者：没接触过自动微分 / CUDA kernel 的人。读完应能自己把每一条梯度公式推一遍，并理解代码为什么这么写。

约定：d<量> 表示 loss 对该量的梯度。代码里的 exp(x) 实际是 （exp2），求导带的 因子被统一推迟到 K1_bwd 处理，本文遇到时会注明。公式用 LaTeX，伪代码用代码块。

---

0. 先建立全局图景

0.1 FlashKDA 把序列切块做线性注意力

序列按 CHUNK 长度切块。块内用注意力，块间用一个状态矩阵 做递推。前向每块做：

反向就是对这一串求伴随（adjoint，即反向传播）。 表示按行/按元素广播乘。

0.2 形状与记号

符号	形状	含义
CHUNK ()	标量	块长，例 16
	标量	head 维度，例 128
		进入本块时的状态。存储是转置的，见 0.3
		前向 K1 生成的派生量（fp32）
		整块门控总和，workspace 里存
		每行一个门控标量（已过 sigmoid）
		，下三角
		输入 ； 是输出 的上游梯度
		中间量（前向没存，反向重算）
		跨块传播的状态梯度累加器（常驻 smem）

0.3 状态的转置存储（极其重要，否则下标全错）

all_states / ds_init / ds_out 在显存里存的是 ，布局 。即：

源码中凡出现 s_in_ptr[d*D + k]，语义都是数学上的 。

0.4 为什么必须从后往前

第 块入口的状态梯度 依赖第 块算出的 。所以反向必须 从最后一块往第一块 迭代（for t = t_tiles-1 ... 0）。

---

1. Kernel 启动配置与线程模型（line 29-70）

template <int CHUNK, int D, int NumThreads, bool IsVarlen = true>
__global__ void __launch_bounds__(NumThreads) _flash_kda_bwd_recurrence(...)

• Grid = (N, H)：blockIdx.x = seq_idx，blockIdx.y = head_idx。每个 block 独占一条序列的一个 head。
• Block = 256 线程，无 warp specialization（注释 line 8）。
• 纯标量的 shared-memory GEMM 循环，不用 tensor core；逻辑直白但有冗余重算（用算力换显存）。

---

2. 序列定位（line 72-89）

把 seq_idx 映射到 token 维范围 和全局 tile 偏移 tile_base。

• Varlen：用 cu_seqlens 取 bos/eos；tile_base 累加前面序列的块数。
• 定长：直接乘除。

---

3. 共享内存布局与 dS 初始化（line 91-117）

float* dS_smem     = (float*)shared_mem;   // [D*D], 整个 kernel 常驻
float* scratch_f32 = dS_smem + D*D;        // 后面所有临时 tile 堆这里

初值来自 dfinal_state（对最终状态的梯度），没有则为 0。转置加载：

---

4. 主循环：加载本块数据（line 120-177）

for (int t = t_tiles - 1; t >= 0; --t) {
    int ws_idx     = head_idx * total_tiles + tile_base + t;
    int actual_len = min(CHUNK, seq_len - t*CHUNK);   // 尾块可能不满

• actual_len：尾块真实行数；越界行读 时取 0。
• 把 （已 exp2）、（读原始再过 sigmoid_tanh_approx）、、 搬进 smem。
• scratch 区堆叠顺序（贯穿全 kernel）：

gT[D] | beta[C] | INV[C*C] | Mqk[C*C] | vcorr[C*D] | U[C*D] | dU[C*D] | dkr[C*D] | dvcorr[C*D] | dL[C*C]

多数 buffer 复用别名，唯独 dL 独立（见第 8 节）。

---

5. Step 1：重算 U（line 179-235）

前向的 未存，反向重算。

5.1 的 是 次乘加，最重。vcorr_smem 复用 kd_S_smem。

---

6. Step 2：输出梯度反传（line 237-303）

前向 ，分两路。

6.1 dU 来自 路（line 251-267）

由 ，对 求偏导得 ：

6.2 dU 加状态路贡献（line 269-287）

由 ， 也影响 loss：

其中 是后一块传回的状态梯度。

6.3 dkr（line 289-303）

对同一式中的 求偏导得 ：

---

7. Step 3：三角求解的伴随（line 305-321）—— 重点之一

前向 ，，等价解 。

本节是全文最难的一处。下面 7.0 是写给线性代数基础薄弱读者的从零详解，已经熟悉矩阵微分的读者可直接跳到 7.1 看结论。

7.0 从零搭起：只用一条核心规则

(a) 三个积木

积木 1 — 矩阵乘法 ： 的第 行第 列 = “A 的第 行"点乘"X 的第 列”。

一句话：matmul 就是一堆点积。

积木 2 — 转置 ：行列互换，（沿对角线翻一下）。

积木 3 — 逆矩阵 ：矩阵版的"倒数"，定义 （ 是单位矩阵，相当于数字 1）。你不需要会手算它，只要知道：

即"解线性方程组"和"求逆再乘"是同一件事。本例 就等价于解 。

(b) 反向传播唯一要背的核心规则

其中 是"loss 对 的梯度"（与 同形状）。推导（只用积木 1 的点积定义 + 链式法则）：

由 ，

转置就是这么冒出来的。 口诀：求输入 的梯度，就把另一个因子 转置乘到 左边；求权重 的梯度，就把 转置乘到 右边。

© 标量类比（建立"为什么有逆/转置"的直觉）

把矩阵换成普通数字，看除法 ：

记住这两条标量结果，矩阵版长得一模一样（只是 变 、注意左右与转置）。

(d) 第一步：求 —— 直接套核心规则

前向 就是 ，对应 。求"对输入 的梯度"：

(e) 第二步：求 —— 需要"逆矩阵的梯度"这块拼图

藏在 里，要多走一步。

先用核心规则求"对权重 的梯度"：

再用逆矩阵反向规则（标量 的矩阵版，推导见 (g)）：

套 ，且 ：

(f) 化简到代码里的简洁式

把 代入并重新分组（结合律）：

（用了 ，以及 。）于是得到

又和标量版结构完全对应。

(g)（选读）逆矩阵反向规则的来历

由 求微分：；再转成伴随（内积配对，见 7.1）得 。

小结：第一步 ，第二步 。全程只用了"matmul 反向"这一条核心规则 + 一次"逆矩阵反向"。和标量 的 一一对应。

7.1 通用公式（务必记住）

对 ，已知 （即 ）：

推导（内积配对）。对 取全微分，用 ：

loss 一阶变化 ，其中 ：

7.2 套进本例（）

---

8. Step 3b：dL（line 324-338）—— 重点之二

由 7.1 的 ，且 ：

（这条公式的完整从零推导见 7.0(e)(f)；直觉：它是"输出敏感度 "与"解 "的外积、带负号，正对应标量 。）

严格下三角（上三角与对角前向被 mask 成 0，那些位置根本不是自由参数），故梯度也 mask：

工程关键：dL 依赖 dvcorr，而 dvcorr_smem 在 Step 5 会被覆盖。所以必须趁 dvcorr、U 还在时立刻把 dL 算进独立不别名的 dL_smem（line 309 注释 “not aliased”）。这就是它被命名为 “Step 3b”、紧贴 Step 3a 之后的原因。

---

9. Step 4：vcorr 反传到 dv、dbeta（line 341-382）

前向 。

代码仅在 时做除法（line 375）。

---

10. Step 5-6：dMqk → dqd、dki（line 384-451）

前向 ，。

dqd 两路（line 411-437）：

dki 来自 Mqk 路（line 439-451）：

---

11. Step 7：L 路反传 → dki、dkd、dbeta（line 453-535）

前向 。记 ，则 。

dki 加 L 贡献（line 472-485）：

dkd 两路（line 487-518）：

此时 dvcorr_smem 已被覆盖，但 inv_smem、dU_smem 还在，于是就地重算 （line 500-505）。

dbeta 加 L 贡献（line 520-534）：

---

12. Step 8：dgT 与 dS 递推更新（line 543-625）

12.1 dgT（必须在更新 dS 之前算，line 567-592）

有两个去处：状态衰减 ，以及 ：

代码里 gT_smem[k] 已是 。用的是当前(后块传来) dS，故必须先于 12.2。

下面 12.1.0 是从零详解（ 是个长度 的向量， 扇出到很多地方，逐路收集）；已熟悉的读者可跳过。

12.1.0 从零推导 dgT

(a) 先看 在前向被用在哪。 是长度 的向量（每个 head 维一个数）。 出现在两个地方，所以它的总梯度 = 两路回传之和（扇出 → 求和，多元链式法则）：

• 路 1 — 状态衰减：（注意只有第 行乘 ，因为衰减是按 这一维广播的）。
• 路 2 — 进入 ：（每个 都用到同一个 ）。

(b) 路 1 的偏导。 固定 ， 影响第 行所有列 的 。用指数求导 （这里同样把公因子 按"dgT 约定"省略，留到 K1_bwd 统一处理）：

乘上游 并对所有列 求和（因为 影响了整行）：

代码里 gT_smem[k] 已经是 ，所以源码先 sum_d dS*S_in 再 *= gT_k（line 581-583）。

© 路 2 的偏导。 出现在每个 的 里，同样 （因 ）。乘上游 并对所有行 求和：

这里的 是本块第 6.3 节刚算出的（dkr_smem）， 从 workspace 读。

(d) 两路相加。

(e) 为什么必须"先于 12.2"。 路 1 用的 是进入本块时（后一块传回来）的 ；而 12.2 会把 dS_smem 原地覆盖成"前一块的 "。若先跑 12.2，路 1 就会读到错的 。所以源码严格按 Step 8a (dgT) → Step 8b (更新 dS) 的顺序（line 567 注释 “compute dgT BEFORE updating dS”）。

符号小提醒： 和 都来自门控， 的指数是 。所以 增大让 变大（ 号，本节路 2）；而 增大让 变小（ 号，那部分梯度走 ，在 K1_bwd 处理，见 3.2(b)）。同一个 对 和 贡献符号相反，别搞混。

12.2 dS 反向递推（line 594-625）

把 更新为前一块的状态梯度，三项相加：

三项偏导依据： ； ； 。

dvcorr 就地重算。写回 dS_smem，进入更前一块。

---

13. 收尾：写出 d_initial_state（line 628-641）

循环结束时 dS_smem 即第一块入口的状态梯度。转置回 写出：

---

14. 全部梯度公式速查表

---

15. 设计要点总结

主题	做法
从后往前	依赖后块，必须逆序迭代
重算换显存	不存，靠 在 smem 重算
状态转置存储	all_states/ds_init/ds_out 都是 ；取 写 s_in_ptr[d*D+k]
buffer 复用与别名	scratch 严格堆叠；dMqk 复用 Mqk，dqd/dki 复用 dvcorr；唯 dL 独立
gT 已是 exp2	workspace 存
尾块掩码	actual_len 控制越界行读 0
数值	workspace fp32；dbeta 除 时加

配套阅读：K1_bwd 的逐行精讲见 k1_bwd_prepare_deep_dive.md。本 kernel 产出的 正是 K1_bwd 的输入。