@triton.jit
def _attn_fwd_inner(acc, l_i, m_i, q1, q2, sm_scale, K1_block_ptr,
    K2_block_ptr, V_block_ptr, start_m, offs_m, offs_n, SEQLEN_K: tl.
    constexpr, WINDOW_SIZE: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N:
    tl.constexpr, EVEN_MN: tl.constexpr, STAGE: tl.constexpr):
    if STAGE == 1:
        hi = start_m * BLOCK_M - WINDOW_SIZE + BLOCK_M
        lo = start_m * BLOCK_M - WINDOW_SIZE
        if hi < 0:
            hi = 0
        if lo < 0:
            lo = 0
    elif STAGE == 2:
        hi = start_m * BLOCK_M
        lo = start_m * BLOCK_M - WINDOW_SIZE + BLOCK_M
        if lo < 0:
            lo = 0
    else:
        lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M
        lo = tl.multiple_of(lo, BLOCK_M)
        hi = min(hi, SEQLEN_K)
    EVEN_MASK_FREE = EVEN_MN & ((STAGE == 1) | (STAGE == 2))
    K1_block_ptr = tl.advance(K1_block_ptr, (0, lo))
    K2_block_ptr = tl.advance(K2_block_ptr, (0, lo))
    V_block_ptr = tl.advance(V_block_ptr, (lo, 0))
    for start_n in range(lo, hi, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        if EVEN_MASK_FREE:
            k1 = tl.load(K1_block_ptr)
        else:
            k1 = tl.load(K1_block_ptr, boundary_check=(1,), padding_option=
                'zero')
        qk = tl.dot(q1, k1) * (sm_scale * RCP_LN2)
        if STAGE == 1:
            if EVEN_MN:
                k2 = tl.load(K2_block_ptr)
            else:
                k2 = tl.load(K2_block_ptr, boundary_check=(1,),
                    padding_option='zero')
            qk2 = tl.dot(q2, k2) * (sm_scale * RCP_LN2)
            mask = offs_m[:, None] <= start_n + WINDOW_SIZE + offs_n[None, :]
            qk = tl.where(mask, qk, qk2)
        elif STAGE == 3:
            mask = offs_m[:, None] >= start_n + offs_n[None, :]
            qk += tl.where(mask, 0, NEGINF)
        if not EVEN_MASK_FREE:
            qk += tl.where((start_n + offs_n)[None, :] < SEQLEN_K, 0, NEGINF)
        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 *= alpha[:, None]
        if EVEN_MASK_FREE:
            v = tl.load(V_block_ptr)
        else:
            v = tl.load(V_block_ptr, boundary_check=(1,), padding_option='zero'
                )
        acc += tl.dot(p.to(V_block_ptr.dtype.element_ty), 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))
    return acc, l_i, m_i
