external-gitcode-ascend-triton-operator-performance-optim

Triton Operator Performance Optimization (Ascend NPU)

Bottom Lines (Must Not Be Violated)

1. Precision: After optimization, rtol=1e-3, atol=1e-3 aligned with PyTorch-NPU. If it does not pass, roll back.
2. Generalization: Support all original input shapes and dtypes; specific sizes must not be hardcoded.

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

Priority: Correctness > Generalization > Performance.

Optimization Workflow

Phase 0: Algorithm Review

Before optimizing, first review the algorithm itself. An inefficient algorithm has inherent shortcomings no matter how much it is optimized.

Phase 1: Layered Evaluation

1. Quick Screening：time.time() Measure end-to-end (covering small/medium/large); if it meets the target, it is complete
2. Precise Diagnosis: When the target is not met, 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

Basic Four Moves(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 are multiples of 16, accumulator FP32
• UB: Total buffer size < 192KB, single-value buffer 32B aligned
• Grid: 1D Grid ≤ number of physical cores, process multiple rows in a core loop
• 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: Validation (MANDATORY)

Precision + Generalization + Performance + End-to-end regression

Anti-Pattern Checklist (NEVER)

• ❌ Make optimization decisions based on a single-scale dataset only
• ❌ Directly optimize the kernel when end-to-end does not meet the target (should first use msprof to confirm the bottleneck)
• ❌ Sacrifice precision for performance / hardcode and break generalization
• ❌ Direct FP16 reduction / non-16-multiple BLOCK for matrix multiplication
• ❌ BLOCK_SIZE exceeds 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 this into masked operations, causing catastrophic performance degradation)
• ❌ 2D Tiling only accounts for the UB of data buffers(must include offset/mask/index arrays)
• ❌ Use precomputed offset tensors for 2D broadcasting(triggers the compiler addptr multi-user assertion)
• ❌ Access auxiliary tensors inside the kernel using broadcast stride(cos/sin, etc.); should be changed to host-side expand+contiguous. The extra memory for expand < the performance loss of non-contiguous memory access
• ❌ Reduction-type operators use two passes(multiple loads of the same data to separately compute statistics and normalization) — should be single pass, after one load perform all computation in UB, see Pitfall 8
• ❌ When comparing with msprof, two kernels have the same name— msprof aggregates by OP Type, and 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 promoted to FP32, matmul 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 two-pass / row-by-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')