clipped_swiglu 算子

概述

clipped_swiglu 是一个基于 Triton 实现的融合 SwiGLU-gate 算子，对应 GPT-OSS 风格的 _apply_gate。算子沿最后一维将输入切分为 gate 与 up（非交错布局：gate=[..., :K]，up=[..., K:]），对 gate 仅做上限 clamp、对 up 做双向 clamp，再融合 SwiGLU 门控计算。该算子利用昇腾 NPU 的并行计算能力实现高性能的逐元素融合运算，内部以 fp32 计算、回写输入 dtype，兼顾精度与性能。

---

函数签名

def clipped_swiglu(
    gate_up: torch.Tensor,
    swiglu_alpha: float = 1.702,
    swiglu_limit: float = 7.0,
) -> torch.Tensor:
    """执行融合 SwiGLU-gate 前向计算

    Args:
        gate_up: 输入张量，末维为偶数，布局为 [gate | up]
        swiglu_alpha: SwiGLU 的 alpha 参数
        swiglu_limit: clamp 上下限 L

    Returns:
        out: SwiGLU-gate 输出，末维为输入末维的一半，其余维度不变
    """

---

参数说明

前向输入

参数	类型	形状	描述	是否必须
gate_up	torch.Tensor	[..., 2K]	输入张量，末维为偶数，前 K 列为 gate、后 K 列为 up	是
swiglu_alpha	float	-	SwiGLU 的 alpha 参数	否，默认 1.702
swiglu_limit	float	-	clamp 上下限 L，gate 上限 clamp 到 L，up 双向 clamp 到 [-L, L]	否，默认 7.0

前向输出

参数	形状	描述
out	[..., K]	SwiGLU-gate 输出，末维为输入末维的一半，其余维度与输入一致，dtype 与输入一致

---

实现原理

前向计算

对于输入末维 2K，切分后对每个元素计算：

g  = min(gate, L)                         # gate 仅上限 clamp（忠实原实现）
u  = clamp(up, -L, L)
out = (u + 1) * g * sigmoid(g * alpha)

算子通过以下步骤实现：

1. 视图展开：将输入 gate_up 视图展开为二维 [M, 2K]，其中 2K 为末维长度，M 为其余维度展平
2. 地址计算：gate 块地址为 row*2K + col，up 块地址为 gate 块 + K，全部用乘加计算，避免除法/取模（昇腾后端上代价高）
3. 内部 fp32 计算：load 后升精到 fp32，依次做 min/clamp/sigmoid/乘加，其中 sigmoid 等超越函数需高精度
4. 回写：结果量化回输入 dtype 后 store 到输出张量
5. 2D launch + 手动分档：grid = (cdiv(M, BLOCK_M), cdiv(K, BLOCK_K))，BLOCK_M/BLOCK_K 由 API 层按 M/K/dtype 手动分档选定（非 autotune，避免模型调用时触发 tune 引入超长等待），且受 ascend910 UB（192KiB，auto-multi-buffer 开启）约束——fp32 的 BLOCK_K 上限 1024、bf16/fp16 上限 4096；常量 alpha/limit 经 constexpr 折叠，换值会重编译

Kernel 设计

设计决策	做法	原因
2D launch + 乘加地址	grid 按 (BLOCK_M, BLOCK_K) 切分，地址用 row*2K + col	避免 1D launch 所需的 //K/%K 除法取模，ascend 后端上代价高
BLOCK_K 按 dtype 分档	bf16/fp16 上限 4096、fp32 上限 1024；取 min(覆盖K的2幂, dtype上限)	ascend910 的 auto-multi-buffer 会复制 tile buffer，fp32（4B/elem）tile 过大会触发 UB 溢出
BLOCK_M 按 dtype+M 分档	按 M 增长（1/2/4/8/16），再按 dtype UB 预算收敛（fp32 预算 2048、bf16/fp16 预算 4096，使 BLOCK_M×BLOCK_K 不超预算）	bf16/fp16 在 BLOCK_K=4096 时 BLOCK_M 收敛到 1，fp32 在 BLOCK_K=1024 时 BLOCK_M≤2，否则 UB 溢出
内部 fp32 计算	load 后 .to(tl.float32)，store 时量化回	sigmoid 等超越函数需高精度；bf16 输入下保证精度
不传 num_warps/num_stages	Config 仅含 BLOCK_M/BLOCK_K	ascend 后端不支持调这两参数

---

使用示例

import torch
from mindspeed_ops.api.triton.clipped_swiglu import clipped_swiglu

B, T, K = 4, 2048, 4096
x = torch.randn(B, T, 2 * K, device='npu', dtype=torch.bfloat16)

y = clipped_swiglu(x, swiglu_alpha=1.702, swiglu_limit=7.0)
assert y.shape == (B, T, K)

---

性能对比

测试环境：Ascend 910_9382（64GB HBM，单卡），torch 2.9.0 + torch_npu，triton 3.2.0 + ascend 后端。

测量方法：host 侧 time.perf_counter 计时（25 warmup + 100 rep 取平均，每次 torch.npu.synchronize()）。device 侧性能另由 ATK performance_device 任务采集（profiler 解析 device 耗时），60 用例平均 device 性能比 4.98×、device 性能达标 Pass。

bf16

shape	eager(ms)	triton(ms)	speedup	bw(GB/s)
(1, 8192)	0.056	0.029	1.96×	863
(128, 8192)	0.066	0.029	2.27×	108912
(1024, 8192)	0.221	0.052	4.26×	484314
(8192, 8192)	1.685	0.394	4.28×	511687
(1, 4096, 24576)	2.902	0.541	5.37×	558400
(1, 8192, 24576)	5.798	1.072	5.41×	563449
(4, 4096, 24576)	11.536	2.138	5.40×	565097

fp16

shape	eager(ms)	triton(ms)	speedup	bw(GB/s)
(1024, 8192)	0.240	0.072	3.35×	352026
(8192, 8192)	1.686	0.487	3.47×	413849
(1, 4096, 24576)	2.687	0.602	4.47×	501783
(4, 4096, 24576)	10.903	2.226	4.90×	542669

fp32

shape	eager(ms)	triton(ms)	speedup	bw(GB/s)
(1024, 8192)	0.296	0.117	2.54×	431975
(8192, 8192)	3.790	0.542	7.00×	743185
(1, 4096, 24576)	5.596	0.747	7.49×	808788
(1, 8192, 24576)	10.569	1.283	8.24×	941380
(4, 4096, 24576)	21.150	2.576	8.21×	937771

3D shape (1, 4096, 24576) 对应 K=12288（LLM hidden 常见尺寸），三 dtype 加速：fp32 8.2× / bf16 5.4× / fp16 4.5×。bf16/fp16 的瓶颈为 fp32 超越函数（sigmoid）计算，fp32 已近该算子上限。

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

---

常见训练 Shape

以下为 LLM 训练时的典型输入形状（末维 2K，K 为 hidden 维度）：

B	T	K	gate_up / out 形状
4	2048	4096	[4, 2048, 8192] / [4, 2048, 4096]
8	2048	4096	[8, 2048, 8192] / [8, 2048, 4096]
1	4096	12288	[1, 4096, 24576] / [1, 4096, 12288]
4	4096	12288	[4, 4096, 24576] / [4, 4096, 12288]

---

注意事项

1. 平台限制：当前算子在 arch35（950 系列）平台上不支持，调用会抛出 NotImplementedError；仅在 arch32（910 系列）上验证
2. 末维约束：输入 gate_up 末维必须为偶数（布局 [gate | up]），否则抛出 ValueError
3. dtype 限制：仅支持 float32/float16/bfloat16，其他 dtype 抛出 ValueError
4. 连续性：kernel 调用前对输入做连续性校验，非连续输入会 .contiguous() 后再下发
5. 常量重编译：swiglu_alpha/swiglu_limit 经 constexpr 折叠，换值会触发重编译，建议按模型固定取值
6. 仅 forward：当前实现仅含前向，无 backward；如需训练反传，需上层自行处理或后续补充
7. ATK 泛化测试说明：tests/atk_tests/triton/clipped_swiglu/ 中的 ATK 用例将非张量参数固定为 alpha=1.702、limit=7.0，仅通过张量输入 shape/dtype 覆盖泛化测试
8. 对于非常小的张量，可能不会观察到明显的性能提升，因为启动内核的开销可能超过计算收益