triton-kernel-optimization

Triton Kernel Optimization

Purpose

Provide production-validated patterns and tuning tactics for performant Triton kernels on AMD MI-series GPUs.

When to Use

• Authoring new Triton kernels for PyTorch or standalone use
• Porting CUDA/HIP concepts into Triton with equivalent performance
• Profiling and benchmarking Triton kernels

Optimization Priority

Phase 1: Foundation (correct and basic performance)

1. Use @triton.autotune with configs covering key block sizes (64/128/256)
2. Use @triton.heuristics for compile-time optimizations (e.g., EVEN_K)
3. Apply tl.assume for stride positivity to help compiler optimize
4. Separate boundary handling from main computation path
5. Use tl.constexpr for all compile-time constants

Phase 2: Memory Optimization 6. Implement cache modifiers (.ca, .cg) for L2 cache control 7. Use split-K for improved L2 reuse on large K dimensions 8. Apply XCD remapping (remap_xcd) for multi-die GPUs (MI250X, MI300) 9. Optimize GROUP_SIZE_M for better L2 locality 10. Pre-shuffle weight layouts for better vectorization

Phase 3: Advanced Techniques 11. Implement persistent kernels for repeated operations 12. Use attention sink for stable long-context attention 13. Fuse quantization with GEMM (e.g., blockscale + matmul) 14. Apply per-token or per-tensor quantization strategies 15. Use grouped GEMM for mixture-of-experts workloads

Anti-patterns:

• Hardcoding block sizes without autotune
• Ignoring tail handling (non-divisible shapes)
• Not using tl.assume for known constraints
• Excessive register pressure from large tile sizes
• Unnecessary synchronization or atomic operations

Core Optimization Patterns

1. Autotuning and Heuristics

Autotune configuration:

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64,
                       'GROUP_SIZE_M': 8}, num_warps=8, num_stages=4),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32,
                       'GROUP_SIZE_M': 8}, num_warps=4, num_stages=5),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32,
                       'GROUP_SIZE_M': 4}, num_warps=4, num_stages=3),
    ],
    key=['M', 'N', 'K'],  # Tune based on problem dimensions
)
@triton.heuristics({
    'EVEN_K': lambda args: args['K'] % args['BLOCK_SIZE_K'] == 0,
    'GRID_MN': lambda args: triton.cdiv(args['M'], args['BLOCK_SIZE_M'])
                          * triton.cdiv(args['N'], args['BLOCK_SIZE_N']),
})
@triton.jit
def gemm_kernel(..., EVEN_K: tl.constexpr, GRID_MN: tl.constexpr):
    # Use EVEN_K to skip boundary checks in hot loop
    if EVEN_K:
        a = tl.load(a_ptrs)  # No mask needed
    else:
        a = tl.load(a_ptrs, mask=mask_k)