external-gitcode-ascend-triton-operator-code-gen

Triton operator code generation

Core principles

Computation logic → Tiling strategy → Code implementation. The order must not be reversed.

Triton API path priority

tl.tanh / tl.erf / tl.sqrt — recommended, best performance
tl.math.tanh / tl.math.erf — alternative
triton.language.extra.ascend.libdevice / triton.language.extra.libdevice — when none of the above are found

Degrees of freedom: medium-high (computation logic must not be wrong; implementation method is flexible)

Reference resource loading

Stage	Must load	Do not load
Design Tiling	`hardware-architecture.md`	templates.md
Generate Kernel	`templates.md`	hardware-architecture.md
Find API signature	`triton-api-reference.md`	—

MANDATORY: read the above files in full before the corresponding stage, with no line limit. Load trigger: when it is necessary to confirm the signature or parameters of a certain tl/libdevice API, read in full triton-api-reference.md。

Workflow

1. Understand requirements → confirm computation logic

Extract the mathematical formula, input/output specifications, and constraint conditions. Describe the computation process in pseudocode;must confirm with the user。

2. Implement the Tiling strategy

Input dependency: if a design document already exists (produced by the design skill), implement according to its Tiling strategy; if not, design it here.

Two principles for inter-core partitioning：

grid = 物理核数（get_npu_aicore_num() or get_npu_vectorcore_num()）
intra-core loop processes multiple tasks, and each core computes the data it is to process (load balancing)
a grid exceeding the number of physical cores will have the excess portion issued serially in a loop, which only adds overhead and cannot increase parallelism

core_num = get_npu_aicore_num()
grid = (core_num,)
# kernel 内：
pid = tl.program_id(0)
num_core = tl.num_programs(0)
blocks_per_core = tl.cdiv(total_blocks, num_core)
for block_idx in range(pid * blocks_per_core, min(...)):
    ...

UB space: safe BLOCK_SIZE = (196608 - 32) / (缓冲区数 × dtype大小) × 0.8

3. Generate Kernel code

Select a template according to the operator type (see templates.md for details):

Operator type	Core type	Template
Reduction	vector core	Template 1
GEMM/attention	AI core	Templates 2,6
Activation/loss/indexing/MoE/post-processing	vector core	Templates 3-5,7-8
Convolution	AI core	Template 9

4. Generate a basic correctness test

Generate abasic correctness test(a smoke test with a single shape × single dtype) to ensure the kernel can compile and run and that the results are correct.Comprehensive precision verificationis the responsibility of the precision-eval skill.

Anti-pattern checklist (NEVER)

❌ Writing code without confirming the computation logic
❌ Ignoring UB size (192KB)
❌ Reduction without increasing precision to FP32 / using int64
❌ grid > 65535 / grid ≠ number of physical cores
❌ Using third-party libraries in the kernel / element-by-element computation
❌ Matrix multiplication degenerating into Vector Core element-by-element multiply-add——Cube Core matrix throughput is far higher than Vector Core (by dozens of times). Even GEMV (N=1) must use tl.dot + pad B to BLOCK_N=16; it is forbidden to use Vector Core to compute inner products row by row
❌ Using GPU-specific parameters (num_stages/num_warps/num_ctas) on the NPU
❌ Implementing the operator with PyTorch instead of Triton
❌ Not testing operator correctness / not testing on the NPU
❌ Using two-pass mode for reduction-type operators(loading the same data multiple times to compute statistics and normalization separately) — should be single-pass: after one load, tl.sum(x, 1) complete the statistics and normalization within UB, see template 1
❌ Not using diagonal scheduling for large matrix multiplication matrices— Using sequential scheduling for large matrices (BLOCK_THRESHOLD ≥ 4) causes L2 cache thrashing; diagonal scheduling should be used, see template 2
❌ for/while inside the loop return / break — Triton IR structured control flow does not support early exit, and compilation will always fail. Alternatives:while + boolean flag, or tl.where mask
❌ tensor[i] indexing operations(read/assign/slice) — Triton tensors do not support Python-style subscripting. Alternatives:tl.where assignment,tl.gather read,tl.extract_slice slice

Common pitfalls

Pitfall	Symptom	Solution
Incorrect computation logic	Output does not match expectations	Describe with pseudocode and confirm
UB overflow	"ub overflow"	Reduce BLOCK_SIZE
coreDim exceeds limit	"coreDim > UINT16_MAX"	Increase BLOCK_SIZE or set `TRITON_ALL_BLOCKS_PARALLEL=1`
Precision loss	FP16 inaccurate	Upcast to FP32 before reduction
Insufficient indexing	D-cache error	Use int64 for very large shapes
Reduction-type operators 5x+ slower	aiv_scalar > 80%	Single pass: one load + `tl.sum(x, 1)` all computation within UB
L2 cache thrashing for large matrices	GEMM throughput lower than expected	Diagonal scheduling (`task_m_idx = block_idx % NUM_BLOCKS_M`), see template 2

PyTorch substitution strictly prohibited

The core computation must be implemented with Triton kernels; the following behaviors are strictly prohibited:

Prohibited: using torch.matmul / torch.mm / torch.bmm / torch.einsum to substitute for a matmul kernel
Prohibited: using F.conv1d / F.conv2d / F.conv3d / F.conv_transpose* to substitute for a conv kernel
Prohibited: using high-level interfaces such as nn.Conv1d / nn.Conv2d / nn.Conv3d / nn.LayerNorm / nn.BatchNorm* / nn.GroupNorm / nn.InstanceNorm
Prohibited: using F.layer_norm / F.batch_norm / F.group_norm / F.instance_norm to substitute for a norm kernel
Prohibited: using torch.cumsum / torch.cumprod / torch.flip, etc., to directly substitute for a scan kernel
Prohibited: writing an @triton.jit kernel but not calling it in forward. The only allowed torch calls: tensor allocation (empty/zeros/ones), shape operations (view/reshape/permute/contiguous), data movement (to/npu)