eva_agg_kernel.py

import math
import torch
import triton
import triton.language as tl

@triton.heuristics(
    {
        "EVEN_M": lambda args: args["seqlen_q"] % args["BLOCK_M"] == 0,
        "EVEN_N": lambda args: args["seqlen_k"] % args["BLOCK_N"] == 0,
        "EVEN_W": lambda args: args["WINDOW_SIZE"] % args["BLOCK_N"] == 0,
        "EVEN_HEADDIM": lambda args: args["headdim"] == args["BLOCK_HEADDIM"],
    }
)
@triton.jit
def _bwd_eva_agg_kernel_dkdv(
    Q,
    K,
    V,
    WindowMask,
    DO,
    LSE,
    DO_T_O,
    DK,
    DV,
    softmax_scale,
    stride_qb, stride_qh, stride_qm,
    stride_kb, stride_kh, stride_kn,
    stride_vb, stride_vh, stride_vn,
    stride_window_mask_b, stride_window_mask_m,
    stride_do_b, stride_do_h, stride_do_m,
    stride_lse_b, stride_lse_h,
    stride_do_t_o_b, stride_do_t_o_h,
    stride_dk_b, stride_dk_h, stride_dk_n,
    stride_dv_b, stride_dv_h, stride_dv_n,
    nheads,
    seqlen_q,
    seqlen_k,
    headdim,
    WINDOW_SIZE: tl.constexpr,
    MASK_TYPE: tl.constexpr,
    BLOCK_HEADDIM: tl.constexpr,
    EVEN_M: tl.constexpr,
    EVEN_N: tl.constexpr,
    EVEN_W: tl.constexpr,
    EVEN_HEADDIM: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    off_bh = tl.program_id(1)
    off_h = off_bh % nheads
    off_b = off_bh // nheads

    start_n = tl.program_id(0)
    # determine which window the current KV block belongs to
    offs_w = (start_n * BLOCK_N) // WINDOW_SIZE
    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_m = tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, BLOCK_HEADDIM)

    # initialize pointers
    q_ptrs = (
        Q + 
        off_b * stride_qb + 
        off_h * stride_qh +
        offs_m[:, None] * stride_qm + offs_d[None, :]
    )
    k_ptrs = (
        K + 
        off_b * stride_kb + 
        off_h * stride_kh +
        offs_n[:, None] * stride_kn + offs_d[None, :]
    )
    v_ptrs = (
        V + 
        off_b * stride_vb + 
        off_h * stride_vh +
        offs_n[:, None] * stride_vn + offs_d[None, :]
    )
    do_ptrs = (
        DO + 
        off_b * stride_do_b + 
        off_h * stride_do_h +
        offs_m[:, None] * stride_do_m + offs_d[None, :]
    )
    do_t_o_ptrs = (
        DO_T_O + 
        off_b * stride_do_t_o_b + 
        off_h * stride_do_t_o_h +
        offs_m[:, None]
    )
    lse_ptrs = (
        LSE + 
        off_b * stride_lse_b + 
        off_h * stride_lse_h +
        offs_m[:, None]
    )
    if MASK_TYPE == 1:
        m_ptrs = (
            WindowMask +
            off_b * stride_window_mask_b +
            (offs_m[:, None] * stride_window_mask_m + offs_n[None, :])
        )
    dk_ptrs = (
        DK + 
        off_b * stride_dk_b + 
        off_h * stride_dk_h +
        offs_n[:, None] * stride_dk_n + offs_d[None, :]
    )
    dv_ptrs = (
        DV + 
        off_b * stride_dv_b + 
        off_h * stride_dv_h +
        offs_n[:, None] * stride_dv_n + offs_d[None, :]
    )

    # 1. for singletons
    # determine start and end of query block
    begin_m = ((start_n * BLOCK_N) // BLOCK_M) * BLOCK_M
    end_m = tl.minimum((offs_w + 1) * WINDOW_SIZE, seqlen_q)

    dk = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
    dv = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
    if EVEN_N & EVEN_M:
        if EVEN_HEADDIM:
            k = tl.load(k_ptrs)
            v = tl.load(v_ptrs)
        else:
            k = tl.load(k_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
            v = tl.load(v_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
    else:
        if EVEN_HEADDIM:
            k = tl.load(k_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
            v = tl.load(v_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
        else:
            k = tl.load(
                k_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0
            )
            v = tl.load(
                v_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0
            )
    for start_m in range(begin_m, end_m, BLOCK_M):
        start_m = tl.multiple_of(start_m, BLOCK_M)
        # load q, do, and lse
        if EVEN_M & EVEN_N:
            if EVEN_HEADDIM:
                q = tl.load(
                    q_ptrs + start_m * stride_qm
                )
                do = tl.load(
                    do_ptrs + start_m * stride_do_m
                )
            else:
                q = tl.load(
                    q_ptrs + start_m * stride_qm, 
                    mask=offs_d[None, :] < headdim, 
                    other=0.0
                )
                do = tl.load(
                    do_ptrs + start_m * stride_do_m, 
                    mask=offs_d[None, :] < headdim, 
                    other=0.0
                )
            do_t_o = tl.load(
                do_t_o_ptrs + start_m
            )
            lse = tl.load(
                lse_ptrs + start_m
            )
        else:
            if EVEN_HEADDIM:
                q = tl.load(
                    q_ptrs + start_m * stride_qm, 
                    mask=(start_m + offs_m)[:, None] < seqlen_q, 
                    other=0.0
                )
                do = tl.load(
                    do_ptrs + start_m * stride_do_m, 
                    mask=(start_m + offs_m)[:, None] < seqlen_q, 
                    other=0.0
                )
            else:
                q = tl.load(
                    q_ptrs + start_m * stride_qm, 
                    mask=((start_m + offs_m)[:, None] < seqlen_q) & (offs_d[None, :] < headdim), 
                    other=0.0
                )
                do = tl.load(
                    do_ptrs + start_m * stride_do_m, 
                    mask=((start_m + offs_m)[:, None] < seqlen_q) & (offs_d[None, :] < headdim), 
                    other=0.0
                )
            do_t_o = tl.load(
                do_t_o_ptrs + start_m, 
                mask=(start_m + offs_m)[:, None] < seqlen_q, 
                other=0.0
            )
            lse = tl.load(
                lse_ptrs + start_m, 
                mask=(start_m + offs_m)[:, None] < seqlen_q, 
                other=0.0
            )
        lse = tl.where(lse == float("-inf"), 0.0, lse)
        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        qk += tl.dot(q, tl.trans(k))
        if not EVEN_M:
            qk += tl.where((start_m + offs_m)[:, None] < seqlen_q, 0, float("-inf"))

        if MASK_TYPE == 1:
            if EVEN_M & EVEN_W:
                mask = tl.load(
                    m_ptrs + (start_m * stride_window_mask_m) - (offs_w * WINDOW_SIZE)
                )
            else:
                mask = tl.load(
                    m_ptrs + (start_m * stride_window_mask_m) - (offs_w * WINDOW_SIZE),
                    mask=((start_m + offs_m)[:, None] < seqlen_q)
                    & (((start_m * stride_window_mask_m) - (offs_w * WINDOW_SIZE) + offs_n)[None, :] < WINDOW_SIZE),
                    other=1,
                )
            # Slightly faster to multiply the softmax_scale in the tl.exp below since the compiler
            # can then fuse the mult and add into an fma instruction. But if we have bias we need to
            # to multiply with softmax_scale here.
            # we assume mask already implies the causal masking
            qk = qk * softmax_scale
            qk = tl.where(mask, float("-inf"), qk)
            p = tl.exp(qk - lse)
        else:
            qk += tl.where((start_m + offs_m)[:, None] >= offs_n[None, :], 0, float("-inf"))
            p = tl.exp(qk * softmax_scale - lse)

        # dp [M, N]
        dp = tl.dot(do, tl.trans(v))
        # p [M, N],  dp [M, N], do_t_o [M, 1] -> ds [M, N]
        ds = (p * (dp - do_t_o) * softmax_scale).to(q.dtype)
        # p is fp32 and [M, N], convert to q.dtype
        # do [M, D] -> dv [N, D]
        dv += tl.dot(tl.trans(p.to(do.dtype)), do)
        # dk [N, D]
        dk += tl.dot(tl.trans(ds), q) 
    if EVEN_N & EVEN_M:
        if EVEN_HEADDIM:
            tl.store(dv_ptrs, dv)
            tl.store(dk_ptrs, dk)
        else:
            tl.store(dv_ptrs, dv, mask=offs_d[None, :] < headdim)
            tl.store(dk_ptrs, dk, mask=offs_d[None, :] < headdim)
    else:
        if EVEN_HEADDIM:
            tl.store(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k)
            tl.store(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k)
        else:
            tl.store(dv_ptrs, dv, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim))
            tl.store(dk_ptrs, dk, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim))

@triton.heuristics(
    {
        "EVEN_M": lambda args: args["seqlen_q"] % args["BLOCK_M"] == 0,
        "EVEN_C": lambda args: args["nchunks"] % args["BLOCK_N"] == 0,
        "EVEN_HEADDIM": lambda args: args["headdim"] == args["BLOCK_HEADDIM"],
    }
)
@triton.jit
def _bwd_eva_agg_kernel_drfa_kv(
    Q,
    RFA_K,
    RFA_V,
    ChunkMask,
    DO,
    LSE,
    DO_T_O,
    D_RFA_K,
    D_RFA_V,
    softmax_scale,
    stride_qb, stride_qh, stride_qm,
    stride_rfa_kb, stride_rfa_kh, stride_rfa_kc,
    stride_rfa_vb, stride_rfa_vh, stride_rfa_vc,
    stride_chunk_mask_b, stride_chunk_mask_m,
    stride_do_b, stride_do_h, stride_do_m,
    stride_lse_b, stride_lse_h,
    stride_do_t_o_b, stride_do_t_o_h,
    stride_d_rfa_k_b, stride_d_rfa_k_h, stride_d_rfa_k_c,
    stride_d_rfa_v_b, stride_d_rfa_v_h, stride_d_rfa_v_c,
    nheads,
    seqlen_q,
    nchunks,
    headdim,
    CHUNKS_PER_WINDOW: tl.constexpr,
    WINDOW_SIZE: tl.constexpr,
    MASK_TYPE: tl.constexpr,
    BLOCK_HEADDIM: tl.constexpr,
    EVEN_M: tl.constexpr,
    EVEN_C: tl.constexpr,
    EVEN_HEADDIM: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    off_bh = tl.program_id(1)
    off_h = off_bh % nheads
    off_b = off_bh // nheads
    start_c = tl.program_id(0)
    # there are 128 chunks per window
    offs_c = start_c * BLOCK_N + tl.arange(0, BLOCK_N)
    # determine which window the current KV block belongs to
    offs_w = (start_c * BLOCK_N) // CHUNKS_PER_WINDOW
    offs_m = tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, BLOCK_HEADDIM)

    # initialize pointers
    q_ptrs = (
        Q + 
        off_b * stride_qb + 
        off_h * stride_qh +
        (offs_m[:, None] * stride_qm + offs_d[None, :])
    )
    do_ptrs = (
        DO + 
        off_b * stride_do_b + 
        off_h * stride_do_h +
        (offs_m[:, None] * stride_do_m + offs_d[None, :])
    )
    do_t_o_ptrs = (
        DO_T_O + 
        off_b * stride_do_t_o_b + 
        off_h * stride_do_t_o_h +
        (offs_m[:, None])
    )
    lse_ptrs = (
        LSE + 
        off_b * stride_lse_b + 
        off_h * stride_lse_h +
        (offs_m[:, None])
    )
    rfa_k_ptrs = (
        RFA_K + 
        off_b * stride_rfa_kb + 
        off_h * stride_rfa_kh +
        (offs_c[:, None] * stride_rfa_kc + offs_d[None, :])
    )
    rfa_v_ptrs = (
        RFA_V + 
        off_b * stride_rfa_vb + 
        off_h * stride_rfa_vh +
        (offs_c[:, None] * stride_rfa_vc + offs_d[None, :])
    )
    if MASK_TYPE == 1:
        rfa_m_ptrs = (
            ChunkMask +
            off_b * stride_chunk_mask_b +
            (offs_m[:, None] * stride_chunk_mask_m + offs_c[None, :])
        )
    d_rfa_k_ptrs = (
        D_RFA_K + 
        off_b * stride_d_rfa_k_b + 
        off_h * stride_d_rfa_k_h +
        (offs_c[:, None] * stride_d_rfa_k_c + offs_d[None, :])
    )
    d_rfa_v_ptrs = (
        D_RFA_V + 
        off_b * stride_d_rfa_v_b + 
        off_h * stride_d_rfa_v_h +
        (offs_c[:, None] * stride_d_rfa_v_c + offs_d[None, :])
    )

    d_rfa_k = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
    d_rfa_v = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
    if EVEN_C & EVEN_M:
        if EVEN_HEADDIM:
            rfa_k = tl.load(rfa_k_ptrs)
            rfa_v = tl.load(rfa_v_ptrs)
        else:
            rfa_k = tl.load(rfa_k_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
            rfa_v = tl.load(rfa_v_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
    else:
        if EVEN_HEADDIM:
            rfa_k = tl.load(rfa_k_ptrs, mask=offs_c[:, None] < nchunks, other=0.0)
            rfa_v = tl.load(rfa_v_ptrs, mask=offs_c[:, None] < nchunks, other=0.0)
        else:
            rfa_k = tl.load(
                rfa_k_ptrs, mask=(offs_c[:, None] < nchunks) & (offs_d[None, :] < headdim), other=0.0
            )
            rfa_v = tl.load(
                rfa_v_ptrs, mask=(offs_c[:, None] < nchunks) & (offs_d[None, :] < headdim), other=0.0
            )
    begin_m = tl.minimum((offs_w + 1) * WINDOW_SIZE, seqlen_q)
    end_m = seqlen_q
    for start_m in range(begin_m, end_m, BLOCK_M):
        start_m = tl.multiple_of(start_m, BLOCK_M)
        # load q, do, and lse
        if EVEN_M:
            if EVEN_HEADDIM:
                q = tl.load(
                    q_ptrs + start_m * stride_qm
                )
                do = tl.load(
                    do_ptrs + start_m * stride_do_m
                )
            else:
                q = tl.load(
                    q_ptrs + start_m * stride_qm, 
                    mask=offs_d[None, :] < headdim, 
                    other=0.0
                )
                do = tl.load(
                    do_ptrs + start_m * stride_do_m, 
                    mask=offs_d[None, :] < headdim, 
                    other=0.0
                )
            do_t_o = tl.load(
                do_t_o_ptrs + start_m
            )
            lse = tl.load(
                lse_ptrs + start_m
            )
        else:
            if EVEN_HEADDIM:
                q = tl.load(
                    q_ptrs + start_m * stride_qm, 
                    mask=(start_m + offs_m)[:, None] < seqlen_q, 
                    other=0.0
                )
                do = tl.load(
                    do_ptrs + start_m * stride_do_m, 
                    mask=(start_m + offs_m)[:, None] < seqlen_q, 
                    other=0.0
                )
            else:
                q = tl.load(
                    q_ptrs + start_m * stride_qm, 
                    mask=((start_m + offs_m)[:, None] < seqlen_q) & (offs_d[None, :] < headdim), 
                    other=0.0
                )
                do = tl.load(
                    do_ptrs + start_m * stride_do_m, 
                    mask=((start_m + offs_m)[:, None] < seqlen_q) & (offs_d[None, :] < headdim), 
                    other=0.0
                )
            do_t_o = tl.load(
                do_t_o_ptrs + start_m, 
                mask=(start_m + offs_m)[:, None] < seqlen_q, 
                other=0.0
            )
            lse = tl.load(
                lse_ptrs + start_m, 
                mask=(start_m + offs_m)[:, None] < seqlen_q, 
                other=0.0
            )
        lse = tl.where(lse == float("-inf"), 0.0, lse)
        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        qk += tl.dot(q, tl.trans(rfa_k))
        if not EVEN_M:
            qk += tl.where((start_m + offs_m)[:, None] < seqlen_q, 0, float("-inf"))

        if MASK_TYPE == 1:
            if EVEN_M & EVEN_C:
                mask = tl.load(
                    rfa_m_ptrs + (start_m * stride_chunk_mask_m)
                )
            else:
                mask = tl.load(
                    rfa_m_ptrs + (start_m * stride_chunk_mask_m),
                    mask=((start_m + offs_m)[:, None] < seqlen_q)
                    & (offs_c[None, :] < nchunks),
                    other=1,
                )
            # Slightly faster to multiply the softmax_scale in the tl.exp below since the compiler
            # can then fuse the mult and add into an fma instruction. But if we have bias we need to
            # to multiply with softmax_scale here.
            # we assume mask already implies the causal masking
            qk = qk * softmax_scale
            qk = tl.where(mask, float("-inf"), qk)
            p = tl.exp(qk - lse)
        else:
            p = tl.exp(qk * softmax_scale - lse)

        dp = tl.dot(do, tl.trans(rfa_v))
        ds = (p * (dp - do_t_o) * softmax_scale).to(q.dtype)
        # p is fp32, convert to q.dtype
        d_rfa_v += tl.dot(tl.trans(p.to(do.dtype)), do)
        # move softmax_scale to ds to save computation
        d_rfa_k += tl.dot(tl.trans(ds), q) 
    if EVEN_C & EVEN_M:
        if EVEN_HEADDIM:
            tl.store(d_rfa_v_ptrs, d_rfa_v)
            tl.store(d_rfa_k_ptrs, d_rfa_k)
        else:
            tl.store(d_rfa_v_ptrs, d_rfa_v, mask=offs_d[None, :] < headdim)
            tl.store(d_rfa_k_ptrs, d_rfa_k, mask=offs_d[None, :] < headdim)
    else:
        if EVEN_HEADDIM:
            tl.store(d_rfa_v_ptrs, d_rfa_v, mask=offs_c[:, None] < nchunks)
            tl.store(d_rfa_k_ptrs, d_rfa_k, mask=offs_c[:, None] < nchunks)
        else:
            tl.store(d_rfa_v_ptrs, d_rfa_v, mask=(offs_c[:, None] < nchunks) & (offs_d[None, :] < headdim))
            tl.store(d_rfa_k_ptrs, d_rfa_k, mask=(offs_c[:, None] < nchunks) & (offs_d[None, :] < headdim))

@triton.heuristics(
    {
        "EVEN_M": lambda args: args["seqlen_q"] % args["BLOCK_M"] == 0,
        "EVEN_N": lambda args: args["seqlen_k"] % args["BLOCK_N"] == 0,
        "EVEN_C": lambda args: args["nchunks"] % args["BLOCK_N"] == 0,
        "EVEN_W": lambda args: args["WINDOW_SIZE"] % args["BLOCK_N"] == 0,
        "EVEN_HEADDIM": lambda args: args["headdim"] == args["BLOCK_HEADDIM"],
    }
)
@triton.jit
def _bwd_eva_agg_kernel_dq(
    Q,
    K,
    V,
    RFA_K,
    RFA_V,
    WindowMask,
    ChunkMask,
    DO,
    LSE,
    DO_T_O,
    DQ,
    softmax_scale,
    stride_qb, stride_qh, stride_qm,
    stride_kb, stride_kh, stride_kn,
    stride_vb, stride_vh, stride_vn,
    stride_rfa_kb, stride_rfa_kh, stride_rfa_kc,
    stride_rfa_vb, stride_rfa_vh, stride_rfa_vc,
    stride_window_mask_b, stride_window_mask_m,
    stride_chunk_mask_b, stride_chunk_mask_m,
    stride_do_b, stride_do_h, stride_do_m,
    stride_lse_b, stride_lse_h,
    stride_do_t_o_b, stride_do_t_o_h,
    stride_dq_b, stride_dq_h, stride_dq_m,
    nheads,
    seqlen_q,
    seqlen_k,
    nchunks,
    headdim,
    CHUNKS_PER_WINDOW: tl.constexpr,
    WINDOW_SIZE: tl.constexpr,
    MASK_TYPE: tl.constexpr,
    EMPTY_RFA_KV: tl.constexpr,
    BLOCK_HEADDIM: tl.constexpr,
    EVEN_M: tl.constexpr,
    EVEN_N: tl.constexpr,
    EVEN_W: tl.constexpr,
    EVEN_C: tl.constexpr,
    EVEN_HEADDIM: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    start_m = tl.program_id(0)
    off_bh = tl.program_id(1)
    off_h = off_bh % nheads
    off_b = off_bh // nheads
    # initialize offsets
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_w = (start_m * BLOCK_M) // WINDOW_SIZE
    offs_n = tl.arange(0, BLOCK_N)
    offs_c = tl.arange(0, BLOCK_N)
    offs_d = tl.arange(0, BLOCK_HEADDIM)
    # TODO: add paratheses or not
    q_ptrs = (
        Q +
        off_b * stride_qb +
        off_h * stride_qh +
        (offs_m[:, None] * stride_qm + offs_d[None, :])
    )
    k_ptrs = (
        K +
        off_b * stride_kb +
        off_h * stride_kh +
        (offs_n[:, None] * stride_kn + offs_d[None, :])
    )
    v_ptrs = (
        V +
        off_b * stride_vb +
        off_h * stride_vh +
        (offs_n[:, None] * stride_vn + offs_d[None, :])
    )
    if EMPTY_RFA_KV == 0:
        rfa_k_ptrs = (
            RFA_K +
            off_b * stride_rfa_kb +
            off_h * stride_rfa_kh +
            (offs_c[:, None] * stride_rfa_kc + offs_d[None, :])
        )
        rfa_v_ptrs = (
            RFA_V +
            off_b * stride_rfa_vb +
            off_h * stride_rfa_vh +
            (offs_c[:, None] * stride_rfa_vc + offs_d[None, :])
        )
    dq_ptrs = (
        DQ + 
        off_b * stride_dq_b +
        off_h * stride_dq_h +
        (offs_m[:, None] * stride_dq_m + offs_d[None, :])
    )
    do_ptrs = (
        DO + 
        off_b * stride_do_b +
        off_h * stride_do_h +
        (offs_m[:, None] * stride_do_m + offs_d[None, :])
    )
    do_t_o_ptrs = (
        DO_T_O + 
        off_b * stride_do_t_o_b +
        off_h * stride_do_t_o_h +
        offs_m[:, None]
    )
    lse_ptrs = (
        LSE + 
        off_b * stride_lse_b +
        off_h * stride_lse_h +
        offs_m[:, None]
    )
    ### load q, do, do_t_o, lse ####
    if EVEN_M:
        if EVEN_HEADDIM:
            q = tl.load(
                q_ptrs
            )
            do = tl.load(
                do_ptrs
            )
        else:
            q = tl.load(
                q_ptrs, 
                mask=offs_d[None, :] < headdim, 
                other=0.0
            )
            do = tl.load(
                do_ptrs, 
                mask=offs_d[None, :] < headdim, 
                other=0.0
            )
        do_t_o = tl.load(
            do_t_o_ptrs
        )
        lse = tl.load(
            lse_ptrs
        )
    else:
        if EVEN_HEADDIM:
            q = tl.load(
                q_ptrs, 
                mask=offs_m[:, None] < seqlen_q, 
                other=0.0
            )
            do = tl.load(
                do_ptrs, 
                mask=offs_m[:, None] < seqlen_q, 
                other=0.0
            )
        else:
            q = tl.load(
                q_ptrs, 
                mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), 
                other=0.0
            )
            do = tl.load(
                do_ptrs, 
                mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), 
                other=0.0
            )
        do_t_o = tl.load(
            do_t_o_ptrs, 
            mask=offs_m[:, None] < seqlen_q, 
            other=0.0
        )
        lse = tl.load(
            lse_ptrs, 
            mask=offs_m[:, None] < seqlen_q, 
            other=0.0
        )
    lse = tl.where(lse == float("-inf"), 0.0, lse)
    lse *= 1.4426950408889634  # log2(e)
    qk_scale = softmax_scale
    qk_scale *= 1.4426950408889634  # log2(e)
    if MASK_TYPE == 1:
        window_mask_ptrs = (
            WindowMask +
            off_b * stride_window_mask_b +
            (offs_m[:, None] * stride_window_mask_m + offs_n[None, :])
        )
        if EMPTY_RFA_KV == 0:
            chunk_mask_ptrs = (
                ChunkMask +
                off_b * stride_chunk_mask_b +
                (offs_m[:, None] * stride_chunk_mask_m + offs_c[None, :])
            )

    dq = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)
    # loop over k, v and update accumulator
    # Iterate over local singletons;
    # so we only iterate over blocks within the current window
    start_idx_n = offs_w * WINDOW_SIZE
    end_idx_n = tl.minimum((start_m + 1) * BLOCK_M, seqlen_k)
    for start_n in range(start_idx_n, end_idx_n, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        if EVEN_N & EVEN_M:
            if EVEN_HEADDIM:
                k = tl.load(
                    k_ptrs + start_n * stride_kn
                )
            else:
                k = tl.load(
                    k_ptrs + start_n * stride_kn,
                    mask=offs_d[None, :] < headdim,
                    other=0.0
                )
        else:
            if EVEN_HEADDIM:
                k = tl.load(
                    k_ptrs + start_n * stride_kn,
                    mask=(start_n + offs_n)[:, None] < seqlen_k,
                    other=0.0,
                )
            else:
                k = tl.load(
                    k_ptrs + start_n * stride_kn,
                    mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim),
                    other=0.0,
                )
        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        qk += tl.dot(q, tl.trans(k))
        # Trying to combine the two masks seem to make the result wrong
        if not EVEN_N:  # Need to mask out otherwise the softmax is wrong
            qk += tl.where((start_n + offs_n)[None, :] < seqlen_k, 0, float("-inf"))

        if MASK_TYPE == 1:
            if EVEN_M & EVEN_W:
                window_mask = tl.load(
                    window_mask_ptrs + start_n - start_idx_n
                )
            else:
                window_mask = tl.load(
                    window_mask_ptrs + start_n - start_idx_n,
                    mask=(offs_m[:, None] < seqlen_q)
                    & ((start_n - start_idx_n + offs_n)[None, :] < WINDOW_SIZE),
                    other=1,
                )
            # Slightly faster to multiply the softmax_scale in the tl.exp below since the compiler
            # can then fuse the mult and add into an fma instruction. But if we have bias we need to
            # to multiply with softmax_scale here.
            # we assume mask already implies the causal masking
            qk = qk * qk_scale
            qk = tl.where(window_mask, float("-inf"), qk)
            p = tl.exp2(qk - lse)
        else:
            qk += tl.where(offs_m[:, None] >= (start_n + offs_n)[None, :], 0, float("-inf"))
            p = tl.exp2(qk * qk_scale - lse)

        if EVEN_N & EVEN_M:
            if EVEN_HEADDIM:
                v = tl.load(
                    v_ptrs + start_n * stride_vn
                )
            else:
                v = tl.load(
                    v_ptrs + start_n * stride_vn,
                    mask=offs_d[None, :] < headdim,
                    other=0.0
                )
        else:
            if EVEN_HEADDIM:
                v = tl.load(
                    v_ptrs + start_n * stride_vn,
                    mask=(start_n + offs_n)[:, None] < seqlen_k,
                    other=0.0,
                )
            else:
                v = tl.load(
                    v_ptrs + start_n * stride_vn,
                    mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim),
                    other=0.0,
                )
        dp = tl.dot(do, tl.trans(v))
        ds = (p * (dp - do_t_o) * softmax_scale).to(q.dtype)
        dq += tl.dot(ds, k)

    if EMPTY_RFA_KV == 0:
        # Iterate over RFA chunks
        # we only iterate over chunks before the current local singleton window
        end_idx_c = tl.minimum(offs_w * CHUNKS_PER_WINDOW, nchunks)
        for start_c in range(0, end_idx_c, BLOCK_N):
            start_c = tl.multiple_of(start_c, BLOCK_N)
            # -- compute qk ----
            if EVEN_C & EVEN_M:
                if EVEN_HEADDIM:
                    rfa_k = tl.load(
                        rfa_k_ptrs + start_c * stride_rfa_kc
                    )
                else:
                    rfa_k = tl.load(
                        rfa_k_ptrs + start_c * stride_rfa_kc,
                        mask=offs_d[None, :] < headdim,
                        other=0.0
                    )
            else:
                if EVEN_HEADDIM:
                    rfa_k = tl.load(
                        rfa_k_ptrs + start_c * stride_rfa_kc,
                        mask=(start_c + offs_c)[:, None] < nchunks,
                        other=0.0,
                    )
                else:
                    rfa_k = tl.load(
                        rfa_k_ptrs + start_c * stride_rfa_kc,
                        mask=((start_c + offs_c)[:, None] < nchunks) & (offs_d[None, :] < headdim),
                        other=0.0,
                    )
            qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
            qk += tl.dot(q, tl.trans(rfa_k))
            # Trying to combine the two masks seem to make the result wrong
            if not EVEN_C:  # Need to mask out otherwise the softmax is wrong
                qk += tl.where((start_c + offs_c)[None, :] < nchunks, 0, float("-inf"))

            if MASK_TYPE == 1:
                if EVEN_C & EVEN_M:
                    chunk_mask = tl.load(
                        chunk_mask_ptrs + start_c
                    )
                else:
                    chunk_mask = tl.load(
                        chunk_mask_ptrs + start_c,
                        mask=(offs_m[:, None] < seqlen_q) & ((start_c + offs_c)[None, :] < nchunks),
                        other=1,
                    )
                # Slightly faster to multiply the softmax_scale in the tl.exp below since the compiler
                # can then fuse the mult and add into an fma instruction. But if we have bias we need to
                # to multiply with softmax_scale here.
                # we assume mask already implies the causal masking
                qk = qk * qk_scale
                qk = tl.where(chunk_mask, float("-inf"), qk)
                p = tl.exp2(qk - lse)
            else:
                p = tl.exp2(qk * qk_scale - lse)

            if EVEN_C & EVEN_M:  
                if EVEN_HEADDIM:
                    rfa_v = tl.load(
                        rfa_v_ptrs + start_c * stride_rfa_vc
                    )
                else:
                    rfa_v = tl.load(
                        rfa_v_ptrs + start_c * stride_rfa_vc,
                        mask=offs_d[None, :] < headdim,
                        other=0.0
                    )
            else:
                if EVEN_HEADDIM:
                    rfa_v = tl.load(
                        rfa_v_ptrs + start_c * stride_rfa_vc,
                        mask=(start_c + offs_n)[:, None] < nchunks,
                        other=0.0,
                    )
                else:
                    rfa_v = tl.load(
                        rfa_v_ptrs + start_c * stride_rfa_vc,
                        mask=((start_c + offs_n)[:, None] < nchunks) & (offs_d[None, :] < headdim),
                        other=0.0,
                    )
            dp = tl.dot(do, tl.trans(rfa_v))
            ds = (p * (dp - do_t_o) * softmax_scale).to(q.dtype)
            dq += tl.dot(ds, rfa_k)

    start_m = tl.program_id(0)
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, BLOCK_HEADDIM)
    dq_ptrs = (
        DQ +
        off_b * stride_dq_b +
        off_h * stride_dq_h +
        (offs_m[:, None] * stride_dq_m + offs_d[None, :])
    )
    if EVEN_M:
        if EVEN_HEADDIM:
            tl.store(
                dq_ptrs, dq
            )
        else:
            tl.store(
                dq_ptrs, dq,
                mask=offs_d[None, :] < headdim
            )
    else:
        if EVEN_HEADDIM:
            tl.store(
                dq_ptrs, dq,
                mask=offs_m[:, None] < seqlen_q
            )
        else:
            tl.store(
                dq_ptrs, dq,
                mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim)
            )

_capability_90_config = {
    "fwd": {
        (torch.bfloat16, 64): (128, 128, 4, 3), 
        (torch.bfloat16, 128): (128, 128, 8, 3),
        (torch.float32, 64): (128, 64, 8, 3),
        (torch.float32, 128): (64, 32, 4, 3),
    },
    "bwd_dq": {
        (torch.bfloat16, 64): (128, 64, 4, 3),
        (torch.bfloat16, 128): (128, 64, 8, 3),
        (torch.float32, 64): (128, 64, 8, 2),
        (torch.float32, 128): (32, 32, 4, 2),
    },
    "bwd_dkdv": {
        (torch.bfloat16, 64): (128, 64, 4, 2),
        (torch.bfloat16, 128): (128, 64, 8, 2),
        (torch.float32, 64): (128, 64, 8, 2),
        (torch.float32, 128): (32, 32, 4, 1),
    },
    "bwd_drfa_kv": {
        (torch.bfloat16, 64): (128, 64, 4, 2),
        (torch.bfloat16, 128): (128, 64, 8, 2),
        (torch.float32, 64): (128, 64, 8, 2),
        (torch.float32, 128): (32, 32, 4, 1),
    }
}

_capability_80_config = {
    "fwd": {
        (torch.bfloat16, 64): (64, 64, 4, 3),
        (torch.bfloat16, 128): (64, 64, 8, 3),
        (torch.float32, 64): (64, 32, 4, 2),
        (torch.float32, 128): (64, 32, 8, 1),
    },
    "bwd_dq": {
        (torch.bfloat16, 64): (64, 64, 4, 3),
        (torch.bfloat16, 128): (64, 32, 4, 2),
        (torch.float32, 64): (32, 32, 4, 2),
        (torch.float32, 128): (32, 32, 4, 2),
    },
    "bwd_dkdv": {
        (torch.bfloat16, 64): (64, 64, 4, 3),
        (torch.bfloat16, 128): (32, 32, 4, 2),
        (torch.float32, 64): (32, 32, 4, 1),
        (torch.float32, 128): (16, 64, 8, 1),
    },
    "bwd_drfa_kv": {
        (torch.bfloat16, 64): (64, 64, 4, 3),
        (torch.bfloat16, 128): (64, 32, 4, 3),
        (torch.float32, 64): (32, 32, 4, 1),
        (torch.float32, 128): (32, 32, 4, 1),
    }
}

def _get_config(dtype, head_dim, mode) -> tuple[int, int, int, int]:
    capability = torch.cuda.get_device_capability()
    if capability >= (9, 0):
        kernel_config = _capability_90_config[mode].get((dtype, head_dim), (32, 32, 4, 1))
    elif capability >= (8, 0):
        kernel_config = _capability_80_config[mode].get((dtype, head_dim), (16, 16, 4, 1))
    else:
        if mode == "fwd":
            if dtype == torch.float32:
                kernel_config = (32, 16, 4, 2)
            else:
                kernel_config = (64, 32, 4, 2)
        else:
            if dtype == torch.float32:
                kernel_config = (16, 16, 4, 1)
            else:
                kernel_config = (32, 32, 4, 1)
    return kernel_config

@triton.heuristics(
    {
        "EVEN_M": lambda args: args["seqlen_q"] % args["BLOCK_M"] == 0,
        "EVEN_N": lambda args: args["seqlen_k"] % args["BLOCK_N"] == 0,
        "EVEN_C": lambda args: args["nchunks"] % args["BLOCK_N"] == 0,
        "EVEN_W": lambda args: args["WINDOW_SIZE"] % args["BLOCK_N"] == 0,
        "EVEN_HEADDIM": lambda args: args["headdim"] == args["BLOCK_HEADDIM"],
    }
)
@triton.jit
def _fwd_eva_agg_kernel(
    Q,
    K,
    V,
    RFA_K,
    RFA_V,
    WindowMask,
    ChunkMask,
    Out,
    LSE,
    softmax_scale,
    stride_qb, stride_qh, stride_qm,
    stride_kb, stride_kh, stride_kn,
    stride_vb, stride_vh, stride_vn,
    stride_rfa_kb, stride_rfa_kh, stride_rfa_kc,
    stride_rfa_vb, stride_rfa_vh, stride_rfa_vc,
    stride_window_mask_b, stride_window_mask_m,
    stride_chunk_mask_b, stride_chunk_mask_m,
    stride_ob, stride_oh, stride_om,
    stride_lse_b, stride_lse_h,
    nheads,
    seqlen_q,
    seqlen_k,
    nchunks,
    headdim,
    CHUNKS_PER_WINDOW: tl.constexpr,
    WINDOW_SIZE: tl.constexpr,
    MASK_TYPE: tl.constexpr,
    EMPTY_RFA_KV: tl.constexpr,
    BLOCK_HEADDIM: tl.constexpr,
    EVEN_M: tl.constexpr,
    EVEN_N: tl.constexpr,
    EVEN_W: tl.constexpr,
    EVEN_C: tl.constexpr,
    EVEN_HEADDIM: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    start_m = tl.program_id(0)
    off_bh = tl.program_id(1)
    off_h = off_bh % nheads
    off_b = off_bh // nheads
    # initialize offsets
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_w = (start_m * BLOCK_M) // WINDOW_SIZE
    offs_n = tl.arange(0, BLOCK_N)
    offs_c = tl.arange(0, BLOCK_N)
    offs_d = tl.arange(0, BLOCK_HEADDIM)
    # TODO: add paratheses or not
    q_ptrs = (
        Q +
        off_b * stride_qb +
        off_h * stride_qh +
        (offs_m[:, None] * stride_qm + offs_d[None, :])
    )
    k_ptrs = (
        K +
        off_b * stride_kb +
        off_h * stride_kh +
        (offs_n[:, None] * stride_kn + offs_d[None, :])
    )
    v_ptrs = (
        V +
        off_b * stride_vb +
        off_h * stride_vh +
        (offs_n[:, None] * stride_vn + offs_d[None, :])
    )
    if EMPTY_RFA_KV == 0:
        rfa_k_ptrs = (
            RFA_K +
            off_b * stride_rfa_kb +
            off_h * stride_rfa_kh +
            (offs_c[:, None] * stride_rfa_kc + offs_d[None, :])
        )
        rfa_v_ptrs = (
            RFA_V +
            off_b * stride_rfa_vb +
            off_h * stride_rfa_vh +
            (offs_c[:, None] * stride_rfa_vc + offs_d[None, :])
        )

    qk_scale = softmax_scale
    qk_scale *= 1.4426950408889634  # log2(e)
    if MASK_TYPE == 1:
        window_mask_ptrs = (
            WindowMask +
            off_b * stride_window_mask_b +
            (offs_m[:, None] * stride_window_mask_m + offs_n[None, :])
        )
        if EMPTY_RFA_KV == 0:
            chunk_mask_ptrs = (
                ChunkMask +
                off_b * stride_chunk_mask_b +
                (offs_m[:, None] * stride_chunk_mask_m + offs_c[None, :])
            )

    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
    d_i = tl.zeros([BLOCK_M], dtype=tl.float32)
    acc_o = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)
    # load q: it will stay in SRAM throughout
    # [2022-10-30] TD: Triton bug - in the case of EVEN_M=True and EVEN_N=False, if we just call
    # tl.load(q_ptrs), we get the wrong output!
    if EVEN_M & EVEN_N:
        if EVEN_HEADDIM:
            q = tl.load(
                q_ptrs
            )
        else:
            q = tl.load(
                q_ptrs,
                mask=offs_d[None, :] < headdim,
                other=0.0
            )
    else:
        if EVEN_HEADDIM:
            q = tl.load(
                q_ptrs,
                mask=offs_m[:, None] < seqlen_q,
                other=0.0
            )
        else:
            q = tl.load(
                q_ptrs,
                mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim),
                other=0.0
            )
    # loop over k, v and update accumulator
    # Iterate over local singletons;
    # so we only iterate over blocks within the current window
    start_idx_n = offs_w * WINDOW_SIZE
    end_idx_n = tl.minimum((start_m + 1) * BLOCK_M, seqlen_k)
    for start_n in range(start_idx_n, end_idx_n, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        # -- compute qk ----
        if EVEN_N & EVEN_M:
            if EVEN_HEADDIM:
                k = tl.load(
                    k_ptrs + start_n * stride_kn
                )
            else:
                k = tl.load(
                    k_ptrs + start_n * stride_kn,
                    mask=offs_d[None, :] < headdim,
                    other=0.0
                )
        else:
            if EVEN_HEADDIM:
                k = tl.load(
                    k_ptrs + start_n * stride_kn,
                    mask=(start_n + offs_n)[:, None] < seqlen_k,
                    other=0.0,
                )
            else:
                k = tl.load(
                    k_ptrs + start_n * stride_kn,
                    mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim),
                    other=0.0,
                )
        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        qk += tl.dot(q, tl.trans(k))
        # Trying to combine the two masks seem to make the result wrong
        if not EVEN_N:  # Need to mask out otherwise the softmax is wrong
            qk += tl.where((start_n + offs_n)[None, :] < seqlen_k, 0, float("-inf"))

        if MASK_TYPE == 1:
            if EVEN_M & EVEN_W:
                window_mask = tl.load(
                    window_mask_ptrs + start_n - start_idx_n
                )
            else:
                window_mask = tl.load(
                    window_mask_ptrs + start_n - start_idx_n,
                    mask=(offs_m[:, None] < seqlen_q)
                    & ((start_n - start_idx_n + offs_n)[None, :] < WINDOW_SIZE),
                    other=1,
                )
            # Slightly faster to multiply the softmax_scale in the tl.exp below since the compiler
            # can then fuse the mult and add into an fma instruction. But if we have bias we need to
            # to multiply with softmax_scale here.
            # we assume mask already implies the causal masking
            qk = qk * qk_scale
            qk = tl.where(window_mask, float("-inf"), qk)
            m_ij = tl.maximum(tl.max(qk, 1), m_i)
            masked_out_rows = (m_ij == float("-inf"))
            m_ij_masked = tl.where(masked_out_rows, 0, m_ij)
            p = tl.exp2(qk - m_ij_masked[:, None])
        else:
            qk += tl.where(offs_m[:, None] >= (start_n + offs_n)[None, :], 0, float("-inf"))
            m_ij = tl.maximum(tl.max(qk, 1) * qk_scale, m_i)
            masked_out_rows = (m_ij == float("-inf"))
            m_ij_masked = tl.where(masked_out_rows, 0, m_ij)
            p = tl.exp2(qk * qk_scale - m_ij_masked[:, None])

        d_ij = tl.sum(p, 1)

        # scale acc_o
        prev_scale = tl.exp2(m_i - m_ij_masked)
        # # -- update output accumulator --
        acc_o = acc_o * prev_scale[:, None]
        # update acc_o
        if EVEN_N & EVEN_M:  # If we just do "if EVEN_N", there seems to be some race condition
            if EVEN_HEADDIM:
                v = tl.load(
                    v_ptrs + start_n * stride_vn
                )
            else:
                v = tl.load(
                    v_ptrs + start_n * stride_vn,
                    mask=offs_d[None, :] < headdim,
                    other=0.0
                )
        else:
            if EVEN_HEADDIM:
                v = tl.load(
                    v_ptrs + start_n * stride_vn,
                    mask=(start_n + offs_n)[:, None] < seqlen_k,
                    other=0.0,
                )
            else:
                v = tl.load(
                    v_ptrs + start_n * stride_vn,
                    mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim),
                    other=0.0,
                )
        p = p.to(v.dtype)
        acc_o = tl.dot(p, v, acc_o)

        # -- update statistics
        d_i = d_i * prev_scale + d_ij
        m_i = m_ij

    if EMPTY_RFA_KV == 0:
        # Iterate over RFA chunks
        # we only iterate over chunks before the current local singleton window
        end_idx_c = tl.minimum(offs_w * CHUNKS_PER_WINDOW, nchunks)
        for start_c in range(0, end_idx_c, BLOCK_N):
            start_c = tl.multiple_of(start_c, BLOCK_N)
            # -- compute qk ----
            if EVEN_C & EVEN_M:
                if EVEN_HEADDIM:
                    rfa_k = tl.load(
                        rfa_k_ptrs + start_c * stride_rfa_kc
                    )
                else:
                    rfa_k = tl.load(
                        rfa_k_ptrs + start_c * stride_rfa_kc,
                        mask=offs_d[None, :] < headdim,
                        other=0.0
                    )
            else:
                if EVEN_HEADDIM:
                    rfa_k = tl.load(
                        rfa_k_ptrs + start_c * stride_rfa_kc,
                        mask=(start_c + offs_c)[:, None] < nchunks,
                        other=0.0,
                    )
                else:
                    rfa_k = tl.load(
                        rfa_k_ptrs + start_c * stride_rfa_kc,
                        mask=((start_c + offs_c)[:, None] < nchunks) & (offs_d[None, :] < headdim),
                        other=0.0,
                    )
            qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
            qk += tl.dot(q, tl.trans(rfa_k))
            # Trying to combine the two masks seem to make the result wrong
            if not EVEN_C:  # Need to mask out otherwise the softmax is wrong
                qk += tl.where((start_c + offs_c)[None, :] < nchunks, 0, float("-inf"))

            if MASK_TYPE == 1:
                if EVEN_C & EVEN_M:
                    chunk_mask = tl.load(
                        chunk_mask_ptrs + start_c
                    )
                else:
                    chunk_mask = tl.load(
                        chunk_mask_ptrs + start_c,
                        mask=(offs_m[:, None] < seqlen_q) & ((start_c + offs_c)[None, :] < nchunks),
                        other=1,
                    )
                # Slightly faster to multiply the softmax_scale in the tl.exp below since the compiler
                # can then fuse the mult and add into an fma instruction. But if we have bias we need to
                # to multiply with softmax_scale here.
                # we assume mask already implies the causal masking
                qk = qk * qk_scale
                qk = tl.where(chunk_mask, float("-inf"), qk)
                m_ij = tl.maximum(tl.max(qk, 1), m_i)
                masked_out_rows = (m_ij == float("-inf"))
                m_ij_masked = tl.where(masked_out_rows, 0, m_ij)
                p = tl.exp2(qk - m_ij_masked[:, None])
            else:
                m_ij = tl.maximum(tl.max(qk, 1) * qk_scale, m_i)
                masked_out_rows = (m_ij == float("-inf"))
                m_ij_masked = tl.where(masked_out_rows, 0, m_ij)
                p = tl.exp2(qk * qk_scale - m_ij_masked[:, None])

            d_ij = tl.sum(p, 1)

            # scale acc_o
            prev_scale = tl.exp2(m_i - m_ij_masked)
            # # -- update output accumulator --
            acc_o = acc_o * prev_scale[:, None]
            # update acc_o
            # TODO: If we just do "if EVEN_N", there seems to be some race condition ?
            if EVEN_C & EVEN_M:  
                if EVEN_HEADDIM:
                    rfa_v = tl.load(
                        rfa_v_ptrs + start_c * stride_rfa_vc
                    )
                else:
                    rfa_v = tl.load(
                        rfa_v_ptrs + start_c * stride_rfa_vc,
                        mask=offs_d[None, :] < headdim,
                        other=0.0
                    )
            else:
                if EVEN_HEADDIM:
                    rfa_v = tl.load(
                        rfa_v_ptrs + start_c * stride_rfa_vc,
                        mask=(start_c + offs_n)[:, None] < nchunks,
                        other=0.0,
                    )
                else:
                    rfa_v = tl.load(
                        rfa_v_ptrs + start_c * stride_rfa_vc,
                        mask=((start_c + offs_n)[:, None] < nchunks) & (offs_d[None, :] < headdim),
                        other=0.0,
                    )
            p = p.to(rfa_v.dtype)
            acc_o = tl.dot(p, rfa_v, acc_o)

            # -- update statistics
            d_i = d_i * prev_scale + d_ij
            m_i = m_ij

    # for rows that are all -inf, set d_i to 1.0
    d_i = tl.where(d_i == 0.0, 1.0, d_i)
    # multiply by log(2)
    lse_m = (m_i + tl.math.log2(d_i)) * 0.6931471805599453
    acc_o = acc_o / d_i[:, None]
    # TODO: understand why rematerialize offsets to save registers?
    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:
        if EVEN_HEADDIM:
            tl.store(
                out_ptrs, acc_o
            )
        else:
            tl.store(
                out_ptrs, acc_o,
                mask=offs_d[None, :] < headdim
            )
    else:
        if EVEN_HEADDIM:
            tl.store(
                out_ptrs, acc_o,
                mask=offs_m[:, None] < seqlen_q
            )
        else:
            tl.store(
                out_ptrs, acc_o,
                mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim)
            )
    lse_ptrs = (
        LSE +
        off_b * stride_lse_b +
        off_h * stride_lse_h +
        offs_m
    )
    if EVEN_M:
        tl.store(
            lse_ptrs, lse_m,
        )
    else:
        tl.store(
            lse_ptrs, lse_m,
            mask=offs_m < seqlen_q
        )

def triton_eva_agg_fwd(
    q, k, v, rfa_k, rfa_v, 
    window_mask, 
    chunk_mask,
    softmax_scale, 
    window_size, 
    chunks_per_window
):
    if rfa_k is None and rfa_v is None:
        empty_rfa_kv = 1

        q, k, v = [
            x if x.stride(-1) == 1 else x.contiguous() 
            for x in [q, k, v]
        ]
    else:
        assert rfa_k is not None and rfa_v is not None, "Both rfa_k and rfa_v must either be None or have values at the same time."
        empty_rfa_kv = 0

        q, k, v, rfa_k, rfa_v = [
            x if x.stride(-1) == 1 else x.contiguous() 
            for x in [q, k, v, rfa_k, rfa_v]
        ]

    # shape constraints
    batch, nheads, seqlen_q, head_dim = q.shape
    _,     _,      seqlen_k, _        = k.shape
    if empty_rfa_kv == 0:
        nchunks = rfa_k.shape[-2]
        assert rfa_k.shape == (batch, nheads, nchunks, head_dim)
        assert rfa_v.shape == (batch, nheads, nchunks, head_dim)
        assert q.dtype == k.dtype == v.dtype == rfa_k.dtype == rfa_v.dtype
    else:
        nchunks = 0
        assert q.dtype == k.dtype == v.dtype, "All tensors must have the same type"
    assert k.shape == (batch, nheads, seqlen_k, head_dim)
    assert v.shape == (batch, nheads, seqlen_k, head_dim)

    assert head_dim <= 128, "We only test head dimensions up to 128"
    # assert q.dtype in [torch.float16, torch.bfloat16], "Only support fp16 and bf16"
    assert q.dtype in [torch.bfloat16, torch.float], "Only support bf16 and fp32 for now"
    assert q.is_cuda and k.is_cuda and v.is_cuda
    softmax_scale = softmax_scale or 1.0 / math.sqrt(head_dim)

    mask_type = 0
    if window_mask is not None:
        mask_type = 1
        assert window_mask.dtype == torch.bool
        assert window_mask.is_cuda
        assert window_mask.dim() == 4
        assert window_mask.shape == (batch, 1, seqlen_q, window_size)
        if window_mask.stride(-1) != 1:
            window_mask = window_mask.contiguous()

        assert chunk_mask is not None
        assert chunk_mask.dtype == torch.bool
        assert chunk_mask.is_cuda
        assert chunk_mask.dim() == 4
        assert chunk_mask.shape == (batch, 1, seqlen_q, nchunks)
        if chunk_mask.stride(-1) != 1:
            chunk_mask = chunk_mask.contiguous()

    chunk_mask_strides = (
        (chunk_mask.stride(0), chunk_mask.stride(2)) 
        if mask_type == 1 else 
        (0, 0)
    )
    window_mask_strides = (
        (window_mask.stride(0), window_mask.stride(2)) 
        if mask_type == 1 else 
        (0, 0)
    )

    rfa_k_strides = (
        (rfa_k.stride(0), rfa_k.stride(1), rfa_k.stride(2))
        if empty_rfa_kv == 0 else
        (0, 0, 0)
    )
    rfa_v_strides = (
        (rfa_v.stride(0), rfa_v.stride(1), rfa_v.stride(2))
        if empty_rfa_kv == 0 else
        (0, 0, 0)
    )

    o = torch.empty_like(q)
    lse = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)

    BLOCK_HEADDIM = max(triton.next_power_of_2(head_dim), 16)
    
    BLOCK_M, BLOCK_N, num_warps, num_stages = _get_config(q.dtype, head_dim, "fwd")

    assert chunks_per_window >= BLOCK_N, "chunks_per_window must be greater than BLOCK" 
    assert chunks_per_window % BLOCK_N == 0, "chunks_per_window must be a multiple of BLOCK_N"

    grid = lambda META: (triton.cdiv(seqlen_q, META["BLOCK_M"]), batch * nheads)
    _fwd_eva_agg_kernel[grid](
        q,
        k,
        v,
        rfa_k,
        rfa_v,
        window_mask,
        chunk_mask,
        o,
        lse,
        softmax_scale,
        q.stride(0), q.stride(1), q.stride(2),
        k.stride(0), k.stride(1), k.stride(2),
        v.stride(0), v.stride(1), v.stride(2),
        rfa_k_strides[0], rfa_k_strides[1], rfa_k_strides[2],
        rfa_v_strides[0], rfa_v_strides[1], rfa_v_strides[2],
        window_mask_strides[0], window_mask_strides[1],
        chunk_mask_strides[0], chunk_mask_strides[1],
        o.stride(0), o.stride(1), o.stride(2),
        lse.stride(0), lse.stride(1),
        nheads,
        seqlen_q,
        seqlen_k,
        nchunks,
        head_dim,
        chunks_per_window,
        window_size,
        mask_type,
        empty_rfa_kv,
        BLOCK_HEADDIM,
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        num_warps=num_warps,
        num_stages=num_stages,
    )
    return o, lse

def triton_eva_agg_bwd(
    do, 
    q, k, v, rfa_k, rfa_v, 
    window_mask, chunk_mask,
    o, lse, 
    dq, dk, dv, d_rfa_k, d_rfa_v, 
    softmax_scale, 
    window_size, 
    chunks_per_window,
    empty_rfa_kv,
    mask_type,
):
    if do.stride(-1) != 1:
        do = do.contiguous()

    # shape constraints
    batch, nheads, seqlen_q, head_dim = q.shape
    _,     _,      seqlen_k, _        = k.shape
    if empty_rfa_kv == 0:
        nchunks = rfa_k.shape[-2]
        assert rfa_k.shape == (batch, nheads, nchunks, head_dim)
        assert rfa_v.shape == (batch, nheads, nchunks, head_dim)
        assert d_rfa_k.stride(-1) == d_rfa_v.stride(-1) == 1
        assert q.dtype == k.dtype == v.dtype == rfa_k.dtype == rfa_v.dtype
    else:
        nchunks = 0
        assert q.dtype == k.dtype == v.dtype, "All tensors must have the same type"

    assert lse.shape == (batch, nheads, seqlen_q)
    assert q.stride(-1) == k.stride(-1) == v.stride(-1) == o.stride(-1) == rfa_k.stride(-1) == rfa_v.stride(-1) == 1
    assert dq.stride(-1) == dk.stride(-1) == dv.stride(-1) == 1
    softmax_scale = softmax_scale or 1.0 / math.sqrt(head_dim)

    assert head_dim <= 128, "We only test head dimensions up to 128"

    window_mask_strides = (
        (window_mask.stride(0), window_mask.stride(2)) 
        if mask_type == 1 else 
        (0, 0)
    )
    chunk_mask_strides = (
        (chunk_mask.stride(0), chunk_mask.stride(2)) 
        if mask_type == 1 else 
        (0, 0)
    )

    rfa_k_strides = (
        (rfa_k.stride(0), rfa_k.stride(1), rfa_k.stride(2))
        if empty_rfa_kv == 0 else
        (0, 0, 0)
    )
    rfa_v_strides = (
        (rfa_v.stride(0), rfa_v.stride(1), rfa_v.stride(2))
        if empty_rfa_kv == 0 else
        (0, 0, 0)
    )

    d_rfa_k_strides = (
        (d_rfa_k.stride(0), d_rfa_k.stride(1), d_rfa_k.stride(2))
        if empty_rfa_kv == 0 else
        (0, 0, 0)
    )
    d_rfa_v_strides = (
        (d_rfa_v.stride(0), d_rfa_v.stride(1), d_rfa_v.stride(2))
        if empty_rfa_kv == 0 else
        (0, 0, 0)
    )

    BLOCK_HEADDIM = max(triton.next_power_of_2(head_dim), 16)

    do_t_o = torch.sum(do.to(torch.float32) * o.to(torch.float32), dim=-1).to(do.dtype)

    BLOCK_M, BLOCK_N, num_warps, num_stages = _get_config(q.dtype, head_dim, "bwd_dq")

    assert chunks_per_window >= BLOCK_N, "chunks_per_window must be greater than BLOCK" 
    assert chunks_per_window % BLOCK_N == 0, "chunks_per_window must be a multiple of BLOCK"
    grid = lambda META: (
        triton.cdiv(seqlen_q, META["BLOCK_M"]),
        batch * nheads,
    )
    _bwd_eva_agg_kernel_dq[grid](
        q,
        k,
        v,
        rfa_k,
        rfa_v,
        window_mask,
        chunk_mask,
        do,
        lse,
        do_t_o,
        dq,
        softmax_scale,
        q.stride(0), q.stride(1), q.stride(2),
        k.stride(0), k.stride(1), k.stride(2),
        v.stride(0), v.stride(1), v.stride(2),
        rfa_k_strides[0], rfa_k_strides[1], rfa_k_strides[2],
        rfa_v_strides[0], rfa_v_strides[1], rfa_v_strides[2],
        window_mask_strides[0], window_mask_strides[1],
        chunk_mask_strides[0], chunk_mask_strides[1],
        do.stride(0), do.stride(1), do.stride(2),
        lse.stride(0), lse.stride(1),
        do_t_o.stride(0), do_t_o.stride(1),
        dq.stride(0), dq.stride(1), dq.stride(2),
        nheads,
        seqlen_q,
        seqlen_k,
        nchunks,
        head_dim,
        chunks_per_window,
        window_size,
        mask_type,
        empty_rfa_kv,
        BLOCK_HEADDIM,
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        num_warps=num_warps,
        num_stages=num_stages,
    )

    BLOCK_M, BLOCK_N, num_warps, num_stages = _get_config(q.dtype, head_dim, "bwd_dkdv")
    grid = lambda META: (
        triton.cdiv(seqlen_k, META["BLOCK_N"]),
        batch * nheads,
    )
    _bwd_eva_agg_kernel_dkdv[grid](
        q,
        k,
        v,
        window_mask,
        do,
        lse,
        do_t_o,
        dk,
        dv,
        softmax_scale,
        q.stride(0), q.stride(1), q.stride(2),
        k.stride(0), k.stride(1), k.stride(2),
        v.stride(0), v.stride(1), v.stride(2),
        window_mask_strides[0], window_mask_strides[1],
        do.stride(0), do.stride(1), do.stride(2),
        lse.stride(0), lse.stride(1),
        do_t_o.stride(0), do_t_o.stride(1),
        dk.stride(0), dk.stride(1), dk.stride(2),
        dv.stride(0), dv.stride(1), dv.stride(2),
        nheads,
        seqlen_q,
        seqlen_k,
        head_dim,
        window_size,
        mask_type,
        BLOCK_HEADDIM,
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        num_warps=num_warps,
        num_stages=num_stages,
    )
    if empty_rfa_kv == 0:
        BLOCK_M, BLOCK_N, num_warps, num_stages = _get_config(q.dtype, head_dim, "bwd_drfa_kv")
        grid = lambda META: (
            triton.cdiv(nchunks, META["BLOCK_N"]),
            batch * nheads,
        )
        _bwd_eva_agg_kernel_drfa_kv[grid](
            q,
            rfa_k,
            rfa_v,
            chunk_mask,
            do,
            lse,
            do_t_o,
            d_rfa_k,
            d_rfa_v,
            softmax_scale,
            q.stride(0), q.stride(1), q.stride(2),
            rfa_k_strides[0], rfa_k_strides[1], rfa_k_strides[2],
            rfa_v_strides[0], rfa_v_strides[1], rfa_v_strides[2],
            chunk_mask_strides[0], chunk_mask_strides[1],
            do.stride(0), do.stride(1), do.stride(2),
            lse.stride(0), lse.stride(1),
            do_t_o.stride(0), do_t_o.stride(1),
            d_rfa_k_strides[0], d_rfa_k_strides[1], d_rfa_k_strides[2],
            d_rfa_v_strides[0], d_rfa_v_strides[1], d_rfa_v_strides[2],
            nheads,
            seqlen_q,
            nchunks,
            head_dim,
            chunks_per_window,
            window_size,
            mask_type,
            BLOCK_HEADDIM,
            BLOCK_M=BLOCK_M,
            BLOCK_N=BLOCK_N,
            num_warps=num_warps,
            num_stages=num_stages,
        )


class EvaAggFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, q, k, v, rfa_k, rfa_v, window_mask, chunk_mask, softmax_scale=None, window_size=None, chunks_per_window=None):
        if rfa_k is None and rfa_v is None:
            empty_rfa_kv = 1
        else:
            assert rfa_k is not None and rfa_v is not None, "Both rfa_k and rfa_v must either be None or have values at the same time."
            empty_rfa_kv = 0
        
        if window_mask is not None:
            mask_type = 1
        else:
            mask_type = 0
        o, lse = triton_eva_agg_fwd(
            q, k, v, rfa_k, rfa_v, window_mask, chunk_mask, softmax_scale, window_size, chunks_per_window
        )
        ctx.save_for_backward(q, k, v, o, lse, rfa_k, rfa_v, window_mask, chunk_mask)
        ctx.softmax_scale = softmax_scale
        ctx.window_size = window_size
        ctx.chunks_per_window = chunks_per_window
        ctx.empty_rfa_kv = empty_rfa_kv
        ctx.mask_type = mask_type
        return o

    @staticmethod
    def backward(ctx, do):
        q, k, v, o, lse, rfa_k, rfa_v, window_mask, chunk_mask = ctx.saved_tensors
        dq = torch.empty_like(q)
        dk = torch.empty_like(k)
        dv = torch.empty_like(v)
        if ctx.empty_rfa_kv == 0:
            d_rfa_k = torch.empty_like(rfa_k)
            d_rfa_v = torch.empty_like(rfa_v)
        else:
            d_rfa_k = None
            d_rfa_v = None
        triton_eva_agg_bwd(
            do,
            q,
            k,
            v,
            rfa_k,
            rfa_v,
            window_mask,
            chunk_mask,
            o,
            lse,
            dq,
            dk,
            dv,
            d_rfa_k,
            d_rfa_v,
            softmax_scale=ctx.softmax_scale,
            window_size=ctx.window_size,
            chunks_per_window=ctx.chunks_per_window,
            empty_rfa_kv=ctx.empty_rfa_kv,
            mask_type=ctx.mask_type,
        )
        return dq, dk, dv, d_rfa_k, d_rfa_v, None, None, None, None, None


def eva_agg_func_triton(
        q, k, v, rfa_k, rfa_v, 
        window_mask, chunk_mask, 
        softmax_scale=None, window_size=None, chunks_per_window=None,
    ):
    return EvaAggFunc.apply(
        q, k, v, rfa_k, rfa_v, 
        window_mask, chunk_mask, 
        softmax_scale, window_size, chunks_per_window,
    )