@triton.jit
def _bwd_kernel(Q, K, V, S, sm_scale, DO, DQ, DK, DV, M, D, stride_qz,
    stride_qh, stride_qm, stride_qk, stride_kz, stride_kh, stride_kn,
    stride_kk, stride_vz, stride_vh, stride_vn, stride_vk, stride_doz,
    stride_doh, stride_dom, stride_dok, stride_dqz, stride_dqh, stride_dqm,
    stride_dqk, stride_dkz, stride_dkh, stride_dkn, stride_dkk, stride_dvz,
    stride_dvh, stride_dvn, stride_dvk, stride_mz, stride_mh, stride_mm,
    stride_dz, stride_dh, stride_dm, M_Q, N_CTX, BLOCK_DMODEL: tl.constexpr,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, EVEN_M: tl.constexpr,
    EVEN_N: tl.constexpr, CAUSAL: tl.constexpr, USE_ALIBI: tl.constexpr):
    off_h = tl.program_id(0)
    off_z = tl.program_id(1)
    Q += off_z * stride_qz + off_h * stride_qh
    K += off_z * stride_kz + off_h * stride_kh
    V += off_z * stride_vz + off_h * stride_vh
    DO += off_z * stride_doz + off_h * stride_doh
    DQ += off_z * stride_dqz + off_h * stride_dqh
    DK += off_z * stride_dkz + off_h * stride_dkh
    DV += off_z * stride_dvz + off_h * stride_dvh
    if USE_ALIBI:
        slope = tl.load(S + off_h)
    for start_n in range(0, N_CTX, BLOCK_N):
        offs_n_curr = start_n + tl.arange(0, BLOCK_N)
        offs_d = tl.arange(0, BLOCK_DMODEL)
        offs_m = tl.arange(0, BLOCK_M)
        k_ptrs = K + (offs_n_curr[:, None] * stride_kn + offs_d[None, :] *
            stride_kk)
        v_ptrs = V + (offs_n_curr[:, None] * stride_vn + offs_d[None, :] *
            stride_vk)
        q_ptrs = Q + (offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk
            )
        dq_ptrs = DQ + (offs_m[:, None] * stride_dqm + offs_d[None, :] *
            stride_dqk)
        do_ptrs = DO + (offs_m[:, None] * stride_dom + offs_d[None, :] *
            stride_dok)
        m_ptrs = M + off_z * stride_mz + off_h * stride_mh + offs_m * stride_mm
        D_ptrs = D + off_z * stride_dz + off_h * stride_dh + offs_m * stride_dm
        dv = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
        dk = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
        if EVEN_N:
            k = tl.load(k_ptrs)
            v = tl.load(v_ptrs)
        else:
            k = tl.load(k_ptrs, mask=offs_n_curr[:, None] < N_CTX, other=0)
            v = tl.load(v_ptrs, mask=offs_n_curr[:, None] < N_CTX, other=0)
        k = k.to(tl.float16)
        v = v.to(tl.float16)
        if CAUSAL:
            begin = start_n + M_Q - N_CTX
            dq_ptrs += begin * stride_dqm
            q_ptrs += begin * stride_qm
            do_ptrs += begin * stride_dom
            m_ptrs += begin * stride_mm
            D_ptrs += begin * stride_dm
        else:
            begin = 0
        for start_m in range(begin, M_Q, BLOCK_M):
            offs_m_curr = start_m + offs_m
            if EVEN_M:
                q = tl.load(q_ptrs)
            else:
                q = tl.load(q_ptrs, mask=offs_m_curr[:, None] < M_Q, other=0)
            q = q.to(tl.float16)
            qk = tl.dot(q, k, trans_b=True)
            qk *= sm_scale
            if USE_ALIBI & CAUSAL:
                qk += causal_alibi_mask(slope, offs_m_curr, offs_n_curr,
                    M_Q, N_CTX, EVEN_M, EVEN_N)
            elif USE_ALIBI:
                qk += symmetric_alibi_mask(slope, offs_m_curr, offs_n_curr,
                    M_Q, N_CTX, EVEN_M, EVEN_N)
            elif CAUSAL:
                qk += causal_mask(offs_m_curr, offs_n_curr, M_Q, N_CTX,
                    EVEN_M, EVEN_N)
            if EVEN_M:
                m = tl.load(m_ptrs)
                Di = tl.load(D_ptrs)
                do = tl.load(do_ptrs)
            else:
                m = tl.load(m_ptrs, mask=offs_m_curr < M_Q, other=0)
                Di = tl.load(D_ptrs, mask=offs_m_curr < M_Q, other=0)
                do = tl.load(do_ptrs, mask=offs_m_curr[:, None] < M_Q, other=0)
            do = do.to(tl.float16)
            p = tl.exp(qk - m[:, None])
            dv += tl.dot(p.to(tl.float16), do, trans_a=True)
            dp = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - Di[:, None]
            dp += tl.dot(do, v, trans_b=True)
            ds = p * dp * sm_scale
            dk += tl.dot(ds.to(tl.float16), q, trans_a=True)
            if EVEN_M:
                dq = tl.load(dq_ptrs, eviction_policy='evict_last')
            else:
                dq = tl.load(dq_ptrs, mask=offs_m_curr[:, None] < M_Q,
                    other=0, eviction_policy='evict_last')
            dq += tl.dot(ds.to(tl.float16), k)
            if EVEN_M:
                tl.store(dq_ptrs, dq, eviction_policy='evict_last')
            else:
                tl.store(dq_ptrs, dq, mask=offs_m_curr[:, None] < M_Q,
                    eviction_policy='evict_last')
            q_ptrs += BLOCK_M * stride_qm
            dq_ptrs += BLOCK_M * stride_dqm
            do_ptrs += BLOCK_M * stride_dom
            m_ptrs += BLOCK_M * stride_mm
            D_ptrs += BLOCK_M * stride_dm
        offs_d = tl.arange(0, BLOCK_DMODEL)
        dv_ptrs = DV + (offs_n_curr[:, None] * stride_dvn + offs_d[None, :] *
            stride_dvk)
        dk_ptrs = DK + (offs_n_curr[:, None] * stride_dkn + offs_d[None, :] *
            stride_dkk)
        if EVEN_N:
            tl.store(dv_ptrs, dv)
            tl.store(dk_ptrs, dk)
        else:
            tl.store(dv_ptrs, dv, mask=offs_n_curr[:, None] < N_CTX)
            tl.store(dk_ptrs, dk, mask=offs_n_curr[:, None] < N_CTX)
