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

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

Degrees of freedom: medium-high (computation logic must not be wrong; implementation approach 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
Look up API signature	triton-api-reference.md	—

MANDATORY: Read the above files in full before the corresponding stage, with no line limit. Load trigger: When you need 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 constraints. Describe the computation process in pseudocode,must be confirmed with the user。

2. Implement 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：

1. grid = 物理核数（get_npu_aicore_num() or get_npu_vectorcore_num()）
2. In-core loops handle multiple tasks, and each core computes the data it needs to process (load balancing)
3. 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 by 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), ensuring the kernel can compile and run and the result is correct.Comprehensive precision validationis handled by the precision-eval skill.

Anti-pattern checklist (NEVER)

• ❌ Writing code without confirming the computation logic
• ❌ Ignoring UB size (192KB)
• ❌ Not upcasting precision to FP32 for reductions / using int64
• ❌ grid > 65535 / grid ≠ number of physical cores
• ❌ Using third-party libraries in the kernel / element-wise computation
• ❌ Matrix multiplication degenerating into Vector Core element-wise multiply-add——Cube Core matrix throughput is far higher than Vector Core (dozens of times). Even GEMV (N=1) must use tl.dot + pad B to BLOCK_N=16, and 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 operators 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 separately compute statistics and normalization) — should be single-pass: after one load, tl.sum(x, 1) completes the statistics and normalization within the UB, see template 1
• ❌ Not using diagonal scheduling for large matrix multiplication matrices— for large matrices (BLOCK_THRESHOLD ≥ 4), using sequential scheduling 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 necessarily fail. Alternative:while + boolean flag, or tl.where mask
• ❌ tensor[i] index operations(read/assignment/slicing) — Triton tensors do not support Python-style subscripting. Alternative:tl.where assignment,tl.gather read,tl.extract_slice slicing

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 is inaccurate	Upcast to FP32 before reduction
Insufficient indexing	D-cache error	Use int64 for very large shapes
Reduction-type operators are 5x+ slower	aiv_scalar > 80%	Single-pass: one load + tl.sum(x, 1) all computation within the 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

A Triton kernel must be used to complete the core computation; the following behaviors are strictly prohibited:

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