@triton.jit
def _fwd_kernel(Q, K, V, S, Out, sm_scale, TMP, L, M, 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_oz, stride_oh,
    stride_om, stride_ok, stride_tz, stride_th, stride_tm, stride_lz,
    stride_lh, stride_lm, stride_mz, stride_mh, stride_mm, M_Q, N_CTX,
    BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.
    constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, CAUSAL: tl.
    constexpr, USE_ALIBI: tl.constexpr):
    start_m = tl.program_id(0) * BLOCK_M
    off_h = tl.program_id(1)
    off_z = tl.program_id(2)
    offs_m = start_m + tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)
    offs_d = tl.arange(0, BLOCK_DMODEL)
    off_q = off_z * stride_qz + off_h * stride_qh + offs_m[:, None
        ] * stride_qm + offs_d[None, :] * stride_qk
    off_k = off_z * stride_kz + off_h * stride_kh + offs_n[:, None
        ] * stride_kn + offs_d[None, :] * stride_kk
    off_v = off_z * stride_vz + off_h * stride_vh + offs_n[:, None
        ] * stride_vn + offs_d[None, :] * stride_vk
    q_ptrs = Q + off_q
    k_ptrs = K + off_k
    v_ptrs = V + off_v
    t_ptrs = TMP + off_z * stride_tz + off_h * stride_th + offs_m * stride_tm
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
    if EVEN_M:
        q = tl.load(q_ptrs)
    else:
        q = tl.load(q_ptrs, mask=offs_m[:, None] < M_Q, other=0)
    q = q.to(tl.float16)
    if USE_ALIBI:
        slope = tl.load(S + off_h)
    if CAUSAL & EVEN_M & EVEN_N:
        bound = start_m + BLOCK_M
    else:
        bound = N_CTX
    for start_n in range(0, bound, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        if EVEN_N:
            k = tl.load(k_ptrs)
        else:
            k = tl.load(k_ptrs, mask=start_n + offs_n[:, None] < N_CTX, other=0
                )
        k = k.to(tl.float16)
        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        qk += tl.dot(q, k, trans_b=True)
        qk *= sm_scale
        if USE_ALIBI & CAUSAL:
            qk += causal_alibi_mask(slope, offs_m, start_n + offs_n, M_Q,
                N_CTX, EVEN_M, EVEN_N)
        elif USE_ALIBI:
            qk += symmetric_alibi_mask(slope, offs_m, start_n + offs_n, M_Q,
                N_CTX, EVEN_M, EVEN_N)
        elif CAUSAL:
            qk += causal_mask(offs_m, start_n + offs_n, M_Q, N_CTX, EVEN_M,
                EVEN_N)
        else:
            qk += bounds_mask(offs_m, start_n + offs_n, M_Q, N_CTX, EVEN_M,
                EVEN_N)
        m_ij = tl.maximum(tl.max(qk, axis=1), -10000)
        p = tl.exp(qk - m_ij[:, None])
        l_ij = tl.sum(p, axis=1)
        m_i_new = tl.maximum(m_i, m_ij)
        alpha = tl.exp(m_i - m_i_new)
        beta = tl.exp(m_ij - m_i_new)
        l_i_new = alpha * l_i + beta * l_ij
        p_scale = beta / l_i_new
        p = p * p_scale[:, None]
        acc_scale = l_i / l_i_new * alpha
        tl.store(t_ptrs, acc_scale)
        acc_scale = tl.load(t_ptrs)
        acc = acc * acc_scale[:, None]
        if EVEN_N:
            v = tl.load(v_ptrs)
        else:
            v = tl.load(v_ptrs, mask=start_n + offs_n[:, None] < N_CTX, other=0
                )
        v = v.to(tl.float16)
        p = p.to(tl.float16)
        acc += tl.dot(p, v)
        l_i = l_i_new
        m_i = m_i_new
        v_ptrs += BLOCK_N * stride_vn
        k_ptrs += BLOCK_N * stride_kn
    offs_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
    l_ptrs = L + off_z * stride_lz + off_h * stride_lh + offs_m * stride_lm
    m_ptrs = M + off_z * stride_mz + off_h * stride_mh + offs_m * stride_mm
    if EVEN_M:
        tl.store(l_ptrs, l_i)
        tl.store(m_ptrs, m_i)
    else:
        tl.store(l_ptrs, l_i, mask=offs_m < M_Q)
        tl.store(m_ptrs, m_i, mask=offs_m < M_Q)
    offs_d = tl.arange(0, BLOCK_DMODEL)
    off_o = off_z * stride_oz + off_h * stride_oh + offs_m[:, None
        ] * stride_om + offs_d[None, :] * stride_ok
    out_ptrs = Out + off_o
    if EVEN_M:
        tl.store(out_ptrs, acc)
    else:
        tl.store(out_ptrs, acc, mask=offs_m[:, None] < M_Q)
