@triton.jit
def _rms_norm_2d_fwd_fused(X, Y, W, B, Rrms, M, C, N, num_blocks, eps,
    BLOCK_SIZE: tl.constexpr):
    m_n = tl.program_id(0)
    m, n = m_n // num_blocks, m_n % num_blocks
    Y += m * C * N
    X += m * C * N
    cols = n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = cols < N
    x_sum_square = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
    for off in range(0, C):
        x = tl.load(X + off * N + cols, mask=mask, other=0.0).to(tl.float32)
        x_sum_square += x * x
    mean_square = x_sum_square / C
    rrms = 1 / tl.sqrt(mean_square + eps)
    tl.store(Rrms + m * N + cols, rrms, mask=mask)
    for off in range(0, C):
        pos = off * N + cols
        w = tl.load(W + off)
        b = tl.load(B + off)
        x = tl.load(X + pos, mask=mask, other=0.0).to(tl.float32)
        x_hat = x * rrms
        y = x_hat * w + b
        tl.store(Y + pos, y, mask=mask)
