| name | rmsnorm-kernel-optimization |
| description | Optimize RMS Normalization kernels in Triton for NVIDIA and AMD GPUs. Covers row-parallel reduction, vectorized loads, warp shuffle patterns, and rsqrt usage. Use when writing or optimizing RMSNorm, LayerNorm, or similar normalization kernels. |
RMSNorm Kernel Optimization
Overview
RMS Normalization: out = x / sqrt(mean(x^2) + eps) * weight
RMSNorm is memory-bound: each element is read once and written once, with minimal compute.
The key metric is GBps (fraction of HBM peak bandwidth).
| Bound | Metric | Target (4090) | Target (MI300X) |
|---|
| Memory | % peak BW | >60% | >60% |
AutoKernel reported ~34.7% of H100 peak bandwidth for a baseline RMSNorm.
Core Technique
Row-Parallel Design
One Triton program per row (token). Each program:
- Loads the entire row into registers (FP32 for numerical stability)
- Computes sum of squares via
tl.sum(x * x)
- Applies
rsqrt(mean_sq + eps) -- single fast instruction
- Multiplies by weight vector
- Stores result
Key Optimizations
- FP32 accumulation: always compute reduction in float32, even for FP16/BF16 inputs
- Vectorized loads: BLOCK_SIZE should be >= hidden_dim, next_power_of_2 for alignment
rsqrt intrinsic: 1/sqrt(x) is a single HW instruction on both NVIDIA and AMD- Multi-row per program: for small hidden dims, process multiple rows per program to improve occupancy
NVIDIA vs AMD Differences
- AMD wavefront=64: fewer programs needed for same thread count
- MI300X LDS = 64 KB/CU (vs 228 KB/SM on H100): less shared memory, but 256 MB L2
- Both: BLOCK_SIZE should be power of 2, >= hidden_dim
Verification
python skills/kernels/rmsnorm/test_rmsnorm.py
- Correctness: compare vs
x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * weight
- Performance: report GBps =
(read_bytes + write_bytes) / latency_s / 1e9
Common Pitfalls
- FP16 reduction: computing
sum(x^2) in FP16 overflows for large hidden dims
- Missing mask: when hidden_dim is not power of 2, unmasked loads cause errors
- Separate sqrt + divide: use
rsqrt instead of sqrt + divide (2x fewer instructions)
- Large BLOCK_SIZE on AMD: LDS is only 64 KB/CU; keep register pressure manageable
Agent Instructions
When optimizing a RMSNorm kernel:
- Profile to confirm memory-bound (GBps should be the focus metric)
- Ensure FP32 accumulation for the reduction
- Try multi-row per program for small hidden dims (< 1024)
- Check vectorized load efficiency with NCU/rocprof
- Run test script after every change
References
