🚀 The feature, motivation and pitch

1. Backgroud: For normalization kernels, e.g. RMS Norm, the inputs are loaded twice, one for reduction, the other for normalization. Inductor generates a persistent reduction which loads inputs to registers for once and cache as float32 for computations. Due to register pressure, this feature is only used when the reduction dimension is small, e.g. <= 1024. Recent blog showed the current SOTA is about 6000GB/s while the max on GB200 is about 7928 GB/s.

2. Proposal: Optimize persistent reduction with recompute from cached inputs 2. 1 Register cache (a) Instead of cache inputs as float32, we directly cache them as float16 and only convert to float32 at the time of computation. (b) Issue multiple load instructions at the same time to ensure enough bytes-in-flights

The final generate kernel (using 2 unrolled loads as an example) should look like:

@triton.jit
def rmsnorm_2tile_kernel(x_ptr, w_ptr, y_ptr, M, stride_x, stride_y,
                         EPS: tl.constexpr, TILE: tl.constexpr, BLOCK_M: tl.constexpr):
    pid = tl.program_id(0)
    rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]  # [BLOCK_M, 1]
    mask = rows < M
    offs = tl.arange(0, TILE)[None, :]  # [1, TILE]
    base = x_ptr + rows * stride_x
    x0 = tl.load(base + offs, mask=mask, other=0.0)
    x1 = tl.load(base + TILE + offs, mask=mask, other=0.0)
    ss = x0.to(tl.float32) * x0.to(tl.float32)
    ss += x1.to(tl.float32) * x1.to(tl.float32)
    inv_rms = tl.rsqrt(tl.sum(ss, axis=1)[:, None] / (2 * TILE) + EPS)
    ybase = y_ptr + rows * stride_y
    w0 = tl.load(w_ptr + offs)
    tl.store(ybase + offs, (x0.to(tl.float32) * inv_rms * w0).to(tl.bfloat16), mask=mask)
    w1 = tl.load(w_ptr + TILE + offs)
    tl.store(ybase + TILE + offs, (x1.to(tl.float32) * inv_rms * w1).to(tl.bfloat16), mask=mask)

2.2 Shared memory cache with TMA Register cache still suffers from register pressure when hidden size is larger than 8192, shared memory can be used, which can also benefit from cluster reduction where multiple SMs are used for one reduction row.

3. Preliminary Results: SOL% = achieved bandwidth / 7928 GB/s peak. Measured with torch.profiler (CUPTI) and L2 cache clearing. Baseline is: torch.compile(rms_norm_ref, mode="max-autotune-no-cudagraphs", fullgraph=True) 3.1 Register cache (best gains at common LLM hidden dims):

Reg-tiled dominates for N = 512–8192. Compiled wins at N >= 32768 (register spilling).

3.2 Shared memory cache with TMA (constant ~22 registers, no performance cliff at large N):

Hi @shunting314 and @eellison, what do you think about this proposal?