菜单
Softmax 归约算子的完整 Triton Ascend 实现示例。展示三阶段归约模式（求 max → 求 sum(exp) → 归一化）、分块累加、标量累加器精度提升等技巧。当生成 reduce 类算子时可参考此示例的代码结构。
默认使用会先检查来源的 Prompt；你也可以切换为直接命令，或下载本地副本。
决定是否安装前，请先阅读 SKILL.md，以及 SkillsMP 当前展示的配套文件。
用 Codex 或 Claude 帮你安装 复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它检查 Skill 页面并帮你完成安装。
直接命令不会经过审查 Prompt；运行前请先检查来源。
npx skills add https://github.com/mindspore-ai/akg --skill triton-ascend-example-softmax命令会保持在同一行。复制前请横向滚动并检查完整内容。
想先保存到本地？可下载 SkillsMP 当前能够提供的文件。
矩阵乘法矩阵乘法 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.
基于 SOC 职业分类
正在显示 SKILL.md
| name | triton-ascend-example-softmax |
| description | Softmax 归约算子的完整 Triton Ascend 实现示例。展示三阶段归约模式（求 max → 求 sum(exp) → 归一化）、分块累加、标量累加器精度提升等技巧。当生成 reduce 类算子时可参考此示例的代码结构。 |
| 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 softmax_kernel(
X_ptr, Y_ptr,
B: tl.constexpr, N: tl.constexpr,
stride_xb: tl.constexpr, stride_xn: tl.constexpr,
stride_yb: tl.constexpr, stride_yn: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr, CORE_NUM: tl.constexpr,
):
pid = tl.program_id(0)
for b in range(pid, B, CORE_NUM):
# Phase 1: max
max_val = -float('inf')
for off in range(0, N, BLOCK_SIZE_N):
n_off = off + tl.arange(0, BLOCK_SIZE_N)
mask = n_off < N
x = tl.load(X_ptr + b * stride_xb + n_off * stride_xn,
mask=mask, other=-float('inf'))
max_val = tl.maximum(max_val, tl.max(x, axis=0))
# Phase 2: sum(exp(x - max))
sum_val = 0.0
for off in range(0, N, BLOCK_SIZE_N):
n_off = off + tl.arange(0, BLOCK_SIZE_N)
mask = n_off < N
x = tl.load(X_ptr + b * stride_xb + n_off * stride_xn,
mask=mask, other=0.0)
exp_x = tl.math.exp(x - max_val)
sum_val += tl.sum(exp_x, axis=0).to(tl.float32)
# Phase 3: normalize
for off in range(0, N, BLOCK_SIZE_N):
n_off = off + tl.arange(0, BLOCK_SIZE_N)
mask = n_off < N
x = tl.load(X_ptr + b * stride_xb + n_off * stride_xn,
mask=mask, other=0.0)
result = tl.math.exp(x - max_val) / sum_val
tl.store(Y_ptr + b * stride_yb + n_off * stride_yn,
result, 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):
if not x.is_contiguous():
x = x.contiguous()
B, N = x.shape
y = torch.empty_like(x)
BLOCK_SIZE_N = 4096
grid = (self.VEC_CORE_NUM,)
softmax_kernel[grid](
x, y, B, N,
x.stride(0), x.stride(1), y.stride(0), y.stride(1),
BLOCK_SIZE_N=BLOCK_SIZE_N, CORE_NUM=self.VEC_CORE_NUM)
return y