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

【免费下载链接】FlashQLAhigh-performance linear attention kernel library built on TileLang项目地址https://gitcode.com/gh_mirrors/fl/FlashQLA点击查看免费下载FlashQLA 是基于 TileLang 构建的高性能线性注意力 Kernel 库为 Qwen 系列 GDNGated DeltaNetChunked 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的组合步骤函数职责① CP 预处理cp_preprocess_bwd仅在开启 Context Parallelism 时生效负责跨分区梯度修正内部会调用 dh 梯度 Kernel② 重算 hfused_gdr_h重新计算前向的 chunk 内状态 h反向需要它③ 主反向fused_gdr_bwd单次 Kernel 输出 dq、dk、dv、dg、db、dh0 六个梯度④ 后处理Python 端GQA 分组归约、dg 的 reverse cumsum 等轻量收尾调用顺序可参考 chunk_gated_delta_rule_bwdPyTorch 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.pyHopper 版。它的组织方式是每个 batch×head 分配一个 CTA512 线程整个序列从最后一个 chunk 往前逐块递推512 个线程被手动划分成 4 个 warp group各自专职寄存器配额见 fused_bwd.pyWarp Group线程范围寄存器配额职责Consumer_S0–127160状态梯度 S 的跨 chunk 递推dS ← g_last·dS Qᵀ(dO·g) Kᵀ(dVg)并累出 dh0Consumer_K128–255144计算 dK、dV 及相关的 dg 分量Consumer_A256–383184计算 dQ、dA、db 及剩余的 dg 分量Producer384–51124用 TMA 异步搬运数据另负责把 dq/dk/dv/dg/db 写回全局内存几个关键设计点状态递推与其余梯度解耦。GDN 反向的核心难点是 dS 只能从后往前逐 chunk 递推数学参考实现见 torch_chunk_gdr_bwd。Consumer_S 独占dh_fragment累加器其他两组只通过共享内存和它交换中间量避免了寄存器在 group 间反复搬运。16 个 barrier 排出的细粒度流水。Kernel 内声明了bar_00…bar_15共 16 个命名 barrierfused_bwd.py对应一个 chunk 内约 13 步计算让数据搬运TMA、Tensor Core GEMM、CUDA Core 逐元素运算三类工作在时间上重叠。单缓冲 提前一块预取。共享内存只保留一份 chunk 数据但 Producer 按领先一个 chunk的节奏提前发起 TMA 预取fused_bwd.py——因为主反向是计算密集型的多留一份缓冲区换不来多少重叠收益。dg/db 的原子累加放在 shared 里完成。三个 Consumer 都把各自负责的 dg 分量写入dg_shared由 Producer 的 store 组统一写回避免多个 warp group 对全局内存做原子操作。三、dh 梯度 Kernel为什么要把状态梯度单独拆一个 Kernelfused_gdr_dh的实现是各架构下的fused_gdr_dh_ws如 Hopper 版 cp_bwd.py它只在CP 开启时由 cp_preprocess_bwd 触发做一件事从 dht 出发只递推状态梯度输出每个 CP 分区头部应有的 dhtoutput_dh0True, output_dhFalse。拆开这个 Kernel 的取舍是本文的重点让 CP 逻辑与主 Kernel 隔离。CP 下每个分区只需处理num_warmup_chunks个预热 chunkgate 指数衰减后远端影响可忽略这个裁剪逻辑全部收敛在 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 完全不同维度fused_gdr_bwdfused_gdr_dhCP 版共享内存缓冲单缓冲2-stage 双缓冲num_stages2流水粒度16 个细粒度 barrier13 步/chunk6 个计算 stage 数据 ready/free 双 barrier预取节奏Producer 提前 1 个 chunkProducer 提前 1 个 stage 写入下一 stage典型负载计算密集整条序列都要算满只跑 warmup 的少量 chunk访存调度开销占比更高dh Kernel 的三类 Consumer 分别负责dh 组持有dh_fragment累加器执行dh ← g_last·dh RgᵀdO XᵀY递推PR 组计算修正后的 QR Q − Lower(P)·X再乘 gateXY 组计算X Ab·K、Y K·dh分工注释见 cp_bwd.py。四、两个 Kernel 的取舍总结与实测收益一句话总结主反向 Kernel 用单缓冲 细粒度 barrier榨干计算重叠dh Kernel 用双缓冲 粗粒度 stage换取 CP 预热循环的调度弹性——两者的 warp group 专门化、寄存器配额set_max_nreg手法一致差异只在流水深度。性能上反向路径的实测数据可参考 benchmark_results_H200.txtH200 上hk16_hv16配置 32k 序列单次更新步flash_qla 反向 1.135ms vs FLA 3.418ms约 3.01×hk32_hv32约 2.56×。跑分脚本为 bench_gated_delta_rule.pyCP 场景另见 bench_gated_delta_rule_cp.py。正确性验证方面单测用 float64 参考实现逐元素对比 bf16 输出相对误差容忍 2%覆盖 dense/varlen/GQA/长序列等场景见 tests/test_gdr_unit.py 与参考实现 tests/ref_gdr.pyCP 端到端验证在 tests/test_cp_e2e.py。想深入阅读建议的路径先看 chunk/init.py 的调用编排 → 再读 fused_bwd.py 四个 warp group 的循环体 → 最后对照 cp_bwd.py 理解 CP 下 dh 递推如何被 warmup 裁剪。赞分享【免费下载链接】FlashQLAhigh-performance linear attention kernel library built on TileLang项目地址https://gitcode.com/gh_mirrors/fl/FlashQLA点击查看免费下载相关推荐OpenProject 开源项目管理软件实战导读免费自托管甘特图与团队协作一次配齐OpenProject 开源项目管理软件实战导读免费自托管甘特图与团队协作一次配齐 当你同时盯多个版本迭代任务散落在群聊和表格里没人说得清哪个里程碑会延后端前端企业应用研发协作协同办公手动梯度计算完全指南掌握反向传播的本质与实现技巧手动梯度计算完全指南掌握反向传播的本质与实现技巧 在深度学习中反向传播算法是训练神经网络的核心技术而理解梯度计算则是掌握这一技术的关键。GitHub 加速示例工程深度学习dia反向传播考量模型训练中的梯度计算dia反向传播考量模型训练中的梯度计算 在dia模型1.6B参数TTS模型的训练过程中梯度计算是影响模型性能与训练效率的核心环节。本文从Transfor人工智能大模型语音音频媒体生成创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考