FlashQLA 反向传播全解：fused_gdr_bwd 与 dh 梯度 Kernel 的设计取舍

FlashQLA 是基于 TileLang 构建的高性能线性注意力 Kernel 库，为 Qwen 系列 GDN（Gated DeltaNet）Chunked Prefill 提供前向/反向加速，反向路径较 FLA Triton 内核快约 2 倍。本文面向想读懂 FlashQLA 反向传播机制的新手，完整拆解 GDN 反向传播的两块核心拼图：主反向 Kernelfused_gdr_bwd（一次性算出 dq/dk/dv/dg/db/dh0 六个梯度）与 CP 场景专用的 dh 梯度 Kernel（状态梯度递推），并解释两者在流水线组织、内存布局和调度粒度上的设计取舍，帮助你快速定位反向 Kernel 源码。

一、先看全景：FlashQLA 反向传播链路

反向入口是chunk_gated_delta_rule_bwd，它与前向一样按 GPU 架构分发（SM90 → hopper、SM100/103 → blackwell、SM120/121 → blackwell_sm120），分发逻辑见 chunk/init.py。一次反向调用实际是"预处理 + 两个主 Kernel"的组合：

调用顺序可参考 chunk_gated_delta_rule_bwd；PyTorch autograd 封装在 ChunkGatedDeltaRuleFunction。

三个值得注意的入口约定：

• 输入 g 是 chunk 内 cumsum 之后的版本（前向阶段算好并保存），反向结束后 Python 端再做一次 reverse cumsum 还原成原始 gate 梯度，见 chunk_gated_delta_rule_bwd 尾部。
• dh0 始终用 fp32 累加、以 fp32 返回（状态梯度的误差对训练影响敏感）；dq/dk/dv 则回到 bf16 存储。
• dht 可选：不提供时视为零，Kernel 内部通过use_dht标志跳过加载，见 fused_gdr_bwd。

二、fused_gdr_bwd：一个 CTA 里 512 线程的"四组流水线"

主反向 Kernel 的 TileLang 源码在 fused_bwd.py（Hopper 版）。它的组织方式是每个 batch×head 分配一个 CTA（512 线程），整个序列从最后一个 chunk 往前逐块递推，512 个线程被手动划分成 4 个 warp group，各自专职（寄存器配额见 fused_bwd.py）：

几个关键设计点：

1. 状态递推与其余梯度解耦。GDN 反向的核心难点是 dS 只能从后往前逐 chunk 递推（数学参考实现见 torch_chunk_gdr_bwd）。Consumer_S 独占dh_fragment累加器，其他两组只通过共享内存和它交换中间量，避免了寄存器在 group 间反复搬运。
2. 16 个 barrier 排出的细粒度流水。Kernel 内声明了bar_00…bar_15共 16 个命名 barrier（fused_bwd.py），对应一个 chunk 内约 13 步计算，让"数据搬运（TMA）、Tensor Core GEMM、CUDA Core 逐元素运算"三类工作在时间上重叠。
3. 单缓冲 + 提前一块预取。共享内存只保留一份 chunk 数据，但 Producer 按"领先一个 chunk"的节奏提前发起 TMA 预取（fused_bwd.py）——因为主反向是计算密集型的，多留一份缓冲区换不来多少重叠收益。
4. dg/db 的原子累加放在 shared 里完成。三个 Consumer 都把各自负责的 dg 分量写入dg_shared，由 Producer 的 store 组统一写回，避免多个 warp group 对全局内存做原子操作。

三、dh 梯度 Kernel：为什么要把状态梯度单独拆一个 Kernel？

fused_gdr_dh的实现是各架构下的fused_gdr_dh_ws（如 Hopper 版 cp_bwd.py），它只在CP 开启时由 cp_preprocess_bwd 触发，做一件事：从 dht 出发，只递推状态梯度，输出每个 CP 分区头部应有的 dht（output_dh0=True, output_dh=False）。

拆开这个 Kernel 的取舍是本文的重点：

• 让 CP 逻辑与主 Kernel 隔离。CP 下每个分区只需处理num_warmup_chunks个"预热 chunk"（gate 指数衰减后远端影响可忽略），这个裁剪逻辑全部收敛在 dh Kernel 内，主fused_gdr_bwd不需要感知 CP。
• 为跨卡聚合留出独立接口。dh Kernel 输出的分区级 dh0 可以直接被aggregate_card_state/correct_terminal_states做 all-gather 修正（见 _inter_correct_bwd），修好后再作为dht喂给主 Kernel。
• warmup 提前量可被 gate 感知。Auto CP 用 gate 的衰减特性估计预热深度，决策模型的说明见 autocp/README.md，拟合系数在 coefs/。

在 CP 场景下它同样采用 Producer/Consumer 结构，但流水线风格与主 Kernel 完全不同：

dh Kernel 的三类 Consumer 分别负责：dh 组持有dh_fragment累加器执行dh ← g_last·dh + RgᵀdO + XᵀY递推；PR 组计算修正后的 Q（R = Q − Lower(P)·X再乘 gate）；XY 组计算X = Ab·K、Y = K·dh（分工注释见 cp_bwd.py）。

四、两个 Kernel 的取舍总结与实测收益

一句话总结：主反向 Kernel 用"单缓冲 + 细粒度 barrier"榨干计算重叠，dh Kernel 用"双缓冲 + 粗粒度 stage"换取 CP 预热循环的调度弹性——两者的 warp group 专门化、寄存器配额（set_max_nreg）手法一致，差异只在流水深度。

性能上，反向路径的实测数据可参考 benchmark_results_H200.txt：H200 上hk16_hv16配置 32k 序列单次更新步，flash_qla 反向 1.135ms vs FLA 3.418ms，约 3.01×；hk32_hv32约 2.56×。跑分脚本为 bench_gated_delta_rule.py，CP 场景另见 bench_gated_delta_rule_cp.py。

正确性验证方面，单测用 float64 参考实现逐元素对比 bf16 输出（相对误差容忍 2%），覆盖 dense/varlen/GQA/长序列等场景，见 tests/test_gdr_unit.py 与参考实现 tests/ref_gdr.py；CP 端到端验证在 tests/test_cp_e2e.py。

想深入阅读建议的路径：先看 chunk/init.py 的调用编排 → 再读 fused_bwd.py 四个 warp group 的循环体 → 最后对照 cp_bwd.py 理解 CP 下 dh 递推如何被 warmup 裁剪。