@triton.autotune(configs=get_configs(), key=['token_size', 'head_size'])
@triton.jit
def _forward(Q_ptr, stride_qb, stride_qh, stride_qm: tl.constexpr,
    stride_qd: tl.constexpr, K_ptr, stride_kb, stride_kh, stride_kn: tl.
    constexpr, stride_kd: tl.constexpr, V_ptr, stride_vb, stride_vh,
    stride_vn: tl.constexpr, stride_vd: tl.constexpr, O_ptr, stride_ob,
    stride_oh, stride_om: tl.constexpr, stride_od: tl.constexpr, R_ptr,
    stride_rb, stride_rh, stride_rm: tl.constexpr, A_ptr, stride_ab,
    stride_ah, stride_am: tl.constexpr, W_ptr, stride_wb, stride_wh,
    stride_wm, stride_wn, logit_scale: tl.constexpr, attend_current: 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, is_compiling: tl.constexpr=False):
    tl.static_assert(BLOCK_M % BLOCK_N == 0)
    batch_id = tl.program_id(0)
    head_pid = tl.program_id(1)
    prog_id = tl.program_id(2)
    tl.num_programs(2)
    seq_length = token_size
    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)
    head_id = head_pid
    seq_prog_id = prog_id
    Q_head_seq_ptr = Q_ptr + stride_qb * batch_id + stride_qh * head_id
    K_head_seq_ptr = K_ptr + stride_kb * batch_id + stride_kh * head_id
    V_head_seq_ptr = V_ptr + stride_vb * batch_id + stride_vh * head_id
    O_head_seq_ptr = O_ptr + stride_ob * batch_id + stride_oh * head_id
    R_head_seq_ptr = R_ptr + stride_rb * batch_id + stride_rh * head_id
    A_head_seq_ptr = A_ptr + stride_ab * batch_id + stride_ah * head_id
    W_head_seq_ptr = W_ptr + stride_wb * batch_id + stride_wh * head_id
    _forward_one_row(seq_prog_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,
        attend_current=attend_current, is_compiling=is_compiling)
