wy_fast 算子

概述

wy_fast 是 DeltaNet / Gated-DeltaNet 等线性注意力(flash-linear-attention 系列)所使用的 wy 表示加速算子。在以 chunk 方式扫描序列时,根据每个 chunk 的下三角因子 A 和缩放/门控信号,重计算(recompute)出 w = A @ (k * β * exp(g)) 与 u = A @ (v * β)。

算子以 torch.autograd.Function 包装对外提供:正向产出 (w, u),反向通过 autograd 自动调用 wy 表示反向 kernel,得到 k / v / β / g 的梯度。

---

函数签名

from mindspeed_ops.api.triton.wy_fast import wy_fast

def wy_fast(
    k: torch.Tensor,
    v: torch.Tensor,
    beta: torch.Tensor,
    A: torch.Tensor,
    g: Optional[torch.Tensor] = None,
    gk: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """重计算 wy 表示 (w, u),支持 autograd。

    Args:
        k:    key,形状 [B, T, Hg, K]
        v:    value,形状 [B, T, H, V]
        beta: 缩放因子 β,形状 [B, T, H]
        A:    下三角因子,形状 [B, T, H, BT](BT 为 chunk 大小)
        g:    可选,逐 token 门控,形状 [B, T, H]
        gk:   可选,逐 K 门控,形状 [B, T, Hg, K](仅作用于正向)

    Returns:
        w: 形状 [B, T, H, K]
        u: 形状 [B, T, H, V]
    """

wy_fast 是一个 autograd 包装(WyFastFunction):

• 正向:返回 (w, u)。
• 反向:对 (w, u) 调 .backward() 即可,autograd 自动调用反向 kernel,梯度写入 k.grad / v.grad / beta.grad / g.grad。
• A 与 gk 视为不可导输入 —— 反向 kernel 是 w/u → A(k,β,g) → (k,β,g) 整条链的手写伴随,梯度只回流到 k / v / β / g。

---

参数说明

输入

参数	形状	描述
k	[B, T, Hg, K]	key;Hg 为 kv-head 数
v	[B, T, H, V]	value;H 为 q-head 数
beta	[B, T, H]	缩放因子 β
A	[B, T, H, BT]	下三角因子;BT 为 chunk 大小,取自 A.shape[-1]
g	[B, T, H]	可选,逐 token 门控
gk	[B, T, Hg, K]	可选,逐 K 门控,仅作用于正向

H % Hg == 0:H == Hg 为 MHA,Hg == 1 为 MQA,其余为 GQA。

正向输出

参数	形状	描述
w	[B, T, H, K]	A @ (k * β * exp(g))
u	[B, T, H, V]	A @ (v * β)

反向梯度(autograd)

对 w / u 反向后,梯度写入各输入张量的 .grad:

梯度	写入位置	形状	描述
dk	k.grad	[B, T, Hg, K]	k 的梯度;GQA 场景下已沿 q-head 维归约回 kv-head 维
dv	v.grad	[B, T, H, V]	v 的梯度
dbeta	beta.grad	[B, T, H]	β 的梯度
dg	g.grad	[B, T, H]	g 的梯度

A / gk 无梯度。

---

支持的数据类型

参数	支持的 dtype	说明
k / v / beta / A / g / gk	float32 / bfloat16 / float16	全部浮点张量入参须为同一种 dtype(已验证路径)

架构差异提示:在 arch35 上,反向 kernel(prepare_wy_repr_bwd)当前仅支持 bfloat16,不支持 float16。若在 arch35 上以 float16 输入调用并触发反向,会因反向 kernel 未覆盖 fp16 路径而出现精度问题。

---

实现原理

重计算(正向)

对于每个 chunk 内的位置 i:

w[i] = sum_j A[i, j] * (k[j] * β[j] * exp(g[j]) [* exp(gk[j])])
u[i] = sum_j A[i, j] * (v[j] * β[j])

GQA 场景下,k 按 Hg 维存储,kernel 内部通过 i_kv = i_h // (H // Hg) 把 q-head 索引折回 kv-head 索引,实现 K/V head 共享。

反向

反向对 A = inv(I + M_lt) 隐含的矩阵-逆求导链做完整链式法则:

M_lt[i, j] = β[i] * <k[i], k[j]> * exp(g[i] - g[j])  for i > j else 0
∂L/∂M = -A^T @ ∂L/∂A @ A^T(矩阵逆求导)

由该公式推出 dk / dβ / dg。dv 直接由 du 通过 A^T @ du * β 得到。

Kernel 设计

设计决策	做法	原因
持久化 kernel	grid 固定为 (24,),每个 program 通过 grid-stride loop 处理 NT/24 个 chunk	对齐 Ascend910 的 24 个 AICore,避免细粒度 grid 的 launch 开销
bwd 单 kernel	一个 prepare_wy_repr_bwd_kernel 完成 K-loop / 矩阵-逆修正 / 第二 K-loop / dg 累加全过程	中间结果驻留 UB,不落 HBM
GQA dk reduce	kernel 内 dk 按 q-head 写出 (B, T, H, K) 工作区,orchestrator 末尾 view().sum(3) 折回 (B, T, Hg, K)	用 H 维独立写入避免 atomic_add,reduction 开销 < 5%
fp32 累加器	所有 tl.zeros(..., dtype=tl.float32) 工作区 + tl.dot 默认 fp32 累加	输入支持 fp32/bf16/fp16,内部精度统一
strict-lower mask	m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t)	只对真正属于下三角的位置计算修正

---

使用示例

训练场景

import torch
from mindspeed_ops.api.triton.wy_fast import wy_fast

B, T, H, K, V, BT = 1, 8192, 32, 128, 128, 64
device = "npu"

k    = (torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) * 0.1).requires_grad_(True)
v    = (torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16) * 0.1).requires_grad_(True)
beta = (torch.randn(B, T, H,    device=device, dtype=torch.bfloat16) * 0.1).requires_grad_(True)
g    = (torch.randn(B, T, H,    device=device, dtype=torch.bfloat16) * 0.01).requires_grad_(True)
A    = torch.zeros(B, T, H, BT, device=device, dtype=torch.bfloat16)
# A 内每个 chunk 是下三角 + 单位阵,业务侧通常由上一阶段算子产出

# 正向
w, u = wy_fast(k, v, beta, A, g=g)

# 反向 —— autograd 自动调用反向 kernel
(w.sum() + u.sum()).backward()
dk, dv, dbeta, dg = k.grad, v.grad, beta.grad, g.grad
assert dk.shape == k.shape

---

性能对比

bwd 在 Ascend910B 与 A100-80GB 实测(单 batch, K=V=128, BT=64, fp32, 3 warmup + 10 iter)。 NPU 端 median 与 min 差距 <1%(无 autotune),GPU 端 FLA mainline 含 autotune 装饰器导致前几个 iter 偏慢, 故 NPU 取 median、GPU 取 min 做"已收敛后稳态"对比。NVIDIA 端使用 NVIDIA_TF32_OVERRIDE=0 强制 fp32。

shape	NPU median (ms)	A100 min (ms)	NPU/A100
T=4K H=32 MHA	2.14	1.25	1.70×
T=8K H=32 MHA	3.94	3.02	1.30×
T=32K H=32 MHA	13.69	9.31	1.47×
T=32K H=32 GQA-4:1	14.57	8.83	1.61×
T=64K H=16 MHA	13.88	10.99	1.26×

---

常见训练 Shape

DeltaNet 训练时的典型输入(K=V=128, BT=64):

B	T	H	Hg
1	4096	32	32
1	32768	32	32
1	32768	32	8
1	65536	16	16

---

注意事项

1. H % Hg == 0 是硬性约束,违反会在 orchestrator 抛 ValueError。
2. 需要反向时 g 必传。无门控场景请传零张量(exp(0) = 1,数值上等价无门控)。
3. K、V 维度当前仅在 K=V=128 下测试过;BT(A.shape[-1])仅验证过 64。
4. 输入张量需在最后一维连续。
5. 反向 GQA(H > Hg)场景下,kernel 先按 q-head 写出 (B, T, H, K) 的 dk 中间工作区(dtype 与 k 一致),orchestrator 再 view().sum(3) reduce 回 (B, T, Hg, K),显存峰值比 MHA 多 H/Hg - 1 倍 dk。
6. 输入幅度需较小:算子内部含 A = inv(I + tril(M)) 矩阵求逆,k / β / g 过大会使其发散;已验证范围为各浮点输入量级 ~0.1。
7. arch35 反向 dtype 限制:在 arch35上,反向 kernel 仅支持 bfloat16,不支持 float16。如需在 arch35 上训练,请使用 bfloat16以避免反向精度问题。