chunk_kda_bwd_wy_dqkg_fused 算子

概述

chunk_kda_bwd_wy_dqkg_fused 是 KDA chunk backward 中用于计算 dq、dk、dv、db、dg 和 dA 的 Triton 融合算子。该算子对应 flash-linear-attention KDA 反向路径中的 wy 表示梯度计算阶段。

当前实现面向 Ascend arch32，将主梯度、dv2/db、dk/dg/db 修正、dA_pre 和 dA 后处理拆成多阶段 kernel，避免在单个 kernel 内完成完整 BT x BT 后处理。

---

函数签名

from mindspeed_ops.api.triton.chunk_kda_bwd import chunk_kda_bwd_wy_dqkg_fused

def chunk_kda_bwd_wy_dqkg_fused(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    v_new: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    A: torch.Tensor,
    h: torch.Tensor,
    do: torch.Tensor,
    dh: torch.Tensor,
    dv: torch.Tensor,
    scale: Optional[float] = None,
    cu_seqlens: Optional[torch.Tensor] = None,
    chunk_size: int = 64,
    chunk_indices: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """计算 KDA chunk backward 的 dq/dk/dv/db/dg/dA。"""

返回顺序固定为 (dq, dk, dv, db, dg, dA)。

---

参数说明

输入

参数	形状	dtype	描述
q	[B, T, H, K]	float16 / float32	query
k	[B, T, H, K]	float16 / float32	key
v	[B, T, H, V]	float16 / float32	value
v_new	[B, T, H, V]	float16 / float32	KDA 前向中间结果
g	[B, T, H, K]	float32	gate log 值
beta	[B, T, H]	float16 / float32	gate beta
A	[B, T, H, BT]	float16 / float32	chunk 内三角矩阵
h	[B, NT, H, K, V]	float32	chunk state
do	[B, T, H, V]	float16 / float32	输出梯度
dh	[B, NT, H, K, V]	float32	chunk state 梯度
dv	[B, T, H, V]	float16 / float32	value 梯度输入
scale	scalar	float	dq 缩放系数
cu_seqlens	[N + 1]	int32 / int64	可选,变长序列累计长度
chunk_size	scalar	int	chunk 长度 BT
chunk_indices	[NT, 2]	int32 / int64	可选,变长序列 chunk 索引

NT = ceil(T / chunk_size)。当前已验证路径要求 q/k/v/v_new/beta/A/do/dv 使用同一种主输入 dtype，g/h/dh 固定为 float32。

输出

输出	形状	dtype	描述
dq	[B, T, H, K]	float32	q 的梯度
dk	[B, T, H, K]	float32	k 的梯度
dv	[B, T, H, V]	同输入 v	v 的梯度
db	[B, T, H]	float32	beta 的梯度
dg	[B, T, H, K]	float32	g 的梯度
dA	[B, T, H, BT]	float32	A 的梯度

---

支持的数据类型

参数	支持的 dtype	说明
q / k / v / v_new / beta / A / do / dv	float16 / float32	主输入需保持同一种 dtype
g / h / dh	float32	当前 UT/ATK 固定验证路径
输出 dv	float16 / float32	dtype 跟随输入 v
其它输出	float32	dq/dk/db/dg/dA 固定为 fp32

---

实现原理

分块计算

算子按 (chunk, batch-head) 维度切分任务。每个 program 处理一个 chunk 和一个 head，内部按 K、V 分块累加。

核心中间量包括:

dq = do @ h^T
dk = v_new @ dh^T
dw = -(dv @ h^T)

随后结合 g、beta、A 对 dq/dk/dg/db/dA 做 KDA wy 反向修正。

多阶段 kernel

阶段	kernel	作用
主梯度	chunk_kda_bwd_kernel_wy_dqkg_fused	计算初始 dq/dk/dg/dw
dv/db	chunk_kda_bwd_kernel_wy_dv2_db	计算 dv2 和第一部分 db
dk/dg/db 修正	chunk_kda_bwd_kernel_wy_dk_dg_db	基于 A @ dw 修正 dk/dg/db
dA 预处理	chunk_kda_bwd_kernel_wy_dA_pre	计算 masked dA_pre
dA 后处理右乘	chunk_kda_bwd_kernel_dA_post_right	计算 dA_pre @ A
dA 后处理左乘	chunk_kda_bwd_kernel_dA_post_left	计算最终 dA = -A @ (dA_pre @ A)

Tiling

当前默认 tiling:

参数	默认值	说明
BK / BV	64	主 kernel 和 dA_pre 的 K/V 分块
BDA	min(BT, 32)	dA 后处理分块,避免 BDA > BT
BM / BN	BDA	dA tile 行列大小

---

使用示例

import torch
import torch.nn.functional as F
from mindspeed_ops.api.triton.chunk_kda_bwd import chunk_kda_bwd_wy_dqkg_fused

B, T, H, K, V, BT = 1, 2048, 32, 128, 128, 64
NT = (T + BT - 1) // BT
device = "npu"
dtype = torch.float16

q = torch.randn(B, T, H, K, device=device, dtype=dtype)
k = torch.randn(B, T, H, K, device=device, dtype=dtype)
v = torch.randn(B, T, H, V, device=device, dtype=dtype)
v_new = torch.randn(B, T, H, V, device=device, dtype=dtype)
g = F.logsigmoid(torch.randn(B, T, H, K, device=device, dtype=torch.float32))
beta = torch.randn(B, T, H, device=device, dtype=dtype).sigmoid()
A = torch.randn(B, T, H, BT, device=device, dtype=dtype)
h = torch.randn(B, NT, H, K, V, device=device, dtype=torch.float32)
do = torch.randn(B, T, H, V, device=device, dtype=dtype)
dh = torch.randn(B, NT, H, K, V, device=device, dtype=torch.float32)
dv = torch.randn(B, T, H, V, device=device, dtype=dtype)

dq, dk, dv_out, db, dg, dA = chunk_kda_bwd_wy_dqkg_fused(
    q, k, v, v_new, g, beta, A, h, do, dh, dv,
    scale=1.0 / (K ** 0.5),
    chunk_size=BT,
)

---

常见训练 Shape

当前模型侧实际关注 shape:

B	T	H	K	V	BT
1	2048	32	128	128	64

---

性能对比

方法	时间消耗
PyTorch 内置逻辑	基准
Triton 实现逻辑	约 5.0x

注：性能提升数据基于昇腾 NPU 环境测试，具体数值可能因硬件配置不同而有所差异。

---

注意事项

1. 主输入 q/k/v/v_new/beta/A/do/dv 需使用同一种 dtype；g/h/dh 固定为 fp32。
2. 当前上库验证覆盖 fp16/fp32，不声明支持 bf16。
3. chunk_size 需与 A.shape[-1] 一致。
4. 当前已验证 BT=16/64，模型 shape 使用 BT=64。
5. 输入张量建议保持 contiguous，最后一维连续。
6. ATK 不包含 model shape；model shape 精度以 UT reference + RMS ratio 标准为准。