@triton.jit
def _rms_norm_2d_bwd_dx_fused(DX, DY, DW, DB, X, W, B, Rrms, M, C, N,
    num_blocks, eps, GROUP_SIZE_M: tl.constexpr, BLOCK_SIZE: tl.constexpr,
    BLOCK_SIZE_C: tl.constexpr):
    m_n = tl.program_id(0)
    m, n = m_n // num_blocks, m_n % num_blocks
    X += m * C * N
    DY += m * C * N
    DX += m * C * N
    Rrms += m * N
    cols = n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = cols < N
    DW = DW + m_n * C
    DB = DB + m_n * C
    rrms = tl.load(Rrms + cols, mask=mask, other=1)
    c1 = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
    for off in range(0, C):
        pos = off * N + cols
        x = tl.load(X + pos, mask=mask, other=0).to(tl.float32)
        dy = tl.load(DY + pos, mask=mask, other=0).to(tl.float32)
        w = tl.load(W + off).to(tl.float32)
        xhat = x * rrms
        wdy = w * dy
        xhat = tl.where(mask, xhat, 0.0)
        wdy = tl.where(mask, wdy, 0.0)
        c1 += xhat * wdy
        tl.store(DW + off, tl.sum((dy * xhat).to(w.dtype), axis=0))
        tl.store(DB + off, tl.sum(dy.to(w.dtype), axis=0))
    c1 /= C
    for off in range(0, C):
        pos = off * N + cols
        x = tl.load(X + pos, mask=mask, other=0).to(tl.float32)
        dy = tl.load(DY + pos, mask=mask, other=0).to(tl.float32)
        w = tl.load(W + off).to(tl.float32)
        xhat = x * rrms
        wdy = w * dy
        dx = (wdy - xhat * c1) * rrms
        tl.store(DX + pos, dx, mask=mask)
