@triton.autotune(configs=get_configs(), key=['head_size'])
@triton.jit
def _forward(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, O_ptr, stride_oh: tl.constexpr, stride_om, stride_od: tl.
    constexpr, R_ptr, stride_rh, stride_rm: tl.constexpr, A_ptr, stride_ah,
    stride_am: tl.constexpr, W_ptr, stride_wh, stride_wm, stride_wn,
    CSL_ptr, logit_scale: tl.constexpr, batch_size, token_size, head_size:
    tl.constexpr, num_heads: tl.constexpr, BLOCK_D: 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, no_grad: tl.constexpr=False, acc_dtype: tl.
    constexpr=tl.float32, return_attention: tl.constexpr=False, use_cumsum:
    tl.constexpr=False, 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
        if not use_cumsum:
            cm = tl.where(N_range[:, None] >= N_range[None, :], 1.0, 0.0).to(
                Q_ptr.type.element_ty)
        else:
            cm = None
        if seq_a_block_id >= 0:
            head_id = fhead_id * 2
            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)
            O_head_seq_ptr = (O_ptr + stride_oh * head_id + stride_om *
                seq_start_offset)
            R_head_seq_ptr = (R_ptr + stride_rh * head_id + stride_rm *
                seq_start_offset)
            A_head_seq_ptr = (A_ptr + stride_ah * head_id + stride_am *
                seq_start_offset)
            W_head_seq_ptr = (W_ptr + stride_wh * head_id + stride_am *
                seq_start_offset)
            _forward_one_row(seq_a_block_id, seq_length, qk_scale, M_range,
                N_range, D_range, D_mask, cm, Q_head_seq_ptr, stride_qm,
                stride_qd, K_head_seq_ptr, stride_kn, stride_kd,
                V_head_seq_ptr, stride_vn, stride_vd, O_head_seq_ptr,
                stride_om, stride_od, R_head_seq_ptr, stride_rm,
                A_head_seq_ptr, stride_am, W_head_seq_ptr, stride_wm,
                stride_wn, BLOCK_D, NO_D_MASK, NO_M_MASK, NO_N_MASK,
                ALLOW_TF32, BLOCK_M, BLOCK_N, no_grad, acc_dtype,
                return_attention, use_cumsum=use_cumsum, attend_current=
                attend_current)
        if seq_b_block_id >= 0 and fhead_id * 2 + 1 < num_heads:
            head_id = fhead_id * 2 + 1
            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)
            O_head_seq_ptr = (O_ptr + stride_oh * head_id + stride_om *
                seq_start_offset)
            R_head_seq_ptr = (R_ptr + stride_rh * head_id + stride_rm *
                seq_start_offset)
            A_head_seq_ptr = (A_ptr + stride_ah * head_id + stride_am *
                seq_start_offset)
            W_head_seq_ptr = (W_ptr + stride_wh * head_id + stride_am *
                seq_start_offset)
            _forward_one_row(seq_b_block_id, seq_length, qk_scale, M_range,
                N_range, D_range, D_mask, cm, Q_head_seq_ptr, stride_qm,
                stride_qd, K_head_seq_ptr, stride_kn, stride_kd,
                V_head_seq_ptr, stride_vn, stride_vd, O_head_seq_ptr,
                stride_om, stride_od, R_head_seq_ptr, stride_rm,
                A_head_seq_ptr, stride_am, W_head_seq_ptr, stride_wm,
                stride_wn, BLOCK_D, NO_D_MASK, NO_M_MASK, NO_N_MASK,
                ALLOW_TF32, BLOCK_M, BLOCK_N, no_grad, acc_dtype,
                return_attention, use_cumsum=use_cumsum, attend_current=
                attend_current)
