@triton.autotune(configs=get_configs(), key=['token_size', 'head_size'],
    reset_to_zero=['DK_ptr', 'DV_ptr'])
@triton.jit
def _backward(DO_ptr, stride_doh: tl.constexpr, stride_dom, stride_dod: tl.
    constexpr, DR_ptr, stride_drh, stride_drm, A_ptr, stride_ah, stride_am,
    Q_ptr, stride_qh: tl.constexpr, stride_qm, stride_qd: tl.constexpr,
    K_ptr, stride_kh: tl.constexpr, stride_kn, stride_kd: tl.constexpr,
    V_ptr, stride_vh: tl.constexpr, stride_vn, stride_vd: tl.constexpr,
    DQ_ptr, stride_dqh: tl.constexpr, stride_dqm, stride_dqd: tl.constexpr,
    DK_ptr, stride_dkh: tl.constexpr, stride_dkn, stride_dkd: tl.constexpr,
    DV_ptr, stride_dvh: tl.constexpr, stride_dvn, stride_dvd: tl.constexpr,
    KV_Lock_ptr, KV_Count_ptr, stride_kvs, stride_kvh, CSL_ptr, logit_scale,
    batch_size, token_size, head_size: tl.constexpr, num_heads: tl.
    constexpr, BLOCK_D: tl.constexpr, BLOCK_CSL: tl.constexpr, NO_D_MASK:
    tl.constexpr, NO_M_MASK: tl.constexpr, NO_N_MASK: tl.constexpr,
    ALLOW_TF32: tl.constexpr, inv_log2: tl.constexpr, BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr, acc_dtype: tl.constexpr=tl.float32,
    attend_current: tl.constexpr=False):
    tl.static_assert(BLOCK_M % BLOCK_N == 0)
    seq_id = tl.program_id(0)
    fhead_id = tl.program_id(1)
    seq_alloc_prog_id = tl.program_id(2)
    num_seq_alloc_progs = tl.num_programs(2)
    if seq_id == 0:
        seq_start_offset = 0
    else:
        seq_start_offset = tl.load(CSL_ptr + seq_id - 1).to(tl.int32)
    seq_end_offset = tl.load(CSL_ptr + seq_id).to(tl.int32)
    seq_length = seq_end_offset - seq_start_offset
    num_seq_blocks = tl.cdiv(seq_length, BLOCK_M)
    seq_a_block_id = num_seq_blocks - seq_alloc_prog_id - 1
    seq_b_block_id = seq_alloc_prog_id - (num_seq_alloc_progs - num_seq_blocks)
    if seq_a_block_id >= 0 or seq_b_block_id >= 0:
        qk_scale = inv_log2 * logit_scale
        M_range = tl.arange(0, BLOCK_M)
        N_range = tl.arange(0, BLOCK_N)
        D_range = tl.arange(0, BLOCK_D)
        D_mask = D_range < head_size
        cm = tl.where(N_range[:, None] >= N_range[None, :], 1.0, 0.0).to(Q_ptr
            .type.element_ty)
        if seq_a_block_id >= 0:
            head_id = fhead_id * 2
            DO_head_seq_ptr = (DO_ptr + stride_doh * head_id + stride_dom *
                seq_start_offset)
            DR_head_seq_ptr = (DR_ptr + stride_drh * head_id + stride_drm *
                seq_start_offset)
            A_head_seq_ptr = (A_ptr + stride_ah * head_id + stride_am *
                seq_start_offset)
            Q_head_seq_ptr = (Q_ptr + stride_qh * head_id + stride_qm *
                seq_start_offset)
            K_head_seq_ptr = (K_ptr + stride_kh * head_id + stride_kn *
                seq_start_offset)
            V_head_seq_ptr = (V_ptr + stride_vh * head_id + stride_vn *
                seq_start_offset)
            DQ_head_seq_ptr = (DQ_ptr + stride_dqh * head_id + stride_dqm *
                seq_start_offset)
            DK_head_seq_ptr = (DK_ptr + stride_dkh * head_id + stride_dkn *
                seq_start_offset)
            DV_head_seq_ptr = (DV_ptr + stride_dvh * head_id + stride_dvn *
                seq_start_offset)
            KV_Lock_head_seq_ptr = (KV_Lock_ptr + stride_kvs * seq_id + 
                stride_kvh * head_id)
            KV_Count_head_seq_ptr = (KV_Count_ptr + stride_kvs * seq_id + 
                stride_kvh * head_id)
            _backward_one_row(seq_a_block_id, seq_length, qk_scale, M_range,
                N_range, D_range, D_mask, cm, DO_head_seq_ptr, stride_dom,
                stride_dod, DR_head_seq_ptr, stride_drm, A_head_seq_ptr,
                stride_am, Q_head_seq_ptr, stride_qm, stride_qd,
                K_head_seq_ptr, stride_kn, stride_kd, V_head_seq_ptr,
                stride_vn, stride_vd, DQ_head_seq_ptr, stride_dqm,
                stride_dqd, DK_head_seq_ptr, stride_dkn, stride_dkd,
                DV_head_seq_ptr, stride_dvn, stride_dvd,
                KV_Lock_head_seq_ptr, KV_Count_head_seq_ptr, logit_scale,
                BLOCK_D, NO_D_MASK, NO_M_MASK, ALLOW_TF32, BLOCK_M, BLOCK_N,
                acc_dtype, attend_current=attend_current)
        if seq_b_block_id >= 0 and fhead_id * 2 + 1 < num_heads:
            head_id = fhead_id * 2 + 1
            DO_head_seq_ptr = (DO_ptr + stride_doh * head_id + stride_dom *
                seq_start_offset)
            DR_head_seq_ptr = (DR_ptr + stride_drh * head_id + stride_drm *
                seq_start_offset)
            A_head_seq_ptr = (A_ptr + stride_ah * head_id + stride_am *
                seq_start_offset)
            Q_head_seq_ptr = (Q_ptr + stride_qh * head_id + stride_qm *
                seq_start_offset)
            K_head_seq_ptr = (K_ptr + stride_kh * head_id + stride_kn *
                seq_start_offset)
            V_head_seq_ptr = (V_ptr + stride_vh * head_id + stride_vn *
                seq_start_offset)
            DQ_head_seq_ptr = (DQ_ptr + stride_dqh * head_id + stride_dqm *
                seq_start_offset)
            DK_head_seq_ptr = (DK_ptr + stride_dkh * head_id + stride_dkn *
                seq_start_offset)
            DV_head_seq_ptr = (DV_ptr + stride_dvh * head_id + stride_dvn *
                seq_start_offset)
            KV_Lock_head_seq_ptr = (KV_Lock_ptr + stride_kvs * seq_id + 
                stride_kvh * head_id)
            KV_Count_head_seq_ptr = (KV_Count_ptr + stride_kvs * seq_id + 
                stride_kvh * head_id)
            _backward_one_row(seq_b_block_id, seq_length, qk_scale, M_range,
                N_range, D_range, D_mask, cm, DO_head_seq_ptr, stride_dom,
                stride_dod, DR_head_seq_ptr, stride_drm, A_head_seq_ptr,
                stride_am, Q_head_seq_ptr, stride_qm, stride_qd,
                K_head_seq_ptr, stride_kn, stride_kd, V_head_seq_ptr,
                stride_vn, stride_vd, DQ_head_seq_ptr, stride_dqm,
                stride_dqd, DK_head_seq_ptr, stride_dkn, stride_dkd,
                DV_head_seq_ptr, stride_dvn, stride_dvd,
                KV_Lock_head_seq_ptr, KV_Count_head_seq_ptr, logit_scale,
                BLOCK_D, NO_D_MASK, NO_M_MASK, ALLOW_TF32, BLOCK_M, BLOCK_N,
                acc_dtype, attend_current=attend_current)
