chunk_gla_fwd_o_gk：Triton 算子 NPU 迁移实践与交付（issue 提交）

交付 MR：Ascend/MindSpeed-Ops !92 — head 17b843e968ab05586b95db62a2c87562a0e52836，base master；当前 open，等待评审与合入。

交付内容：生产源码（mindspeed_ops/api/triton/、mindspeed_ops/arch32/triton/）+ 仓内 UT（tests/unit_tests/triton/）+ ATK 三件套（tests/atk_tests/triton/）+ 文档（docs/triton/chunk_gla_fwd_o_gk.md）。

验收证据（四项）

针对 review comment #192727253：

需要补充以下信息：
（1）ATK通过截图
（2）UT通过截图
（3）NPU与GPU性能对比达标数据
（4）NPU上的triton算子与GPU上的triton算子的精度对比情况或者GPU上该pytorch算子与triton算子的精度对齐情况

已在远端 head 17b843e968ab05586b95db62a2c87562a0e52836（本 MR 当前 head，与复验工作树逐字节一致）上逐项复验，证据如下。

（1）ATK 通过截图

ATK 26.9.24，物理 NPU 6（Ascend910_9382 physical=6/card=3/chip=0），40 例：

说明：memory_device 已定量归因为框架启动 workspace —— Triton-Ascend 驱动按 workspace_size * grid 每次 launch 分配，编译 workspace_size 22016–81920B；候选/参考内存比随 grid 单调上升（grid≤2 为 0.85–1.02、grid=4 为 1.03–1.11、grid 6..64 为 1.15–1.44），不是候选可移除缓冲；ATK 三项均不在合并流水线门禁内（流水线只跑仓库原生 UT）。

（2）UT 通过截图

仓库原生 tests/unit_tests/triton/test_chunk_gla_fwd_o_gk.py：19 passed, 1 warning in 3.71s（物理 NPU 6，torch_npu 先于 triton 导入；即仓库 CI stage3 的合并门禁）。

（3）NPU 与 GPU 性能对比达标数据

口径：同一评测记录、同一窗口背靠背；任务书“相对原 Triton 算子提升 20%”的判据分母是 NPU 侧 triton_ascend_reference。

• 窗口：候选 eval_20260929_080625，参考刷新 baseline_20260929_075143（warmup 3 / repeat 5）。
• geomean(原 triton(NPU) / 候选) = 1.4754，min = 1.0648，≥1.2 覆盖 18/20。
• 与 GPU(H100 triton_cuda_reference) 的绝对对照：geomean(候选 / H100) = 1.8164（候选比 H100 慢约 1.8×，为 NPU/H100 硬件代差，不是任务书判据）。

（4）精度对齐

以 fp64 CPU golden（torch） 为共同高精度基准（独立于 NPU 与 GPU 两条执行路径）：

• NPU 候选（本 MR 算子）：20/20 通过 domain-aware 精度检查。
• GPU(H100) triton 参考：15/20 通过；未通过的 5 例为 checker 的 special-domain（相消）判定，normal_mismatch_count=0，无正常值失配。
• 逐 case 相对该 golden 的 max_diff/mere 两侧一致（下表）：NPU 候选与 GPU triton 对同一高精度基准的误差同量级。
• 另：ATK accuracy（NPU 候选 vs 独立 NPU torch 小算子参考）40/40 Pass。

证据：截图存放在提交者 fork 的独立 evidence/pr92-chunk-gla-atk 分支（commit e15f671a8a2390681741e72c0992470092f4d179），不属于本 PR 的源码分支与交付文件 diff；ATK/UT 生成物与日志未提交到算子目录。

实践文档（PR 内 docs/triton/chunk_gla_fwd_o_gk.md 全文）

chunk_gla_fwd_o_gk 算子

概述

chunk_gla_fwd_o_gk 是 GLA（Gated Linear Attention）chunk forward 的输出阶段。
算子将 recurrent state 项与 chunk 内因果 attention 项相加，输出
o[B, T, H, V]：

o = scale * (q * exp_or_exp2(g)) @ h + tril(A) @ v

当前实现面向 arch32 NPU，使用 Triton-Ascend kernel。arch35 暂未提供
实现，调用时会明确抛出 NotImplementedError。

模型接口

def chunk_gla_fwd_o_gk(
    q: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    A: torch.Tensor,
    h: torch.Tensor,
    scale: float = 0.125,
    chunk_size: int = 64,
    use_exp2: bool = False,
) -> torch.Tensor:
    """Run the dense GLA chunk forward output stage."""

输入说明

输入必须位于 NPU，q 与 v 的 dtype 必须一致。当前公开接口只支持 dense、
non-varlen 路径；varlen metadata 不属于本次交付范围。

输出说明

实现原理

每个 Triton program 负责一个 V tile、一个 sequence chunk 和一个 batch/head
组合，按以下步骤计算：

1. 分 K tile 读取 q、g 和 h，计算 (q * exp_or_exp2(g)) @ h，其中 scale 由
   按算子声明的 scale 与输入 dtype 三选一的载载路径引入（见「迁移要点」第 1 条）。
2. 读取当前 chunk 的 A 和 v，通过下三角 mask 计算 tril(A) @ v。
3. 将两部分结果在 fp32 accumulator 中相加，再转换为输出 dtype 写回。
4. sequence、K 和 V 尾块均使用边界检查或显式 mask。

生产 dispatch 只依赖运行时 tensor metadata：

迁移要点

原始 Triton 实现使用一个通用 chunk_gla_fwd_kernel_o，在同一 kernel 中完成
recurrent state 和 chunk-local 两部分计算，并通过 Triton autotune 选择 K/V tile。
本 PR 保持相同数学语义，针对 Ascend 做了以下调整：

1. scale 不再在 half-precision 上预先融合（那样会把已缩放的 qg 操作数舍入两次后送进
   cube）。改为按算子声明的 scale 与输入 dtype三选一：|scale| > 1 走累加后置乘
   （避免折叠后的操作数越出半精度范围）；float16 输入用 qg 的 hi/lo 半精度对驱动两次
   tl.dot（22 位有效位）；bfloat16 输入把缩放后的 qg 留在 fp32 并用 fp32 cube 操作数
   （8 位尾数在两项半精度拆分里只剩 16 位）。
2. 根据 shape metadata 增加 H=1 和 B=1/K=64/V<=64 两个实现分支。
3. 使用确定性的 BK、BV 和 multibuffer 配置，不依赖运行时 autotune 缓存。BK/BV
   只按运行时 K/V 范围选择：K<=32 用 32-lane、K 为 128 的倍数时用 128-lane（把
   reduction 收进单次迭代）、其余 64-lane；V 轴在 V > 64 且 32|V 且 128∤V 时用
   32-lane（该 tile 无 padding），否则 V > 64 用 128-lane、其余 64-lane。
4. 保留原始来源版权声明；UT 使用独立 Torch 小算子 reference 验证数学语义。

生产 API 只启动本 PR 自有 Triton-Ascend kernel，不导入测试 reference，不调用
Torch、CPU、供应商整算子或其它后端 fallback。

使用示例

import math

import torch

from mindspeed_ops.api.triton.chunk_gla_fwd_o_gk import chunk_gla_fwd_o_gk

B, T, H, K, V, C = 1, 128, 4, 64, 64, 64
q = torch.randn((B, T, H, K), device="npu", dtype=torch.float16)
v = torch.randn((B, T, H, V), device="npu", dtype=torch.float16)
g = torch.randn((B, T, H, K), device="npu", dtype=torch.float32) * 0.08
A = torch.randn((B, T, H, C), device="npu", dtype=torch.float32) * 0.08
h = torch.randn((B, math.ceil(T / C), H, K, V), device="npu", dtype=torch.float32) * 0.08

o = chunk_gla_fwd_o_gk(q, v, g, A, h, scale=K**-0.5, chunk_size=C)
print(o.shape)

适用场景

适用于 GLA dense、non-varlen forward 中，按固定 chunk 合并 recurrent state
贡献和 chunk 内因果 attention 贡献的输出阶段。

注意事项

1. 当前仅支持 arch32 NPU；arch35 明确报不支持。
2. q/v 仅支持 float16 或 bfloat16，g/A/h 必须为 float32。
3. chunk_size 只接受非布尔整数 32 或 64，use_exp2 只接受布尔值。
4. UT/ATK reference 只用于测试，生产路径不存在 Torch、CPU、reference、
   供应商整算子或其它后端 fallback。