@triton.jit
def _layer_norm_bwd_dx_fused(DX, DY, DSc, DSh, Y, Sc, Sh, Mean, Rstd, Lock,
    stride, N, GROUP_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK_SIZE_N)
    mask = cols < N
    Y += row * stride
    DY += row * stride
    DX += row * stride
    lock_id = row % GROUP_SIZE_M
    Lock += lock_id
    Count = Lock + GROUP_SIZE_M
    DSc = DSc + lock_id * N + cols
    DSh = DSh + lock_id * N + cols
    y = tl.load(Y + cols, mask=mask, other=0).to(tl.float32)
    dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32)
    sc = tl.load(Sc + cols, mask=mask).to(tl.float32)
    sh = tl.load(Sh + cols, mask=mask).to(tl.float32)
    mean = tl.load(Mean + row)
    rstd = tl.load(Rstd + row)
    xhat = (y - sh) / sc
    scdy = sc * dy
    xhat = tl.where(mask, xhat, 0.0)
    scdy = tl.where(mask, scdy, 0.0)
    c1 = tl.sum(xhat * scdy, axis=0) / N
    c2 = tl.sum(scdy, axis=0) / N
    dx = (scdy - (xhat * c1 + c2)) * rstd
    tl.store(DX + cols, dx, mask=mask)
    partial_dsc = (dy * xhat).to(sc.dtype)
    partial_dsh = dy.to(sc.dtype)
    while tl.atomic_cas(Lock, 0, 1) == 1:
        pass
    count = tl.load(Count)
    if count == 0:
        tl.atomic_xchg(Count, 1)
    else:
        partial_dsc += tl.load(DSc, mask=mask)
        partial_dsh += tl.load(DSh, mask=mask)
    tl.store(DSc, partial_dsc, mask=mask)
    tl.store(DSh, partial_dsh, mask=mask)
    tl.atomic_xchg(Lock, 0)
