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 + QUANT 融合算子及其性能优化，支撑 xxx 系列模型 A5 MXFP8 量化训练性能提升。

Origin（信息来源）

2012

Benefit / Necessity （价值/作用）

支撑 xxx 系列模型 A5 MXFP8 量化训练性能提升。

Design（设计方案）

将 sigmoid + broadcast Mul × 2 + Add + MXFP8 quant 融合为一次算子下发；
不在 HBM 中物化两路 gated attention 大张量和量化前的完整融合结果；
双轴版本一次生成 row-wise 和 col-wise MXFP8 数据，供需要保留输入列方向数据的训练反向使用；推理等仅需 row-wise 输出的场景通过单轴模式（dual_axis_flag=false）复用同一融合算子，跳过 col-wise 链路。

支持两种量化模式（由 dual_axis_flag 参数选择）：
双轴模式（dual_axis_flag=true）：一次生成 row-wise（[1,32] block）及 col-wise（[32,1] block）两套量化数据，返回 row_data_out、row_scale_out、col_data_out、col_scale_out；
单轴模式（dual_axis_flag=false，默认）：仅生成 row-wise（[1,32] block）一套量化数据，返回 row_data_out、row_scale_out，col_data_out、col_scale_out 返回空 tensor，算子内部跳过 col-wise 量化链路。单轴模式面向推理等仅需 row-wise 输出的场景，避免 col-wise 量化与输出的额外计算和搬运开销。
torch接口
row_data_out, row_scale_out, col_data_out, col_scale_out = torch_npu.npu_cla_gate_quant(
global_attn, local_attn, global_gate_logits, local_gate_logits,
dst_type=torch.float8_e4m3fn, round_mode="rint", scale_alg=1, dual_axis_flag=false,
)