菜单
LayerNorm 归约算子的完整 Triton Ascend 实现示例。展示两阶段归约模式（统计量计算 → 归一化输出）、标量累加器、分块遍历等技巧。当生成 reduce/normalize 类算子时可参考此示例的代码结构。
用 Codex 或 Claude 帮你安装 复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它检查 Skill 页面并帮你完成安装。
基于 SOC 职业分类
| name | triton-ascend-example-layernorm |
| description | LayerNorm 归约算子的完整 Triton Ascend 实现示例。展示两阶段归约模式（统计量计算 → 归一化输出）、标量累加器、分块遍历等技巧。当生成 reduce/normalize 类算子时可参考此示例的代码结构。 |
| category | example |
| version | 1.0.0 |
| metadata | {"backend":"ascend","dsl":"triton_ascend","hardware":"Atlas A2, Atlas A3","operator_type":"reduce"} |
import torch
import triton
import triton.language as tl
@triton.jit
def layernorm_kernel(
X_ptr, Y_ptr,
batch_size: tl.constexpr, feature_size: tl.constexpr,
eps: tl.constexpr,
BLOCK_SIZE: tl.constexpr, CORE_NUM: tl.constexpr,
):
core_id = tl.program_id(0)
for batch_idx in range(core_id, batch_size, CORE_NUM):
batch_offset = batch_idx * feature_size
# Phase 1: compute mean & variance
mean_acc = 0.0
var_acc = 0.0
for i in range(0, feature_size, BLOCK_SIZE):
offsets = batch_offset + i + tl.arange(0, BLOCK_SIZE)
mask = offsets < batch_offset + feature_size
x = tl.load(X_ptr + offsets, mask=mask, other=0.0)
mean_acc += tl.sum(x, axis=0)
var_acc += tl.sum(x * x, axis=0)
mean_val = mean_acc / feature_size
std_val = tl.sqrt(var_acc / feature_size - mean_val * mean_val + eps)
# Phase 2: normalize
for i in range(0, feature_size, BLOCK_SIZE):
offsets = batch_offset + i + tl.arange(0, BLOCK_SIZE)
mask = offsets < batch_offset + feature_size
x = tl.load(X_ptr + offsets, mask=mask, other=0.0)
tl.store(Y_ptr + offsets, (x - mean_val) / std_val, mask=mask)
class ModelNew(torch.nn.Module):
def __init__(self):
super().__init__()
try:
self.VEC_CORE_NUM = torch_npu.npu.npu_config.get_device_limit(0).get("vector_core_num", 40)
except:
self.VEC_CORE_NUM = 40
def forward(self, x):
shape = x.shape
batch_size = shape[0]
feature_size = 1
for s in shape[1:]:
feature_size *= s
if not x.is_contiguous():
x = x.contiguous()
y = torch.empty_like(x)
BLOCK_SIZE = 1024
grid = (self.VEC_CORE_NUM,)
layernorm_kernel[grid](x, y, batch_size, feature_size, 1e-5,
BLOCK_SIZE=BLOCK_SIZE, CORE_NUM=self.VEC_CORE_NUM)
return y
矩阵乘法矩阵乘法 A[M, K] @ B[K, N] = C[M, N]中，大K维度矩阵乘法(K>>M,N)优化：针对M/N较小但K极大(如M=N=256,K=131072)的场景，Split-K切分K维度并行化、Workspace+Reduce替代全局同步，实现显著性能提升
Triton Ascend hard API restrictions and forbidden syntax. MUST-follow rules that apply to every kernel: forbidden control flow (return/break/continue/lambda/while), tensor slice/index restrictions, scalar conversion rules, BLOCK_SIZE upper bound. Violating any of these produces a compile or runtime error on Ascend.
Triton Ascend 性能优化通用策略: BLOCK_SIZE 选择 (1024-2048 for elementwise, must be <65536), grid configuration (use VEC_CORE_NUM / CUBE_CORE_NUM, 2D/3D grid for matmul / conv / reduce, 1D grid + inner loop for elementwise / pointwise), 256B alignment for memory transfers, autotune block-size patterns, fp16 / fp32 precision conversion. Bind via keywords like matmul, elementwise, reduce, block_size, grid, autotune, alignment, fp16, fp32, tile, interleaved-loop, cube-core, vec-core.
通过 adaptive_search 或 evolve 搜索式 workflow 生成优化算子。 后台 silent mode 执行，轮询监控进度。
适用于归约(reduce)类算子和含归约子步骤的复合算子（如归一化）的优化指南。典型算子包括：sum, mean, max, min, prod, argmax, argmin, cumsum, cumprod, softmax, logsoftmax, layernorm, rmsnorm, groupnorm, instancenorm, batchnorm, l1norm, l2norm, frobeniusnorm, var, std, average_pooling, sum_pooling 等。特别重要：当归约维度不是最后一维（如 dim=1 归约 shape=[B,F,D1,D2]），需要正确处理多维索引和两阶段归约。包含 PyTorch normalized_shape 多轴归一化语义说明。不适用于纯逐元素运算或矩阵乘法。如果算子是损失函数（先逐元素计算再全局归约），应选择 elementwise-reduce-fused 指南。
CPU C++ 算子核心概念、标准结构模式、KernelBench 代码规范和内嵌扩展方法