@triton.jit
def _fwd_kernel(Q, K, V, sm_scale, L, M, Out, stride_head, stride_seq,
    hidden_dim, seq_len, Br: tl.constexpr, Bt: tl.constexpr, hidden: tl.
    constexpr):
    start_m = tl.program_id(0)
    off_hz = tl.program_id(1)
    offs_br = start_m * Br + tl.arange(0, Br)
    offs_bt = tl.arange(0, Bt)
    offs_d = tl.arange(0, hidden)
    off_q = off_hz * stride_head + offs_br[:, None] * stride_seq + offs_d[
        None, :] * hidden_dim
    off_kv = off_hz * stride_head + offs_bt[:, None] * stride_seq + offs_d[
        None, :] * hidden_dim
    q_ptrs = Q + off_q
    k_ptrs = K + off_kv
    v_ptrs = V + off_kv
    m_i = tl.zeros([Br], dtype=tl.float16) - float('inf')
    l_i = tl.zeros([Br], dtype=tl.float16)
    acc = tl.zeros([Br, hidden], dtype=tl.float16)
    q = tl.load(q_ptrs)
    for start_n in tl.range(0, (start_m + 1) * Br, Bt):
        k = tl.load(k_ptrs + start_n * stride_seq)
        v = tl.load(v_ptrs + start_n * stride_seq)
        qk = tl.dot(q, tl.trans(k)) * sm_scale
        qk += tl.where(offs_br[:, None] >= start_n + offs_bt[None, :], 0,
            float('-inf'))
        m_ij = tl.max(qk, 1)
        p = tl.exp(qk - m_ij[:, None])
        l_ij = tl.sum(p, 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 = p * (beta / l_i_new)[:, None]
        p = p.to(v.dtype)
        acc_scale_factor = l_i / l_i_new * alpha
        acc = acc * acc_scale_factor[:, None] + tl.dot(p, v)
        acc = acc.to(v.dtype)
        l_i = l_i_new.to(l_i.dtype)
        m_i = m_i_new.to(m_i.dtype)
    l_ptrs = L + off_hz * seq_len + offs_br
    m_ptrs = M + off_hz * seq_len + offs_br
    out_ptrs = Out + off_q
    tl.store(l_ptrs, l_i)
    tl.store(m_ptrs, m_i)
    tl.store(out_ptrs, acc)
