@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, Out, L, softmax_scale, stride_qb,
    stride_qh, stride_qm, stride_kb, stride_kh, stride_kn, stride_vb,
    stride_vh, stride_vn, stride_ob, stride_oh, stride_om, nheads, seqlen_q,
    seqlen_k, CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K, WINDOW_SIZE: tl.
    constexpr, BLOCK_HEADDIM: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N:
    tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, WRITE_LSE:
    tl.constexpr):
    start_m = tl.program_id(0)
    off_hb = tl.program_id(1)
    off_b = off_hb // nheads
    off_h = off_hb % nheads
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)
    offs_d = tl.arange(0, BLOCK_HEADDIM)
    Q1_block_ptr = tl.make_block_ptr(base=Q1 + (off_b * stride_qb + off_h *
        stride_qh), shape=(seqlen_q, BLOCK_HEADDIM), strides=(stride_qm, 1),
        offsets=(start_m * BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_HEADDIM
        ), order=(1, 0))
    Q2_block_ptr = tl.make_block_ptr(base=Q2 + (off_b * stride_qb + off_h *
        stride_qh), shape=(seqlen_q, BLOCK_HEADDIM), strides=(stride_qm, 1),
        offsets=(start_m * BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_HEADDIM
        ), order=(1, 0))
    K1_block_ptr = tl.make_block_ptr(base=K1 + (off_b * stride_kb + off_h *
        stride_kh), shape=(BLOCK_HEADDIM, seqlen_k), strides=(1, stride_kn),
        offsets=(0, 0), block_shape=(BLOCK_HEADDIM, BLOCK_N), order=(0, 1))
    K2_block_ptr = tl.make_block_ptr(base=K2 + (off_b * stride_kb + off_h *
        stride_kh), shape=(BLOCK_HEADDIM, seqlen_k), strides=(1, stride_kn),
        offsets=(0, 0), block_shape=(BLOCK_HEADDIM, BLOCK_N), order=(0, 1))
    V_block_ptr = tl.make_block_ptr(base=V + (off_b * stride_vb + off_h *
        stride_vh), shape=(seqlen_k, BLOCK_HEADDIM), strides=(stride_vn, 1),
        offsets=(0, 0), block_shape=(BLOCK_N, BLOCK_HEADDIM), order=(1, 0))
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) + NEGINF
    acc = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)
    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')
    acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q1, q2, softmax_scale,
        K1_block_ptr, K2_block_ptr, V_block_ptr, start_m, offs_m, offs_n,
        seqlen_k, WINDOW_SIZE, BLOCK_M, BLOCK_N, EVEN_M & EVEN_N, 1)
    acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q1, q2, softmax_scale,
        K1_block_ptr, K2_block_ptr, V_block_ptr, start_m, offs_m, offs_n,
        seqlen_k, WINDOW_SIZE, BLOCK_M, BLOCK_N, EVEN_M & EVEN_N, 2)
    acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q1, q2, softmax_scale,
        K1_block_ptr, K2_block_ptr, V_block_ptr, start_m, offs_m, offs_n,
        seqlen_k, WINDOW_SIZE, BLOCK_M, BLOCK_N, EVEN_M & EVEN_N, 3)
    if WRITE_LSE:
        l_ptrs = L + off_hb * seqlen_q + offs_m
        tl.store(l_ptrs, m_i + tl.math.log2(l_i))
    acc = acc / l_i[:, None]
    start_m = tl.program_id(0)
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, BLOCK_HEADDIM)
    out_ptrs = Out + off_b * stride_ob + off_h * stride_oh + (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)
