@triton.jit
def fused_attention_kernel(q_ptr, k_ptr, v_ptr, tmp_ptr, l_ptr, m_ptr,
    out_ptr, stride_qz: tl.constexpr, stride_qh: tl.constexpr, stride_qm:
    tl.constexpr, stride_qk: tl.constexpr, stride_kz: tl.constexpr,
    stride_kh: tl.constexpr, stride_kk: tl.constexpr, stride_kn: tl.
    constexpr, stride_vz: tl.constexpr, stride_vh: tl.constexpr, stride_vk:
    tl.constexpr, stride_vn: tl.constexpr, stride_oz: tl.constexpr,
    stride_oh: tl.constexpr, stride_om: tl.constexpr, stride_on: tl.
    constexpr, z: tl.constexpr, h: tl.constexpr, n_ctx: tl.constexpr,
    block_m: tl.constexpr, block_dmodel: tl.constexpr, block_n: tl.constexpr):
    """Flash attention kernel."""
    start_qm = tl.program_id(0)
    off_hz = tl.program_id(1)
    offs_m = start_qm * block_m + tl.arange(0, block_m)
    offs_n = tl.arange(0, block_n)
    offs_d = tl.arange(0, block_dmodel)
    off_q = off_hz * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :
        ] * stride_qk
    off_k = off_hz * stride_qh + offs_n[None, :] * stride_kn + offs_d[:, None
        ] * stride_kk
    off_v = off_hz * stride_qh + offs_n[:, None] * stride_qm + offs_d[None, :
        ] * stride_qk
    q_ptrs = q_ptr + off_q
    k_ptrs = k_ptr + off_k
    v_ptrs = v_ptr + off_v
    t_ptrs = tmp_ptr + off_hz * n_ctx + offs_m
    acc = tl.zeros([block_m, block_dmodel], dtype=tl.float32)
    m_i = tl.zeros([block_m], dtype=tl.float32) - float('inf')
    l_i = tl.zeros([block_m], dtype=tl.float32)
    q = tl.load(q_ptrs)
    for start_n in range(0, start_qm + 1):
        k = tl.load(k_ptrs)
        qk = tl.dot(q, k)
        qk += tl.where(offs_m[:, None] >= start_n * block_n + offs_n[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_scale = beta / l_i_new
        p = p * p_scale[:, None]
        p = p.to(tl.float16)
        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]
        v = tl.load(v_ptrs)
        acc += tl.dot(p, v)
        k_ptrs += block_n * stride_kn
        v_ptrs += block_n * stride_vk
        l_i = l_i_new
        m_i = m_i_new
    start_qm = tl.program_id(0)
    offs_m = start_qm * block_m + tl.arange(0, block_m)
    l_ptrs = l_ptr + off_hz * n_ctx + offs_m
    m_ptrs = m_ptr + off_hz * n_ctx + offs_m
    tl.store(l_ptrs, l_i)
    tl.store(m_ptrs, m_i)
    offs_n = tl.arange(0, block_dmodel)
    off_out = off_hz * stride_oh + offs_m[:, None] * stride_om + offs_n[None, :
        ] * stride_on
    out_ptrs = out_ptr + off_out
    tl.store(out_ptrs, acc)
