@triton.heuristics({'EVEN_M': lambda args: args['seqlen_q'] % args[
    'BLOCK_M'] == 0, 'EVEN_N': lambda args: args['seqlen_k'] % args[
    'BLOCK_N'] == 0})
@triton.jit
def _fwd_kernel(Q1, Q2, K1, K2, V, sm_scale, Out, stride_qz, stride_qh,
    stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk,
    stride_vz, stride_vh, stride_vk, stride_vn, stride_oz, stride_oh,
    stride_om, stride_on, seqlen_q, seqlen_k, BLOCK_M: tl.constexpr,
    BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr, IS_CAUSAL: tl.
    constexpr, WINDOW: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr
    ):
    start_m = tl.program_id(0)
    off_hz = tl.program_id(1)
    q_offset = off_hz * stride_qh
    kv_offset = off_hz * stride_kh
    Q1_block_ptr = tl.make_block_ptr(base=Q1 + q_offset, shape=(seqlen_q,
        BLOCK_DMODEL), strides=(stride_qm, stride_qk), offsets=(start_m *
        BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_DMODEL), order=(1, 0))
    Q2_block_ptr = tl.make_block_ptr(base=Q2 + q_offset, shape=(seqlen_q,
        BLOCK_DMODEL), strides=(stride_qm, stride_qk), offsets=(start_m *
        BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_DMODEL), order=(1, 0))
    K1_block_ptr = tl.make_block_ptr(base=K1 + kv_offset, shape=(
        BLOCK_DMODEL, seqlen_k), strides=(stride_kk, stride_kn), offsets=(0,
        0), block_shape=(BLOCK_DMODEL, BLOCK_N), order=(0, 1))
    K2_block_ptr = tl.make_block_ptr(base=K2 + kv_offset, shape=(
        BLOCK_DMODEL, seqlen_k), strides=(stride_kk, stride_kn), offsets=(0,
        0), block_shape=(BLOCK_DMODEL, BLOCK_N), order=(0, 1))
    V_block_ptr = tl.make_block_ptr(base=V + kv_offset, shape=(seqlen_k,
        BLOCK_DMODEL), strides=(stride_vk, stride_vn), offsets=(0, 0),
        block_shape=(BLOCK_N, BLOCK_DMODEL), order=(1, 0))
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)
    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)
    qk_scale = sm_scale * 1.44269504
    if EVEN_M:
        q1 = tl.load(Q1_block_ptr)
        q2 = tl.load(Q2_block_ptr)
    else:
        q1 = tl.load(Q1_block_ptr, boundary_check=(0,), padding_option='zero')
        q2 = tl.load(Q2_block_ptr, boundary_check=(0,), padding_option='zero')
    lo = 0
    hi = tl.minimum((start_m + 1) * BLOCK_M, seqlen_k
        ) if IS_CAUSAL else seqlen_k
    for start_n in range(lo, hi, BLOCK_N):
        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        if IS_CAUSAL:
            qk = tl.where(offs_m[:, None] >= start_n + offs_n[None, :], qk,
                float('-inf'))
        if not EVEN_N:
            qk = tl.where(start_n + offs_n[None, :] < seqlen_k, qk, float(
                '-inf'))
        if start_n <= start_m * BLOCK_M - WINDOW - BLOCK_N or start_n >= (
            start_m + 1) * BLOCK_M + WINDOW:
            if EVEN_N & EVEN_M:
                k2 = tl.load(K2_block_ptr)
            else:
                k2 = tl.load(K2_block_ptr, boundary_check=(1,),
                    padding_option='zero')
            qk += tl.dot(q2, k2)
        elif start_n > (start_m + 1
            ) * BLOCK_M - WINDOW and start_n < start_m * BLOCK_M + WINDOW - BLOCK_N:
            if EVEN_N & EVEN_M:
                k1 = tl.load(K1_block_ptr)
            else:
                k1 = tl.load(K1_block_ptr, boundary_check=(1,),
                    padding_option='zero')
            qk += tl.dot(q1, k1)
        else:
            if EVEN_N & EVEN_M:
                k1 = tl.load(K1_block_ptr)
                k2 = tl.load(K2_block_ptr)
            else:
                k1 = tl.load(K1_block_ptr, boundary_check=(1,),
                    padding_option='zero')
                k2 = tl.load(K2_block_ptr, boundary_check=(1,),
                    padding_option='zero')
            qk1 = tl.dot(q1, k1)
            qk2 = tl.dot(q2, k2)
            qk += tl.where(tl.abs(offs_m[:, None] - (start_n + offs_n[None,
                :])) < WINDOW, qk1, qk2)
        qk *= qk_scale
        m_i_new = tl.maximum(m_i, tl.max(qk, 1))
        alpha = tl.math.exp2(m_i - m_i_new)
        p = tl.math.exp2(qk - m_i_new[:, None])
        acc_scale = l_i * 0 + alpha
        acc *= acc_scale[:, None]
        if EVEN_N & EVEN_M:
            v = tl.load(V_block_ptr)
        else:
            v = tl.load(V_block_ptr, boundary_check=(0,), padding_option='zero'
                )
        v = tl.load(V_block_ptr)
        acc += tl.dot(p.to(v.dtype), v)
        l_i = l_i * alpha + tl.sum(p, 1)
        m_i = m_i_new
        K1_block_ptr = tl.advance(K1_block_ptr, (0, BLOCK_N))
        K2_block_ptr = tl.advance(K2_block_ptr, (0, BLOCK_N))
        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
    acc = acc / l_i[:, None]
    O_block_ptr = tl.make_block_ptr(base=Out + q_offset, shape=(seqlen_q,
        BLOCK_DMODEL), strides=(stride_om, stride_on), offsets=(start_m *
        BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_DMODEL), order=(1, 0))
    offs_d = tl.arange(0, BLOCK_DMODEL)
    out_ptrs = Out + q_offset + (offs_m[:, None] * stride_om + offs_d[None, :])
    if EVEN_M:
        tl.store(out_ptrs, acc)
    else:
        tl.store(out_ptrs, acc, mask=offs_m[:, None] < seqlen_q)
