第六章 Dao 反向 CUDA 实现

反向一次完成哪些事

反向需要获取 $x,w,\beta,g$，可选地重建 $z$ 并计算 $\delta$，随后生成 $dx$，归约 $dw,d\beta$。带状态时还要接收 $dF$、输出 $dS$。

融合的价值是重复使用已经搬入的输入和梯度。如果拆成多个 kernel 并将 $z$ 或 $\delta$ 全部写到全局内存，会增加往返；如果融合过多，寄存器和片上内存也可能成为限制。Dao 使用融合路径，但 MLU 上仍需要结合资源模型选择。

以下结构对应 v1.5.0 反向 CUDA 源码。

先推导反向 tile 的数据范围

令本任务负责 $dx$ 的时间区间 $[a,b]$。根据第三章：

$$ dx_s=\sum_{j=0}^{K-1}w_{K-1-j}\delta_{s+j}, $$

它需要 $\delta$ 的区间 $[a,b+K-1]$。这是右 halo。若融合 SiLU 反向，这些 $\delta_t$ 又需要相应预激活 $z_t$，而重建 $z_t$ 需要输入 $[t-K+1,t]$。合起来，输入范围是：

$$ x:\ [a-K+1,b+K-1],\qquad g:\ [a,b+K-1]. $$

因此，带 SiLU 重计算的独立反向 tile 可能需要输入左右两侧 halo。只照着前向多搬左侧数据，无法覆盖 $dx$ 右端所需的激活导数。

这里讲的是独立 tile 的依赖集合。实现可以用跨 chunk 保留中间值等方式减少重复加载，但不能省掉数学依赖。

时间连续反向为什么倒着扫

普通路径仍以 (B,C) 划分 block，但从最后一个时间 chunk 向前扫描。先处理较晚位置后，可以把其梯度信息保留给前一个 chunk，满足 $dx$ 对未来梯度的需求。

每个线程将负责位置的 $dw_k$ 贡献累积到 FP32 寄存器数组，并累计偏置梯度。完成所有 chunks 后，通过 CUB BlockReduce 汇总同一 block 的贡献。随后由指定线程对全局 FP32 参数梯度做 atomicAdd，合并不同 batch 的结果。

于是两种方向同时出现：$dx$ 需要从右侧收集梯度；$dw$ 则在同一循环中把所有时间位置的乘积逐步求和。不能把“倒着扫”理解成将权重梯度公式的输入序列简单反转。

通道连续反向的独立 tiles

channel-last 反向继续沿 batch、时间 tile 和通道 tile 划分 grid。该版本在 $L\le128$ 时用长度 64 的时间 tile，较长时用 128。这是源码中的配置选择，不是通用硬件定律。

共享内存中 dout_smem 带右侧 $K-1$ 行，x_smem 为左右 halo 预留空间。重建预激活时应用与前向一致的 bias 和序列边界规则，然后使用：

$$ \delta=g\sigma(z)[1+z(1-\sigma(z))]. $$

每个线程拥有自己的一段时间位置。参数梯度先在这段内累加，再用 Allreduce<kNThreadsPerRow> 汇合同一通道对应线程的结果，最后一个代表线程原子加到全局。$dx$ 则使用右侧 $\delta$ 与反序索引的权重进行 gather，并写回自己的唯一输出区域。

halo 参与 dx 不应被重复计入 dw

假设 tile A 负责输出梯度位置 $[0,63]$，为了计算它的 $dx$，额外读了位置 $[64,66]$ 的梯度。tile B 将负责这些位置的参数梯度贡献。

如果 A 也把 halo 的贡献加进 $dw$，同一输出位置就会被计算两次。正确区分是：依赖读取范围可以重叠，参数梯度贡献的所有权范围不能重叠。 这是移植独立时间 tiles 时非常重要的检查项。

为什么原子归约可能有微小数值变化

浮点加法不满足严格结合律：$(a+b)+c$ 可能不等于 $a+(b+c)$。多个 block 到达全局原子累加的先后顺序可能变化，因此 FP32 累加仍可能出现末位差异。

如果业务需要稳定的归约顺序，可以让各任务把部分和写到独立 workspace，再用固定顺序的第二阶段归约。代价包括 workspace、额外启动和读写。是否采用应由正确性契约和性能测量决定；不能只为了逐 bit 对齐一个并行实现而牺牲所有性能。

如果低精度参数梯度最终溢出，FP32 中间累加也无法保证最终低精度结果有限。混合精度训练还需要从 loss scaling、参数 dtype 和优化器状态整体分析。

状态梯度是另一条连接

最终状态 $F$ 是输入尾部的拷贝，因此它的梯度应直接加回对应输入，不乘卷积权重或 SiLU 导数。初始状态参与了开头若干窗口，所以其梯度还包含普通卷积链式法则的贡献。

Dao 的 channel-last 反向在写回 $dx$ 前添加 dfinal_states 的直接路径，并在首个时间 tile 计算 dinitial_states。当 $L<K-1$ 时，最终状态中仍可能含有初始状态的元素，这些元素还必须接收 $dF$ 的直接贡献。下一章给出统一公式。

从症状判断错误位置

症状	优先检查
每个 tile 末尾 $K-1$ 个 dx 错误	右 halo、越界梯度清零
只有 SiLU 模式错误	预激活重建、导数、输入右 halo
dx 正确但 dw 随 tile 数偏大	halo 被重复计入参数梯度
batch=1 正确，batch 增大后错误	全局归约、初始化、竞争写
只有分块训练结果不一致	dF 到 dx/dS 的直接路径、状态是否 detach
packed 边界附近错误	前反向 mask 是否使用一致的源目标位置

下一章：状态分块与增量推理。