@triton.jit
def _bwd_preprocess(Out, DO, NDO, L, Delta, M_Q, stride_oz, stride_oh,
    stride_om, stride_od, stride_doz, stride_doh, stride_dom, stride_dod,
    stride_ndoz, stride_ndoh, stride_ndom, stride_ndod, stride_lz,
    stride_lh, stride_lm, stride_dz, stride_dh, stride_dm, BLOCK_DMODEL: tl
    .constexpr, BLOCK_M: tl.constexpr, EVEN_M: tl.constexpr):
    off_h = tl.program_id(1)
    off_z = tl.program_id(2)
    offs_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, BLOCK_DMODEL)
    Out = Out + off_z * stride_oz + off_h * stride_oh + offs_m[:, None
        ] * stride_om + offs_d[None, :] * stride_od
    DO = DO + off_z * stride_doz + off_h * stride_doh + offs_m[:, None
        ] * stride_dom + offs_d[None, :] * stride_dod
    NDO = NDO + off_z * stride_ndoz + off_h * stride_ndoh + offs_m[:, None
        ] * stride_ndom + offs_d[None, :] * stride_ndod
    L = L + off_z * stride_lz + off_h * stride_lh + offs_m * stride_lm
    Delta = Delta + off_z * stride_dz + off_h * stride_dh + offs_m * stride_dm
    if EVEN_M:
        o = tl.load(Out).to(tl.float32)
        do = tl.load(DO).to(tl.float32)
        denom = tl.load(L).to(tl.float32)
    else:
        o = tl.load(Out, mask=offs_m[:, None] < M_Q).to(tl.float32)
        do = tl.load(DO, mask=offs_m[:, None] < M_Q).to(tl.float32)
        denom = tl.load(L, mask=offs_m < M_Q).to(tl.float32)
    do = do / denom[:, None]
    delta = tl.sum(o * do, axis=1)
    if EVEN_M:
        tl.store(NDO, do)
        tl.store(Delta, delta)
    else:
        tl.store(NDO, do, mask=offs_m[:, None] < M_Q)
        tl.store(Delta, delta, mask=offs_m < M_Q)
