triton-operator-performance-optim

Triton Operator Performance Optimization (Ascend NPU)

Bottom Line (Must Not Be Breached)

1. Accuracy: After optimization, rtol=1e-3, atol=1e-3 aligned with PyTorch-NPU. If it does not pass, roll back.
2. Generality: Support all original input shapes and dtypes, no hardcoding of specific sizes.

Performance Ratio Definition：Ratio = torch_npu 耗时 / Triton 耗时(reciprocal of elapsed time). Ratio > 1.0 means Triton is faster.

Priority: Correctness > Generality > Performance.

Optimization Workflow

Phase 0: Algorithm Review

Review the algorithm itself before optimizing. Even with optimization, an inefficient algorithm has inherent shortcomings.

Phase 1: Layered Evaluation

1. Quick Screening：time.time() Measure end-to-end (covering small/medium/large), if it meets the target then done
2. Precise Diagnosis: When it does not meet the target, use msprof to measure kernel-side elapsed time and locate the real bottleneck

Phase 2: Bottleneck Optimization

Bottleneck	Optimization Focus
Memory-Bound	Vectorized memory access, UB cache reuse, operator fusion
Compute-Bound	Cube adaptation, Block size tuning
Latency-Bound	Increase parallelism, reduce synchronization

The Basic Four Axes(in order): Block/Grid Size → Contiguous memory access → UB reuse → Compile-time constants

Load：optimization-patterns.md, ascend-terminology.md

Phase 3: Hardware Specialization

• Cube: BLOCK_M/N/K as multiples of 16, accumulator FP32
• UB: Total buffer size < 192KB, single-value buffer 32B aligned
• Grid: 1D Grid ≤ number of physical cores, loop within the core to process multiple rows
• Diagonal Scheduling: For large matrices (above BLOCK_THRESHOLD) use diagonal Grid scheduling to improve L2 cache hit rate
• Multiple Vector Cores: post-dot operations use tl.parallel(bind_sub_block=True) to distribute to 2 vector cores

Load：triton-ascend-api.md, tiling-strategies.md, triton-api-reference.md

Phase 4: Advanced Optimization (As Needed)

Operator fusion, Double Buffer (tl.multibuffer(tensor, 2) or tl.compile_hint(tensor, "multi_buffer", 2)）

Phase 5: Verification (MANDATORY)

Accuracy + Generality + Performance + End-to-end regression

Anti-pattern Checklist (NEVER)

• ❌ Make optimization decisions based on data of a single scale only
• ❌ Directly optimize the kernel when end-to-end does not meet the target (should first confirm the bottleneck with msprof)
• ❌ Sacrifice accuracy for performance / hardcode to break generality
• ❌ FP16 direct reduction / Non-16-multiple BLOCK for matrix multiplication
• ❌ BLOCK_SIZE exceeding UB (192KB) / Non-contiguous memory access
• ❌ Use in the hot path tensor.item()(triggers CPU-NPU synchronization)
• ❌ Use if branches inside a loop to modify variables(Triton compiles to masked operations, catastrophic performance degradation)
• ❌ For 2D Tiling only count the UB of the data buffer(must include offset/mask/index arrays)
• ❌ Use precomputed offset tensors for 2D broadcasting(triggers the compiler addptr multi-user assertion)
• ❌ Use broadcast stride inside the kernel to access auxiliary tensors(cos/sin, etc.), should be changed to host-side expand+contiguous. The extra memory from expand < the performance loss of non-contiguous memory access
• ❌ Reduction-type operators use double pass(loading the same data multiple times to separately compute statistics and normalization) — should be single pass, all computation within UB after one load, see Pitfall 8
• ❌ When comparing with msprof, two kernels have the same name— msprof aggregates by OP Type, same names will be mixed together
• ❌ Large matrices do not use diagonal scheduling(L2 cache thrashing, must be enabled above BLOCK_THRESHOLD)

Checklist

• Precision aligned with PyTorch-NPU (rtol=1e-3, atol=1e-3)
• Non-aligned dimensions and boundaries pass
• Performance tests cover small/medium/large
• grid ≤ number of physical cores, BLOCK_SIZE is a compile-time constant
• Buffer < 192KB, all load/store have Mask
• Reduction promotes to FP32, matrix multiplication BLOCK is a multiple of 16
• Whether reduction-type operators are single pass (required when D ≤ UB)

Common Bottleneck Quick Reference

msprof Metric	Bottleneck	Typical Optimization
aiv_scalar > 80%	Scalar Bound	Check for double pass / per-row loop + tl.where accumulation, change to single pass
aiv_mte2 > 50%	Memory Bound	Contiguous memory access, expand+contiguous, increase BLOCK
aiv_vec > 50%	Compute Bound	Algorithm optimization, reduce redundant computation
aic_cube_ratio < 50%	Low Cube utilization	Check alignment (512B/element size), whether BLOCK is a multiple of 16, use compile_hint('dot_pad_only_k')