Closed
[Requirement|需求建议]: 新增cla_gate_backward算子 #5982
jin-yao312Add on Sep 18Closed at Sep 21
jin-yao312
Sep 18 Comments:
Sep 18 Comments:
/assign @jin-yao312


Sep 18 Assign jin-yao312 as responsible person
jin-yao312
Sep 18 Comments:
Sep 18 Comments:
/assign @fengqiuyue


Sep 18 Assign fengqiuyue as responsible person
Sep 18 Associated pull request
Sep 18 Associated pull request
Sep 18 Deleted associated pull request: 新增cla_gate_backward算子
Sep 18 Modify description of issue
Sep 21 Close issue
Sep 21 Added label: resolved
Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.
Backgroud(背景信息)
模型采用 Cross-Layer Attention(CLA)+ gated SWA 结构,实现 CLA gate 反向融合算子及其性能优化,支撑模型 A5 MXFP8 量化训练性能提升。
Origin(信息来源)
2012
Benefit / Necessity (价值/作用)
支撑 xxx 系列模型 A5 MXFP8 量化训练性能提升。
Design(设计方案)
反向融合边界覆盖 CLA gate merge 的反向:输入上游梯度 G = ∂L/∂O、两路原始 Attention 输出 O_g / O_l 与两路 gate logits z_g / z_l,一次性计算四个梯度并返回;
将 sigmoid×2 + 广播乘 + 逐元素乘 + 沿 D sum×2 融合为一次算子下发;
本接口不做梯度量化:输出保持输入精度(BF16/FP16),不生成 FP8 梯度。反向算子不接收前向 sigmoid 中间结果,s_g / s_l 由算子内重新计算。
torch接口:
grad_global_attn_out, grad_local_attn_out, grad_global_gate_logits_out, grad_local_gate_logits_out = cann_ops_nn.cla_gate_backward(
grad_merged, global_attn, local_attn,
global_gate_logits, local_gate_logits,
input_attn_layout="TND",
)