Update model and kernels for training support

eva.py

@@ -2,8 +2,8 @@ from typing import Dict, Optional, Tuple, List, Any, Union
 import torch
 from torch import nn
 import torch.nn.functional as F
-from .eva_agg_kernel import triton_eva_agg_fwd
-from .eva_prep_kv_kernel import triton_eva_prep_kv_fwd
+from .eva_agg_kernel import eva_agg_func_triton
+from .eva_prep_kv_kernel import eva_prep_kv_func_triton
 try:
     import triton
     USE_TRITON_IMPL = True
@@ -129,10 +129,10 @@ class EvaAttention(nn.Module):
         assert not output_attentions
         bsz, q_len, _ = hidden_states.size()
 
-        if use_cache and past_key_value is None:
-            raise ValueError
-
-        assert isinstance(attention_mask, tuple)
+        if use_cache:
+            if past_key_value is None:
+                raise ValueError
+            assert isinstance(attention_mask, tuple)
 
         # infer the model's running mode
         is_prefilling = use_cache and past_key_value.get_seq_length(self.layer_idx) == 0
@@ -141,13 +141,16 @@ class EvaAttention(nn.Module):
         if is_prefilling:
             assert len(attention_mask) == 2
             window_mask, intra_chunk_mask = attention_mask
-            chunk_dummpy_mask = None
+            chunk_mask = None
         elif is_decoding:
             assert len(attention_mask) == 3
-            window_mask, intra_chunk_mask, chunk_dummpy_mask = attention_mask
+            window_mask, intra_chunk_mask, chunk_mask = attention_mask
         else:
-            window_mask, intra_chunk_mask = attention_mask
-            chunk_dummpy_mask = None
+            if attention_mask is not None:
+                assert isinstance(attention_mask, tuple) and len(attention_mask) == 3
+                window_mask, chunk_mask, intra_chunk_mask = attention_mask
+            else:
+                window_mask, chunk_mask, intra_chunk_mask = None, None, None
 
         ############################################
         # compute q, k, v from hidden states
@@ -201,7 +204,7 @@ class EvaAttention(nn.Module):
             #   k/v: [b, h, w, d]
             #   rfa_k/rfa_v: [b, h, w//c, d]
             # 3. in forward inference; the seq_len is already divisible
-            rfa_k, rfa_v = triton_eva_prep_kv_fwd(
+            rfa_k, rfa_v = eva_prep_kv_func_triton(
                 dump_k, dump_v, 
                 self.adaptive_mu_k, self.adaptive_phi, 
                 dump_rf_mask, self.head_dim_scaling, self.chunk_size
@@ -227,10 +230,11 @@ class EvaAttention(nn.Module):
             #   q: [b, h, n, d]
             #   k/v: [b, h, n, d]
             #   rfa_k/rfa_v: [b, h, n // c, d]
-            attn_output = triton_eva_agg_fwd(
+            attn_output = eva_agg_func_triton(
                 q, s_k, s_v, 
                 rfa_k, rfa_v, 
-                singleton_mask, self.head_dim_scaling, self.window_size, self.chunks_per_window
+                singleton_mask, chunk_mask,
+                self.head_dim_scaling, self.window_size, self.chunks_per_window
             )
         elif is_decoding:
             # 2. in decoding, the input shape is 
@@ -258,8 +262,8 @@ class EvaAttention(nn.Module):
                     agg_k = torch.cat([s_k, rfa_k[..., :num_windows_seen_so_far * self.chunks_per_window, :]], dim=-2)
                     agg_v = torch.cat([s_v, rfa_v[..., :num_windows_seen_so_far * self.chunks_per_window, :]], dim=-2)
                     if singleton_mask is not None:
-                        assert chunk_dummpy_mask is not None
-                        attn_mask = torch.cat([singleton_mask, chunk_dummpy_mask], dim=-1)
+                        assert chunk_mask is not None
+                        attn_mask = torch.cat([singleton_mask, chunk_mask], dim=-1)
                     else:
                         attn_mask = singleton_mask
             else:
@@ -275,10 +279,11 @@ class EvaAttention(nn.Module):
             )
         else:
             # 3. in single-forward inference
-            attn_output = triton_eva_agg_fwd(
+            attn_output = eva_agg_func_triton(
                 q, s_k, s_v, 
                 rfa_k, rfa_v, 
-                singleton_mask, self.head_dim_scaling, self.window_size, self.chunks_per_window
+                singleton_mask, chunk_mask,
+                self.head_dim_scaling, self.window_size, self.chunks_per_window
             )
         if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
             raise ValueError(

eva_agg_kernel.py

@@ -4,15 +4,969 @@ import torch
 import triton
 import triton.language as tl
 
-# Disabling autotune for now, set num_warps=4 if headdim=64 and num_warps=8 if headdim=128
-# @triton.autotune(
-#     configs=[
-#         triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=4, num_stages=1),
-#         # This config has a race condition when EVEN_M == False, disabling it for now.
-#         # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=4, num_stages=1),
-#     ],
-#     key=['CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL', 'BLOCK_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_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,
@@ -30,23 +984,24 @@ def _fwd_eva_agg_kernel(
     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_mb, stride_mm,
+    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,
-    CACHE_KEY_SEQLEN_Q, # TODO: why keeping this
-    CACHE_KEY_SEQLEN_K, # TODO: why keeping this
-    CACHE_KEY_NCHUNKS, # TODO: why keeping this
     CHUNKS_PER_WINDOW: tl.constexpr,
     WINDOW_SIZE: tl.constexpr,
     MASK_TYPE: tl.constexpr,
@@ -106,11 +1061,18 @@ def _fwd_eva_agg_kernel(
     qk_scale = softmax_scale
     qk_scale *= 1.4426950408889634  # log2(e)
     if MASK_TYPE == 1:
-        m_ptrs = (
+        window_mask_ptrs = (
             WindowMask +
-            off_b * stride_mb +
-            (offs_m[:, None] * stride_mm + offs_n[None, :])
+            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)
@@ -181,32 +1143,37 @@ def _fwd_eva_agg_kernel(
 
         if MASK_TYPE == 1:
             if EVEN_M & EVEN_W:
-                mask = tl.load(
-                    m_ptrs + start_n - start_idx_n
-                ).to(tl.float32)
+                window_mask = tl.load(
+                    window_mask_ptrs + start_n - start_idx_n
+                )
             else:
-                mask = tl.load(
-                    m_ptrs + start_n - start_idx_n,
+                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=0.0,
-                ).to(tl.float32)
+                    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 + mask
+            qk = qk * qk_scale
+            qk = tl.where(window_mask, float("-inf"), qk)
             m_ij = tl.maximum(tl.max(qk, 1), m_i)
-            p = tl.exp2(qk - m_ij[:, None])
+            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)
-            p = tl.exp2(qk * qk_scale - m_ij[:, None])
+            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)
+        prev_scale = tl.exp2(m_i - m_ij_masked)
         # # -- update output accumulator --
         acc_o = acc_o * prev_scale[:, None]
         # update acc_o
@@ -278,13 +1245,37 @@ def _fwd_eva_agg_kernel(
             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"))
 
-            m_ij = tl.maximum(tl.max(qk, 1) * qk_scale, m_i)
-            p = tl.exp2(qk * qk_scale - m_ij[:, None])
+            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)
+            prev_scale = tl.exp2(m_i - m_ij_masked)
             # # -- update output accumulator --
             acc_o = acc_o * prev_scale[:, None]
             # update acc_o
@@ -320,7 +1311,10 @@ def _fwd_eva_agg_kernel(
             d_i = d_i * prev_scale + d_ij
             m_i = m_ij
 
-    # BUG: have to store and immediately load
+    # 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)
@@ -353,8 +1347,30 @@ def _fwd_eva_agg_kernel(
                 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, softmax_scale, window_size, chunks_per_window):
+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
 
@@ -394,13 +1410,27 @@ def triton_eva_agg_fwd(q, k, v, rfa_k, rfa_v, window_mask, softmax_scale, window
     mask_type = 0
     if window_mask is not None:
         mask_type = 1
-        assert window_mask.dtype == q.dtype, torch.float
+        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()
-    mask_strides = (
+
+        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)
@@ -416,20 +1446,16 @@ def triton_eva_agg_fwd(q, k, v, rfa_k, rfa_v, window_mask, softmax_scale, window
         if empty_rfa_kv == 0 else
         (0, 0, 0)
     )
-    assert chunks_per_window > 0, "chunks_per_window must be greater than 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)
-    if q.dtype == torch.float:
-        BLOCK = 64
-    else:
-        BLOCK = 128
-    num_warps = 4 if head_dim <= 64 else 8
-    assert chunks_per_window >= BLOCK, "chunks_per_window must be greater than BLOCK" 
-    # WINDOW_MASK_TYPE:
-    # - 0: regular causal mask, simply None
-    # - 1: the shape must be B, 1, W, I, J
+    
+    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](
@@ -439,31 +1465,302 @@ def triton_eva_agg_fwd(q, k, v, rfa_k, rfa_v, window_mask, softmax_scale, window
         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],
-        mask_strides[0], mask_strides[1],
+        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,
-        seqlen_q // 32,
-        seqlen_k // 32,
-        nchunks // 32,
         chunks_per_window,
         window_size,
         mask_type,
         empty_rfa_kv,
         BLOCK_HEADDIM,
-        BLOCK_M=BLOCK,
-        BLOCK_N=BLOCK,
+        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=1,
+        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,
     )
-    return o

eva_prep_kv_kernel.py

@@ -16,7 +16,7 @@ def _fwd_eva_prep_kv_kernel(
     V, # [b, h, n, d]
     PARAM_MU, # [1, h, 1, 1, d]
     PARAM_PHI,  # [1, h, 1, 1, d]
-    ChunkMask, # [b, h, n, 1]
+    Mask, # [b, h, n, 1]
     Out_RFA_K, # [b, h, c, d]
     Out_RFA_V, # [b, h, c, d]
     softmax_scale,
@@ -31,8 +31,6 @@ def _fwd_eva_prep_kv_kernel(
     seqlen,
     nchunks,
     headdim,
-    CACHE_KEY_SEQLEN, # TODO: why keeping this
-    CACHE_KEY_NCHUNKS, # TODO: why keeping this
     CHUNKS_PER_BLOCK: tl.constexpr,
     CHUNK_SIZE: tl.constexpr,
     MASK_TYPE: tl.constexpr,
@@ -91,7 +89,7 @@ def _fwd_eva_prep_kv_kernel(
     log2e = 1.4426950408889634
     if MASK_TYPE == 1:
         m_ptrs = (
-            ChunkMask +
+            Mask +
             offs_b * stride_mb +
             (
                 (
@@ -144,7 +142,7 @@ def _fwd_eva_prep_kv_kernel(
         if EVEN_N:
             mask = tl.load(
                 m_ptrs
-            ).to(tl.float32)
+            )
         else:
             mask = tl.load(
                 m_ptrs,
@@ -153,12 +151,17 @@ def _fwd_eva_prep_kv_kernel(
                         offs_c[:, None] * CHUNK_SIZE + 
                         offs_m[None, :]
                     ) < seqlen,
-                other=0.0,
-            ).to(tl.float32)
-        rfa_k_c_w = rfa_k_c_w + mask
-
-    rfa_k_c_w = tl.exp2(rfa_k_c_w - tl.max(rfa_k_c_w, axis=-1)[:, None])
-    rfa_k_c_w = rfa_k_c_w / tl.sum(rfa_k_c_w, axis=-1)[:, None]
+                other=1,
+            )
+        rfa_k_c_w = tl.where(mask, float("-inf"), rfa_k_c_w)
+        
+    m_rfa_k_c_w = tl.max(rfa_k_c_w, axis=-1)
+    masked_out_rows_rfa_k = (m_rfa_k_c_w == float("-inf"))
+    m_rfa_k_c_w_masked = tl.where(masked_out_rows_rfa_k, 0, m_rfa_k_c_w)
+    rfa_k_c_w = tl.exp2(rfa_k_c_w - m_rfa_k_c_w_masked[:, None])
+    denom_k = tl.sum(rfa_k_c_w, axis=-1)
+    denom_k = tl.where(denom_k == 0.0, 1.0, denom_k)
+    rfa_k_c_w = rfa_k_c_w / denom_k[:, None]
     rfa_k_c = tl.sum(k * rfa_k_c_w[:, :, None].to(k.dtype), axis=-2)
     # TODO: understand why rematerialize offsets to save registers?
     offs_out_c = start_n * CHUNKS_PER_BLOCK + tl.arange(0, CHUNKS_PER_BLOCK)
@@ -209,7 +212,7 @@ def _fwd_eva_prep_kv_kernel(
         )
 
     if MASK_TYPE == 1:
-        rfa_v_c_w = rfa_v_c_w + mask
+        rfa_v_c_w = tl.where(mask, float("-inf"), rfa_v_c_w)
 
     if EVEN_N:
         if EVEN_HEADDIM:
@@ -246,8 +249,14 @@ def _fwd_eva_prep_kv_kernel(
                 other=0.0
             )
     
-    rfa_v_c_w = tl.exp2(rfa_v_c_w - tl.max(rfa_v_c_w, axis=-1)[:, None])
-    rfa_v_c_w = rfa_v_c_w / tl.sum(rfa_v_c_w, axis=-1)[:, None]
+
+    m_rfa_v_c_w = tl.max(rfa_v_c_w, axis=-1)
+    masked_out_rows_rfa_v = (m_rfa_v_c_w == float("-inf"))
+    m_rfa_v_c_w_masked = tl.where(masked_out_rows_rfa_v, 0, m_rfa_v_c_w)
+    rfa_v_c_w = tl.exp2(rfa_v_c_w - m_rfa_v_c_w_masked[:, None])
+    denom_v = tl.sum(rfa_v_c_w, axis=-1)
+    denom_v = tl.where(denom_v == 0.0, 1.0, denom_v)
+    rfa_v_c_w = rfa_v_c_w / denom_v[:, None]
     rfa_v_c = tl.sum(v * rfa_v_c_w[:, :, None].to(v.dtype), axis=-2)
 
     offs_out_c = start_n * CHUNKS_PER_BLOCK + tl.arange(0, CHUNKS_PER_BLOCK)
@@ -279,7 +288,529 @@ def _fwd_eva_prep_kv_kernel(
                 mask=(offs_out_c[:, None] < nchunks) & (offs_d[None, :] < headdim)
             )
 
-def triton_eva_prep_kv_fwd(k, v, param_mu, param_phi, chunk_mask, softmax_scale, chunksize):
+
+
+@triton.heuristics(
+    {
+        "EVEN_N": lambda args: args["seqlen"] % args["BLOCK_N"] == 0,
+        "EVEN_HEADDIM": lambda args: args["headdim"] == args["BLOCK_HEADDIM"],
+    }
+)
+@triton.jit
+def _bwd_eva_prep_kv_kernel(
+    RFA_K, # [b, h, c, d]
+    RFA_V, # [b, h, c, d]
+    K, # [b, h, n, d]
+    V, # [b, h, n, d]
+    PARAM_MU, # [1, h, 1, 1, d]
+    PARAM_PHI,  # [1, h, 1, 1, d]
+    Mask, # [b, h, n, 1]
+    D_RFA_K, # [b, h, c, d]
+    D_RFA_V, # [b, h, c, d]
+    D_K, # [b, h, n, d]
+    D_V, # [b, h, n, d]
+    D_PARAM_MU_PARTIAL, # [b, h, g, d]
+    D_PARAM_PHI_PARTIAL, # [b, h, g, d]
+    softmax_scale,
+    stride_rfa_k_b, stride_rfa_k_h, stride_rfa_k_c,
+    stride_rfa_v_b, stride_rfa_v_h, stride_rfa_v_c,
+    stride_kb, stride_kh, stride_kn,
+    stride_vb, stride_vh, stride_vn,
+    stride_mu_h,
+    stride_phi_h,
+    stride_mb, stride_mn,
+    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,
+    stride_d_k_b, stride_d_k_h, stride_d_k_n,
+    stride_d_v_b, stride_d_v_h, stride_d_v_n,
+    stride_d_mu_b, stride_d_mu_h, stride_d_mu_g,
+    stride_d_phi_b, stride_d_phi_h, stride_d_phi_g,
+    nheads,
+    seqlen,
+    nchunks,
+    headdim,
+    CHUNKS_PER_BLOCK: tl.constexpr,
+    CHUNK_SIZE: tl.constexpr,
+    MASK_TYPE: tl.constexpr,
+    BLOCK_HEADDIM: tl.constexpr,
+    EVEN_N: tl.constexpr,
+    EVEN_HEADDIM: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    start_n = tl.program_id(0)
+    offs_bh = tl.program_id(1)
+    offs_h = offs_bh % nheads
+    offs_b = offs_bh // nheads
+    # initialize offsets
+    # we load BLOCK_N keys and values each time, and
+    # reshape it to [CHUNKS_PER_BLOCK, CHUNK_SIZE]
+    offs_c = tl.arange(0, CHUNKS_PER_BLOCK)
+    offs_m = tl.arange(0, CHUNK_SIZE)
+    offs_d = tl.arange(0, BLOCK_HEADDIM)
+
+    offs_rfa_c = start_n * CHUNKS_PER_BLOCK + offs_c
+
+    k_ptrs = (
+        K +
+        offs_b * stride_kb +
+        offs_h * stride_kh +
+        (
+            (
+                start_n * BLOCK_N + 
+                offs_c[:, None, None] * CHUNK_SIZE + 
+                offs_m[None, :, None]
+            ) * stride_kn + 
+            offs_d[None, None, :]
+        )
+    )
+    rfa_k_ptrs = (
+        RFA_K +
+        offs_b * stride_rfa_k_b +
+        offs_h * stride_rfa_k_h +
+        (offs_rfa_c[:, None] * stride_rfa_k_c + offs_d[None, :])
+    )
+    rfa_v_ptrs = (
+        RFA_V +
+        offs_b * stride_rfa_v_b +
+        offs_h * stride_rfa_v_h +
+        (offs_rfa_c[:, None] * stride_rfa_v_c + offs_d[None, :])
+    )
+
+    d_rfa_k_ptrs = (
+        D_RFA_K +
+        offs_b * stride_d_rfa_k_b +
+        offs_h * stride_d_rfa_k_h +
+        (offs_rfa_c[:, None] * stride_d_rfa_k_c + offs_d[None, :])
+    )
+    d_rfa_v_ptrs = (
+        D_RFA_V +
+        offs_b * stride_d_rfa_v_b +
+        offs_h * stride_d_rfa_v_h +
+        (offs_rfa_c[:, None] * stride_d_rfa_v_c + offs_d[None, :])
+    )
+
+    param_mu_ptrs = (
+        PARAM_MU +
+        offs_h * stride_mu_h +
+        offs_d[None, None, :]
+    )
+    param_phi_ptrs = (
+        PARAM_PHI +
+        offs_h * stride_phi_h +
+        offs_d[None, None, :]
+    )
+    
+    log2e = 1.4426950408889634
+    if MASK_TYPE == 1:
+        m_ptrs = (
+            Mask +
+            offs_b * stride_mb +
+            (
+                (
+                    start_n * BLOCK_N +
+                    offs_c[:, None] * CHUNK_SIZE + 
+                    offs_m[None, :]
+                ) * stride_mn
+            )
+        )
+    if EVEN_N:
+        if EVEN_HEADDIM:
+            k = tl.load(
+                k_ptrs
+            )
+        else:
+            k = tl.load(
+                k_ptrs,
+                mask=offs_d[None, None, :] < headdim,
+                other=0.0
+            )
+    else:
+        if EVEN_HEADDIM:
+            k = tl.load(
+                k_ptrs,
+                mask=(
+                        start_n * BLOCK_N +
+                        offs_c[:, None, None] * CHUNK_SIZE + 
+                        offs_m[None, :, None]
+                    ) < seqlen,
+                other=0.0
+            )
+        else:
+            k = tl.load(
+                k_ptrs,
+                mask=(
+                        (
+                            start_n * BLOCK_N +
+                            offs_c[:, None, None] * CHUNK_SIZE + 
+                            offs_m[None, :, None]
+                        ) < seqlen
+                    ) & (offs_d[None, None, :] < headdim),
+                other=0.0
+            )
+
+    if EVEN_N:
+        if EVEN_HEADDIM:
+            rfa_k = tl.load(
+                rfa_k_ptrs
+            )
+        else:
+            rfa_k = tl.load(
+                rfa_k_ptrs,
+                mask=offs_d[None, :] < headdim,
+                other=0.0
+            )
+    else:
+        if EVEN_HEADDIM:
+            rfa_k = tl.load(
+                rfa_k_ptrs,
+                mask=offs_rfa_c[:, None] < nchunks,
+                other=0.0
+            )
+        else:
+            rfa_k = tl.load(
+                rfa_k_ptrs,
+                mask=(offs_rfa_c[:, None] < nchunks) & (offs_d[None, :] < headdim),
+                other=0.0
+            )
+    
+    if EVEN_N:
+        if EVEN_HEADDIM:
+            d_rfa_k = tl.load(
+                d_rfa_k_ptrs
+            )
+        else:
+            d_rfa_k = tl.load(
+                d_rfa_k_ptrs,
+                mask=offs_d[None, :] < headdim,
+                other=0.0
+            )
+    else:
+        if EVEN_HEADDIM:
+            d_rfa_k = tl.load(
+                d_rfa_k_ptrs,
+                mask=offs_rfa_c[:, None] < nchunks,
+                other=0.0
+            )
+        else:
+            d_rfa_k = tl.load(
+                d_rfa_k_ptrs,
+                mask=(offs_rfa_c[:, None] < nchunks) & (offs_d[None, :] < headdim),
+                other=0.0
+            )
+    
+    param_mu = tl.load(param_mu_ptrs).to(k.dtype)
+    mu_c_w = tl.zeros([CHUNKS_PER_BLOCK, CHUNK_SIZE], dtype=tl.float32)
+    mu_c_w += tl.sum(k * param_mu, axis=-1)
+    mu_c_w *= log2e
+
+    if not EVEN_N:  # Need to mask out otherwise the softmax is wrong
+        mu_c_w += tl.where(
+            (
+                start_n * BLOCK_N +
+                offs_c[:, None] * CHUNK_SIZE + 
+                offs_m[None, :]
+            ) < seqlen, 
+            0, 
+            float("-inf")
+        )
+
+    if MASK_TYPE == 1:
+        if EVEN_N:
+            mask = tl.load(
+                m_ptrs
+            )
+        else:
+            mask = tl.load(
+                m_ptrs,
+                mask=(
+                        start_n * BLOCK_N +
+                        offs_c[:, None] * CHUNK_SIZE + 
+                        offs_m[None, :]
+                    ) < seqlen,
+                other=1,
+            )
+        mu_c_w = tl.where(mask, float("-inf"), mu_c_w)
+
+    # [c, w]
+    m_mu_c_w = tl.max(mu_c_w, axis=-1)
+    masked_out_rows_mu = (m_mu_c_w == float("-inf"))
+    m_mu_c_w_masked = tl.where(masked_out_rows_mu, 0, m_mu_c_w)
+    mu_c_w = tl.exp2(mu_c_w - m_mu_c_w_masked[:, None])
+    denom_mu = tl.sum(mu_c_w, axis=-1)
+    denom_mu = tl.where(denom_mu == 0.0, 1.0, denom_mu)
+    mu_tilde_c_w = mu_c_w / denom_mu[:, None]
+    mu_tilde_c_w = mu_tilde_c_w.to(k.dtype)
+    # [c, d] [c, w, d] -> [c, w]
+    d_mu_tilde_c_w = tl.sum(d_rfa_k[:, None, :] * k, axis=-1)
+    # [c, d] [c, d] -> [c]
+    d_out_rfa_k_t_rfa_k = tl.sum(d_rfa_k * rfa_k, axis=-1)[:, None]
+    d_mu_c_w = (d_mu_tilde_c_w - d_out_rfa_k_t_rfa_k) * mu_tilde_c_w
+
+    # [c, w] [c, w, d] -> [d]
+    d_param_mu = tl.sum(tl.sum(d_mu_c_w[:, :, None] * k, axis=0), axis=0)
+    # [c, w] [c, d] + [c, w] [1, 1, d] -> [c, w, d]
+    d_k = mu_tilde_c_w[:, :, None] * d_rfa_k[:, None, :] + d_mu_c_w[:, :, None] * param_mu
+
+    d_param_mu_partial_ptrs = (
+        D_PARAM_MU_PARTIAL +
+        offs_b * stride_d_mu_b +
+        offs_h * stride_d_mu_h +
+        start_n * stride_d_mu_g +
+        offs_d
+    )
+    if EVEN_HEADDIM:
+        tl.store(
+            d_param_mu_partial_ptrs, d_param_mu
+        )
+    else:
+        tl.store(
+            d_param_mu_partial_ptrs, d_param_mu,
+            mask=offs_d < headdim
+        )
+
+
+    v_ptrs = (
+        V +
+        offs_b * stride_vb +
+        offs_h * stride_vh +
+        (
+            (
+                start_n * BLOCK_N +
+                offs_c[:, None, None] * CHUNK_SIZE + 
+                offs_m[None, :, None]
+            ) * stride_vn + 
+            offs_d[None, None, :]
+        )
+    )
+    if EVEN_N:
+        if EVEN_HEADDIM:
+            v = tl.load(
+                v_ptrs
+            )
+        else:
+            v = tl.load(
+                v_ptrs,
+                mask=offs_d[None, None, :] < headdim,
+                other=0.0
+            )
+    else:
+        if EVEN_HEADDIM:
+            v = tl.load(
+                v_ptrs,
+                mask=(
+                        start_n * BLOCK_N +
+                        offs_c[:, None, None] * CHUNK_SIZE + 
+                        offs_m[None, :, None]
+                    ) < seqlen,
+                other=0.0
+            )
+        else:
+            v = tl.load(
+                v_ptrs,
+                mask=(
+                        (
+                            start_n * BLOCK_N +
+                            offs_c[:, None, None] * CHUNK_SIZE + 
+                            offs_m[None, :, None]
+                        ) < seqlen
+                    ) & (offs_d[None, None, :] < headdim),
+                other=0.0
+            )
+
+
+    if EVEN_N:
+        if EVEN_HEADDIM:
+            rfa_v = tl.load(
+                rfa_v_ptrs
+            )
+        else:
+            rfa_v = tl.load(
+                rfa_v_ptrs,
+                mask=offs_d[None, :] < headdim,
+                other=0.0
+            )
+    else:
+        if EVEN_HEADDIM:
+            rfa_v = tl.load(
+                rfa_v_ptrs,
+                mask=offs_rfa_c[:, None] < nchunks,
+                other=0.0
+            )
+        else:
+            rfa_v = tl.load(
+                rfa_v_ptrs,
+                mask=(offs_rfa_c[:, None] < nchunks) & (offs_d[None, :] < headdim),
+                other=0.0
+            )
+    
+    if EVEN_N:
+        if EVEN_HEADDIM:
+            d_rfa_v = tl.load(
+                d_rfa_v_ptrs
+            )
+        else:
+            d_rfa_v = tl.load(
+                d_rfa_v_ptrs,
+                mask=offs_d[None, :] < headdim,
+                other=0.0
+            )
+    else:
+        if EVEN_HEADDIM:
+            d_rfa_v = tl.load(
+                d_rfa_v_ptrs,
+                mask=offs_rfa_c[:, None] < nchunks,
+                other=0.0
+            )
+        else:
+            d_rfa_v = tl.load(
+                d_rfa_v_ptrs,
+                mask=(offs_rfa_c[:, None] < nchunks) & (offs_d[None, :] < headdim),
+                other=0.0
+            )
+    
+    param_phi = tl.load(param_phi_ptrs).to(k.dtype)
+    phi_c_w = tl.zeros([CHUNKS_PER_BLOCK, CHUNK_SIZE], dtype=tl.float32)
+    phi_c_w += tl.sum(k * param_phi, axis=-1)
+    phi_c_w -= (0.5 * tl.sum(k * k, axis=-1))
+    phi_c_w *= log2e * softmax_scale
+    if not EVEN_N:  # Need to mask out otherwise the softmax is wrong
+        phi_c_w += tl.where(
+            (
+                start_n * BLOCK_N +
+                offs_c[:, None] * CHUNK_SIZE + 
+                offs_m[None, :]
+            ) < seqlen, 
+            0, 
+            float("-inf")
+        )
+
+    if MASK_TYPE == 1:
+        phi_c_w = tl.where(mask, float("-inf"), phi_c_w)
+
+
+    m_phi_c_w = tl.max(phi_c_w, axis=-1)
+    masked_out_rows_phi = (m_phi_c_w == float("-inf"))
+    m_phi_c_w_masked = tl.where(masked_out_rows_phi, 0, m_phi_c_w)
+    phi_c_w = tl.exp2(phi_c_w - m_phi_c_w_masked[:, None])
+    denom_phi = tl.sum(phi_c_w, axis=-1)
+    denom_phi = tl.where(denom_phi == 0.0, 1.0, denom_phi)
+    phi_tilde_c_w = phi_c_w / denom_phi[:, None]
+    # phi_c_w = tl.exp2(phi_c_w - tl.max(phi_c_w, axis=-1)[:, None])
+    # phi_tilde_c_w = phi_c_w / tl.sum(phi_c_w, axis=-1)[:, None]
+    phi_tilde_c_w = phi_tilde_c_w.to(k.dtype)
+    d_phi_tilde_c_w = tl.sum(d_rfa_v[:, None, :] * v, axis=-1)
+    d_out_rfa_v_t_rfa_v = tl.sum(d_rfa_v * rfa_v, axis=-1)[:, None]
+    d_phi_c_w = (d_phi_tilde_c_w.to(tl.float32) - d_out_rfa_v_t_rfa_v.to(tl.float32)) * phi_tilde_c_w
+
+    d_param_phi = tl.sum(tl.sum(d_phi_c_w[:, :, None] * k * softmax_scale, axis=0), axis=0)
+    d_v = phi_tilde_c_w[:, :, None] * d_rfa_v[:, None, :]
+    # [c, w, d] + [c, w] * [1, 1, d] - [c, w, d]
+    d_k = d_k + softmax_scale * d_phi_c_w[:, :, None] * (param_phi - k)
+
+    d_k_ptrs = (
+        D_K +
+        offs_b * stride_d_k_b +
+        offs_h * stride_d_k_h +
+        (
+            (
+                start_n * BLOCK_N +
+                offs_c[:, None, None] * CHUNK_SIZE + 
+                offs_m[None, :, None]
+            ) * stride_d_k_n + 
+            offs_d[None, None, :]
+        )
+    )
+    d_v_ptrs = (
+        D_V +
+        offs_b * stride_d_v_b +
+        offs_h * stride_d_v_h +
+        (
+            (
+                start_n * BLOCK_N +
+                offs_c[:, None, None] * CHUNK_SIZE + 
+                offs_m[None, :, None]
+            ) * stride_d_v_n + 
+            offs_d[None, None, :]
+        )
+    )
+    if EVEN_N:
+        if EVEN_HEADDIM:
+            tl.store(
+                d_k_ptrs, d_k
+            )
+            tl.store(
+                d_v_ptrs, d_v
+            )
+        else:
+            tl.store(
+                d_k_ptrs, d_k,
+                mask=offs_d[None, None, :] < headdim
+            )
+            tl.store(
+                d_v_ptrs, d_v,
+                mask=offs_d[None, None, :] < headdim
+            )
+    else:
+        if EVEN_HEADDIM:
+            tl.store(
+                d_k_ptrs, d_k,
+                mask=(
+                        (
+                            start_n * BLOCK_N +
+                            offs_c[:, None, None] * CHUNK_SIZE + 
+                            offs_m[None, :, None]
+                        ) < seqlen
+                    ),
+            )
+            tl.store(
+                d_v_ptrs, d_v,
+                mask=(
+                        (
+                            start_n * BLOCK_N +
+                            offs_c[:, None, None] * CHUNK_SIZE + 
+                            offs_m[None, :, None]
+                        ) < seqlen
+                    ),
+            )
+        else:
+            tl.store(
+                d_k_ptrs, d_k,
+                mask=(
+                        (
+                            start_n * BLOCK_N +
+                            offs_c[:, None, None] * CHUNK_SIZE + 
+                            offs_m[None, :, None]
+                        ) < seqlen
+                    ) & (offs_d[None, None, :] < headdim),
+            )
+            tl.store(
+                d_v_ptrs, d_v,
+                mask=(
+                        (
+                            start_n * BLOCK_N +
+                            offs_c[:, None, None] * CHUNK_SIZE + 
+                            offs_m[None, :, None]
+                        ) < seqlen
+                    ) & (offs_d[None, None, :] < headdim),
+            )
+    d_param_phi_partial_ptrs = (
+        D_PARAM_PHI_PARTIAL +
+        offs_b * stride_d_phi_b +
+        offs_h * stride_d_phi_h +
+        start_n * stride_d_phi_g +
+        offs_d
+    )
+    if EVEN_HEADDIM:
+        tl.store(
+            d_param_phi_partial_ptrs, d_param_phi
+        )
+    else:
+        tl.store(
+            d_param_phi_partial_ptrs, d_param_phi,
+            mask=offs_d < headdim
+        )
+
+def triton_eva_prep_kv_fwd(k, v, param_mu, param_phi, mask, softmax_scale, chunksize):
     k, v, param_mu, param_phi = [
         x if x.stride(-1) == 1 else x.contiguous() 
         for x in [k, v, param_mu, param_phi]
@@ -300,16 +831,16 @@ def triton_eva_prep_kv_fwd(k, v, param_mu, param_phi, chunk_mask, softmax_scale,
     softmax_scale = softmax_scale or 1.0 / math.sqrt(head_dim)
 
     mask_type = 0
-    if chunk_mask is not None:
+    if mask is not None:
         mask_type = 1
-        assert chunk_mask.dtype == k.dtype
-        assert chunk_mask.is_cuda
-        assert chunk_mask.dim() == 4
-        assert chunk_mask.shape == (batch, 1, seqlen, 1)
-        if chunk_mask.stride(-1) != 1:
-            chunk_mask = chunk_mask.contiguous()
+        assert mask.dtype == torch.bool
+        assert mask.is_cuda
+        assert mask.dim() == 4
+        assert mask.shape == (batch, 1, seqlen, 1)
+        if mask.stride(-1) != 1:
+            mask = mask.contiguous()
     mask_strides = (
-        (chunk_mask.stride(0), chunk_mask.stride(2)) 
+        (mask.stride(0), mask.stride(2)) 
         if mask_type == 1 else 
         (0, 0)
     )
@@ -329,7 +860,7 @@ def triton_eva_prep_kv_fwd(k, v, param_mu, param_phi, chunk_mask, softmax_scale,
         v,
         param_mu,
         param_phi,
-        chunk_mask,
+        mask,
         out_rfa_k,
         out_rfa_v,
         softmax_scale,
@@ -344,8 +875,6 @@ def triton_eva_prep_kv_fwd(k, v, param_mu, param_phi, chunk_mask, softmax_scale,
         seqlen,
         nchunks,
         head_dim,
-        seqlen // 32,
-        nchunks // 32,
         chunks_per_block,
         chunksize,
         mask_type,
@@ -355,3 +884,134 @@ def triton_eva_prep_kv_fwd(k, v, param_mu, param_phi, chunk_mask, softmax_scale,
         num_stages=1,
     )
     return out_rfa_k, out_rfa_v
+
+def triton_eva_prep_kv_bwd(
+        d_rfa_k, d_rfa_v,
+        k, v, param_mu, param_phi, 
+        mask, 
+        rfa_k, rfa_v,
+        d_k, d_v, d_param_mu, d_param_phi,
+        softmax_scale, 
+        mask_type,
+        chunksize
+    ):
+    d_rfa_k, d_rfa_v = [
+        x if x.stride(-1) == 1 else x.contiguous() 
+        for x in [d_rfa_k, d_rfa_v]
+    ]
+
+    # shape constraints
+    batch, nheads, seqlen, head_dim = k.shape
+    assert seqlen % chunksize == 0, "seqlen must be divisible by chunksize"
+    nchunks = seqlen // chunksize
+    softmax_scale = softmax_scale or 1.0 / math.sqrt(head_dim)
+
+    mask_strides = (
+        (mask.stride(0), mask.stride(2))
+        if mask_type == 1 else 
+        (0, 0)
+    )
+
+    BLOCK_HEADDIM = max(triton.next_power_of_2(head_dim), 16)
+    BLOCK = 128
+    num_warps = 4 if head_dim <= 64 else 8
+    
+    assert (BLOCK > chunksize) & (BLOCK % chunksize) == 0, "BLOCK must be divisible by chunksize"
+    chunks_per_block = BLOCK // chunksize
+
+    partial_groups = triton.cdiv(seqlen, BLOCK)
+    d_param_mu_partial = torch.zeros((batch, nheads, partial_groups, head_dim), dtype=torch.float32, device=d_rfa_k.device)
+    d_param_phi_partial = torch.zeros((batch, nheads, partial_groups, head_dim), dtype=torch.float32, device=d_rfa_k.device)
+    grid = lambda META: (partial_groups, batch * nheads)
+    _bwd_eva_prep_kv_kernel[grid](
+        rfa_k, # [b, h, c, d]
+        rfa_v, # [b, h, c, d]
+        k, # [b, h, n, d]
+        v, # [b, h, n, d]
+        param_mu, # [1, h, 1, 1, d]
+        param_phi,  # [1, h, 1, 1, d]
+        mask, # [b, h, n, 1]
+        d_rfa_k, # [b, h, c, d]
+        d_rfa_v, # [b, h, c, d]
+        d_k, # [b, h, n, d]
+        d_v, # [b, h, n, d]
+        d_param_mu_partial, # [b, h, g, d]
+        d_param_phi_partial, # [b, h, g, d]
+        softmax_scale,
+        rfa_k.stride(0), rfa_k.stride(1), rfa_k.stride(2),
+        rfa_v.stride(0), rfa_v.stride(1), rfa_v.stride(2),
+        k.stride(0), k.stride(1), k.stride(2),
+        v.stride(0), v.stride(1), v.stride(2),
+        param_mu.stride(1),
+        param_phi.stride(1),
+        mask_strides[0], mask_strides[1],
+        d_rfa_k.stride(0), d_rfa_k.stride(1), d_rfa_k.stride(2),
+        d_rfa_v.stride(0), d_rfa_v.stride(1), d_rfa_v.stride(2),
+        d_k.stride(0), d_k.stride(1), d_k.stride(2),
+        d_v.stride(0), d_v.stride(1), d_v.stride(2),
+        d_param_mu_partial.stride(0), d_param_mu_partial.stride(1), d_param_mu_partial.stride(2),
+        d_param_phi_partial.stride(0), d_param_phi_partial.stride(1), d_param_phi_partial.stride(2),
+        nheads,
+        seqlen,
+        nchunks,
+        head_dim,
+        chunks_per_block,
+        chunksize,
+        mask_type,
+        BLOCK_HEADDIM,
+        BLOCK_N=BLOCK,
+        num_warps=num_warps,
+        num_stages=1,
+    )
+    d_param_mu.copy_(d_param_mu_partial.sum(dim=(0, -2), keepdim=True).unsqueeze(-2).to(d_param_mu.dtype))
+    d_param_phi.copy_(d_param_phi_partial.sum(dim=(0, -2), keepdim=True).unsqueeze(-2).to(d_param_phi.dtype))
+
+
+
+class EvaPrepKVFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, k, v, param_mu, param_phi, mask, softmax_scale=None, chunksize=None):
+        if mask is not None:
+            mask_type = 1
+        else:
+            mask_type = 0
+        rfa_k, rfa_v = triton_eva_prep_kv_fwd(
+            k, v, param_mu, param_phi, mask, softmax_scale, chunksize
+        )
+        ctx.save_for_backward(k, v, param_mu, param_phi, mask, rfa_k, rfa_v)
+        ctx.softmax_scale = softmax_scale
+        ctx.chunksize = chunksize
+        ctx.mask_type = mask_type
+        return rfa_k, rfa_v
+
+    @staticmethod
+    def backward(ctx, d_rfa_k, d_rfa_v):
+        k, v, param_mu, param_phi, mask, rfa_k, rfa_v = ctx.saved_tensors
+        d_k = torch.empty_like(k)
+        d_v = torch.empty_like(v)
+        d_param_mu = torch.empty_like(param_mu)
+        d_param_phi = torch.empty_like(param_phi)
+        triton_eva_prep_kv_bwd(
+            d_rfa_k, d_rfa_v,
+            k, v, param_mu, param_phi, 
+            mask, 
+            rfa_k, rfa_v,
+            d_k, d_v, d_param_mu, d_param_phi,
+            ctx.softmax_scale, 
+            ctx.mask_type,
+            ctx.chunksize
+        )
+        return d_k, d_v, d_param_mu, d_param_phi, None, None, None
+
+def eva_prep_kv_func_triton(
+        k, v, 
+        param_mu, param_phi,
+        mask, 
+        softmax_scale=None, chunksize=None
+    ):
+    return EvaPrepKVFunc.apply(
+        k, v, 
+        param_mu, param_phi,
+        mask, 
+        softmax_scale, chunksize
+    )

eva_pt_ref.py

@@ -263,7 +263,6 @@ class EvaAttention(nn.Module):
                 v, 
                 self.layer_idx,
                 self.window_size,
-                self.singleton_update
             )
         else:
             prev_w_q = self.window_partition(q) # [b, h, w, i, d]
@@ -289,10 +288,9 @@ class EvaAttention(nn.Module):
                 layer_idx=self.layer_idx,
                 window_size=self.window_size,
                 chunk_size=self.chunk_size,
-                singleton_update=self.singleton_update
             )
         else:
-            prev_s_mask = window_causal_mask # [1, 1, w, i, j]
+            prev_s_mask = self.window_partition(prev_causal_mask) # [1, 1, w, i, j]
             cur_s_mask = None
             prev_chunk_mask = self.window_partition(chunk_causal_mask)
             cur_chunk_mask = None

modeling_evabyte.py

@@ -148,7 +148,7 @@ class EvaByteRMSNorm(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.config = config
-        self.fp32_ln = config.fp32_ln
+        self.fp32_ln = True
         self.variance_epsilon = config.rms_norm_eps
         self.add_unit_offset = config.norm_add_unit_offset
         if self.add_unit_offset:
@@ -157,18 +157,14 @@ class EvaByteRMSNorm(nn.Module):
             self.weight = nn.Parameter(torch.ones(config.hidden_size))
 
     def forward(self, hidden_states):
-        if hasattr(self, 'config'):
-            fp32_ln = self.config.fp32_ln
-        else:
-            fp32_ln = self.fp32_ln
-        hidden_states = hidden_states.to(torch.float32 if fp32_ln else torch.bfloat16)
+        _hidden_states = hidden_states.to(torch.float32 if self.fp32_ln else torch.bfloat16)
 
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        variance = _hidden_states.pow(2).mean(-1, keepdim=True)
+        _hidden_states = _hidden_states * torch.rsqrt(variance + self.variance_epsilon)
         if self.add_unit_offset:
-            return (1 + self.weight) * hidden_states
+            return ((1 + self.weight) * _hidden_states).type_as(hidden_states)
         else:
-            return self.weight * hidden_states
+            return (self.weight * _hidden_states).type_as(hidden_states)
 
 class EvaByteRotaryEmbedding(torch.nn.Module):
     def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
@@ -313,7 +309,7 @@ class EvaByteDecoderLayer(nn.Module):
                                                                             cos=cos,
                                                                             sin=sin,
                                                                             multibyte_decoding=multibyte_decoding)
-        hidden_states = residual + hidden_states
+        hidden_states = (residual + hidden_states).to(hidden_states.dtype)
 
         # Fully Connected
         residual = hidden_states
@@ -321,7 +317,7 @@ class EvaByteDecoderLayer(nn.Module):
             residual = residual.float()
         hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + hidden_states
+        hidden_states = (residual + hidden_states).to(hidden_states.dtype)
 
         outputs = (hidden_states, )
 
@@ -653,7 +649,7 @@ class EvaByteModel(EvaBytePreTrainedModel):
                     )
             else:
                 assert self.training
-                assert seq_len % self.config.window_size == 0
+                assert seq_len % self.config.window_size == 0, "Training is only tested for sequences that are a multiple of window_size"
                 # for training, we need to pass in the attention mask
                 # usually calculated by _prepare_training_attn_mask()
                 causal_mask = attention_mask
@@ -683,31 +679,6 @@ class EvaByteModel(EvaBytePreTrainedModel):
         cos = cos.unsqueeze(1)
         sin = sin.unsqueeze(1)
 
-        if USE_TRITON_IMPL and (not multibyte_decoding):
-            # the masks generated above for triton kernels are boolean. Convert them to floats
-            if (
-                (not use_cache) or
-                (use_cache and past_seen_tokens == 0)
-            ):
-                window_mask, intra_chunk_mask = causal_mask
-                
-                if window_mask is not None:
-                    assert window_mask.dtype == torch.bool
-                    window_mask_float = window_mask.to(torch.float)
-                    window_mask_float = window_mask_float.masked_fill(window_mask.to(torch.bool), MASK_MIN_VALUE)
-                    window_mask_float = window_mask_float.reshape(batch_size, 1, -1, self.config.window_size)
-                    window_mask = window_mask_float.to(hidden_states.dtype)
-
-                if intra_chunk_mask is not None:
-                    assert intra_chunk_mask.dtype == torch.bool
-                    intra_chunk_mask_float = intra_chunk_mask.to(torch.float)
-                    intra_chunk_mask_float = intra_chunk_mask_float.masked_fill(intra_chunk_mask.to(torch.bool), MASK_MIN_VALUE)
-                    intra_chunk_mask = intra_chunk_mask_float.to(hidden_states.dtype)
-                causal_mask = (window_mask, intra_chunk_mask)
-
-        if self.config.fp32_skip_add:
-            hidden_states = hidden_states.float()
-
         # decoder layers
         all_hidden_states = () if output_hidden_states else None
         all_self_attns = () if output_attentions else None
@@ -718,20 +689,17 @@ class EvaByteModel(EvaBytePreTrainedModel):
                 all_hidden_states += (hidden_states, )
 
             if self.gradient_checkpointing and self.training:
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        # None for past_key_value
-                        return module(*inputs, output_attentions, use_cache=None)
-
-                    return custom_forward
-
                 layer_outputs = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(decoder_layer),
+                    decoder_layer.__call__,
                     hidden_states,
                     causal_mask,
                     position_ids,
-                    None,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    cos,
+                    sin,
+                    multibyte_decoding,
                 )
             else:
                 layer_outputs = decoder_layer(
@@ -806,154 +774,6 @@ class EvaByteForCausalLM(EvaBytePreTrainedModel, MultiByteDecodingMixin):
     def get_decoder(self):
         return self.model
 
-    def _prepare_training_attn_mask(
-        self, 
-        target_token_type_ids,
-        use_doc_boundary_attention, 
-        EOS_TOKEN_TYPE_ID=None,
-        PAD_TOKEN_TYPE_ID=None,
-    ):
-        '''
-        This function prepares the attention mask for training byte models.
-            target_token_type_ids:
-                Tensor of shape (batch_size, seq_len), marking the token type ids 
-                for the target sequence. In particular, we should have
-                    - target_token_type_ids[i, j] = EOS_TOKEN_TYPE_ID 
-                        if the j-th token in the i-th sequence is the end of an article.
-                    - target_token_type_ids[i, j] = PAD_TOKEN_TYPE_ID 
-                        if the j-th token in the i-th sequence is the padding token.
-            use_doc_boundary_attention: bool, 
-                whether to enable doc boundary attention.
-            EOS_TOKEN_TYPE_ID: int, 
-                the token type id for the end of an article.
-            PAD_TOKEN_TYPE_ID: int, 
-                the token type id for the padding token.
-        '''
-        assert self.training
-        batch_size, num_tokens = target_token_type_ids.shape
-
-        chunk_causal_mask, window_causal_mask = prepare_eva_attention_mask(
-            num_tokens, 
-            target_token_type_ids.device, 
-            chunk_size=self.config.chunk_size, 
-            window_size=self.config.window_size,
-            use_cache=False,
-            cache=None
-        )
-        if use_doc_boundary_attention:
-            #### step 1: mark each document with a unique id
-            end_token_ids = {EOS_TOKEN_TYPE_ID, PAD_TOKEN_TYPE_ID}
-            token_types = torch.zeros(batch_size, num_tokens)
-            for sequence_idx, sequence in enumerate(target_token_type_ids):
-                num_articles = 0
-                start_index = 0
-                # for each sample in the batch, the collapsed attention mask looks like:
-                # [1, 1, .... 1, 0, 2, 2, ... 2, 0, ... n, n ..... n], assuming there are n articles in the sequence.
-                # Each of the n articles are separated by 0.
-                for token_idx, token_type_id in enumerate(sequence):
-                    if start_index is not None and token_type_id.item() in end_token_ids:
-                        num_articles += 1
-                        end_index = token_idx if token_type_id == PAD_TOKEN_TYPE_ID else token_idx + 1
-                        token_types[sequence_idx][start_index:end_index] = num_articles
-                        start_index = None
-                    elif start_index is None and token_type_id not in end_token_ids:
-                        start_index = token_idx + 1
-
-            assert num_tokens % self.config.chunk_size == 0, "Number of tokens must be divisible by chunk size"
-            assert num_tokens % self.config.window_size == 0, "Number of tokens must be divisible by window size"
-            num_chunks = num_tokens // self.config.chunk_size
-            num_windows = num_tokens // self.config.window_size
-
-            article_separator = 0
-
-            #### step 2: generate attention masks for each window
-            #### NOTE: we perform exact attention within each window, 
-            ####       so we only need to mask out different documents
-            ####       for each window.
-            token_types_windows = token_types.reshape(batch_size, num_windows, self.config.window_size, 1)
-            token_types_windows_t = token_types_windows.transpose(-1, -2)
-            # replace all elements in TOKEN_SEPS with -1
-            token_types_windows = torch.where(token_types_windows == article_separator, -1, token_types_windows)
-            window_3d_mask = (token_types_windows == token_types_windows_t)
-            window_3d_mask = ~window_3d_mask
-
-            #### step 3: generate chunk-level 3D masks
-            #### NOTE: this is a bit tricky, as we aim to mask out different 
-            ####       documents to avoid cross-doc attention across chunks.
-            #### Example: suppose we have a sequence of length 12 with 3 documents:
-            ####       [1, 1, 1, 1, 1, 2, 2, 3, 3, 3, 3, 3].
-            ####       The chunk-size and window-size are both 4.
-            ####       The chunk-level mask of shape (batch_size, seq_len, num_chunks) is:
-            ####       [
-            ####         [0, 0, 0],
-            ####         [0, 0, 0],
-            ####         [0, 0, 0],
-            ####         [0, 0, 0],
-            ####
-            ####         [1, 0, 0],
-            ####         [0, 0, 0],
-            ####         [0, 0, 0],
-            ####         [0, 0, 0],
-            ####
-            ####         [0, 1, 0],
-            ####         [0, 1, 0],
-            ####         [0, 1, 0],
-            ####         [0, 1, 0],
-            ####       ]
-            ####       Explanation:
-            ####       - Tokens will not attend to their own and future chunks.
-            ####         (as tokens within a chunk are captured by the window-level exact attention)
-            ####       - Tokens will attend to a chunk only if there are tokens 
-            ####         from the same document in that chunk.
-            ####       The mask within each chunk of shape (batch_size, num_chunks, chunk_size) is:
-            ####       [
-            ####         [1, 1, 1, 1],
-            ####         [0, 0, 0, 1],
-            ####         [1, 1, 1, 1],
-            ####       ]
-            ####       Explanation:
-            ####       - If all tokens in a chunk are from the same document, 
-            ####         no tokens will be masked out.
-            ####       - If there are tokens from different documents in a chunk, 
-            ####         only tokens from the rightmost document will be kept.
-            ####         (b/c the future chunks might contain tokens from the rightmost document,
-            ####         but all the remaining docs will never get attended by other docs)
-            token_types_chunks = token_types.reshape(batch_size, num_chunks, self.config.chunk_size)
-            inter_chunk_mask = torch.zeros((batch_size, num_tokens, num_chunks), dtype=torch.bool)
-            intra_chunk_mask = torch.ones_like(token_types_chunks, dtype=torch.bool)
-            
-            for chunk_idx in range(num_chunks):
-                for batch_idx in range(batch_size):
-                    # Identify tokens in the current chunk belonging to each sequence
-                    chunk = token_types_chunks[batch_idx, chunk_idx]
-                    unique_elements = torch.unique(chunk, sorted=True).tolist()
-                    
-                    # Create a mask for whether each token can attend to the current chunk
-                    for token_type in unique_elements:
-                        if token_type == article_separator:
-                            continue
-                        token_mask = (token_types[batch_idx] == token_type)
-                        inter_chunk_mask[batch_idx, :, chunk_idx] |= token_mask
-
-                    # Create a mask within each chunk
-                    unique_elements = [x for x in unique_elements if x != article_separator]
-                    if len(unique_elements) > 1 and chunk[-1] != article_separator:
-                        intra_chunk_mask[batch_idx, chunk_idx] = (chunk == unique_elements[-1])
-            
-            inter_chunk_mask = ~inter_chunk_mask
-            intra_chunk_mask = ~intra_chunk_mask
-
-            window_mask = torch.logical_or(window_causal_mask, window_3d_mask.unsqueeze(1))
-            inter_chunk_mask = torch.logical_or(chunk_causal_mask, inter_chunk_mask.unsqueeze(1))
-            intra_chunk_mask = intra_chunk_mask.unsqueeze(1).unsqueeze(-1)
-
-            joint_mask = torch.cat([window_mask, inter_chunk_mask.reshape(*window_mask.shape)], dim=-1)
-            attention_mask = (joint_mask, intra_chunk_mask)
-        else:
-            joint_mask = torch.cat([window_causal_mask, chunk_causal_mask.reshape(*window_causal_mask.shape)], dim=-1)
-            attention_mask = (joint_mask, None)
-        return attention_mask
-
     def forward(
             self,
             input_ids: torch.LongTensor = None,