Add files using upload-large-folder tool

fla2/ops/mask_gated_delta_rule_t/wy_fast.py

@@ -0,0 +1,541 @@
+# -*- coding: utf-8 -*-
+import pdb
+import torch
+import triton
+import triton.language as tl
+from einops import rearrange
+# from ...utils import autocast_custom_bwd, autocast_custom_fwd, contiguous
+from ...utils import autocast_custom_bwd, autocast_custom_fwd, contiguous
+# Inspired by "THE WY REPRESENTATION FOR PRODUCTS OF HOUSEHOLDER MATRICES" https://epubs.siam.org/doi/pdf/10.1137/0908009
+# o: cumprod
+# o2: cumprodsum
+from typing import Optional
+@triton.jit
+def safe_exp(x):
+    return tl.exp(tl.where(x <= 0, x, float('-inf')))
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16)
+    ],
+    key=["BT", "BK", "BV"],
+)
+@triton.jit
+def gated_fwd_recompute_w_u_kernel(
+    k,
+    v,
+    beta,
+    mask_ij,
+    w,
+    u,
+    Aw,
+    Au,
+    s_qk_h,
+    s_qk_t,
+    s_qk_d,
+    s_vo_h,
+    s_vo_t,
+    s_vo_d,
+    T,
+    K,
+    V,
+    r: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    dk = K//r
+    p_beta = tl.make_block_ptr(beta + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+    b_beta = tl.load(p_beta, boundary_check=(0,))
+    p_Aw = tl.make_block_ptr(Aw + i_bh*T*BT*r*r ,(T*r,BT*r), (BT*r,1), (i_t*BT*r,0), (BT*r,BT*r),(1,0))
+    b_Aw = tl.load(p_Aw, boundary_check=(0, 1)).to(k.dtype.element_ty)
+    for i_r in range(r):
+        p_mask = tl.make_block_ptr(mask_ij + i_bh * T*r*r,(T,r,r),(r*r,r,1),(i_t*BT,0,i_r),(BT,r,1),(2,1,0))
+        b_mask = tl.load(p_mask)#BT r 1
+        for i_k in range(tl.cdiv(dk, BK)):
+            p_k = tl.make_block_ptr(k + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_r*dk + i_k * BK), (BT, BK), (1, 0))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            b_kb = (b_k * b_beta[:, None]).to(b_k.dtype)[:,None,:]*b_mask.to(b_k.dtype)#BT*r*d
+            b_kb = tl.reshape(b_kb,(BT*r,BK))
+            b_w = tl.dot(b_Aw, b_kb, allow_tf32=False)#get BT*r *BK
+            p_w = tl.make_block_ptr(w + i_bh * s_qk_h*r, (T*r, K), (s_qk_t, s_qk_d), (i_t * BT * r, i_r*dk + i_k * BK), (BT*r, BK), (1, 0))
+            tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
+    tl.debug_barrier()
+    b_Aw = None
+    p_Au = tl.make_block_ptr(Au + i_bh*T*BT*r*r ,(T*r,BT*r), (BT*r,1), (i_t*BT*r,0), (BT*r,BT*r),(1,0))
+    b_Au = tl.load(p_Au, boundary_check=(0, 1)).to(k.dtype.element_ty)
+
+    for i_v in range(tl.cdiv(V, BV)):#no need for 任意mask不使用 #无需for 循环 ，这里也不存在mask
+        p_v = tl.make_block_ptr(v + i_bh * s_vo_h, (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        b_vb = (b_v * b_beta[:, None]).to(b_v.dtype)[:,None,:]*tl.full([r],1, dtype=b_v.dtype)[None,:,None]
+        b_vb = tl.reshape(b_vb,(BT*r,BV))
+        b_u = tl.dot(b_Au, b_vb, allow_tf32=False)
+        p_u = tl.make_block_ptr(u + i_bh * s_vo_h*r, (T*r, V), (s_vo_t, s_vo_d), (i_t * BT*r, i_v * BV), (BT*r, BV), (1, 0))
+        tl.store(p_u, (b_u).to(p_u.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16)
+    ],
+    key=["BT", "BK","r"],
+)
+@triton.jit
+def gated_chunk_scaled_dot_kkt_fwd_kernel(        
+    k,
+    beta,
+    g_cumsum,
+    mask_ij,
+    A,
+    Ag,
+    s_qk_h,
+    s_qk_t,
+    s_qk_d,
+    T,
+    K,
+    r:  tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    b_A = tl.zeros([BT,BT,r,r], dtype=tl.float32)#r*BT r*BT
+    dk = K//r
+    p_beta = tl.make_block_ptr(beta + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+    b_beta = tl.load(p_beta, boundary_check=(0,))
+    for i_r in range(r):
+        r_mask = tl.arange(0, r) == i_r 
+        p_mask = tl.make_block_ptr(mask_ij + i_bh * T*r*r,(T,r,r),(r*r,r,1),(i_t*BT,0,i_r),(BT,r,1),(2,1,0))
+        b_mask = tl.load(p_mask)#BT r 1
+        ij_mask = b_mask*r_mask[None,None,:]#行数 #BT [r,r]
+
+        for i_k in range(tl.cdiv(dk, BK)):#分块k读取计算
+            p_k = tl.make_block_ptr(k + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_r * dk + i_k * BK), (BT, BK), (1, 0))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            b_kb = (b_k * b_beta[:, None]).to(b_k.dtype)
+            dot = tl.dot(b_kb, tl.trans(b_k), allow_tf32=False)#BT BT 
+            b_A += dot[:,:,None,None]*ij_mask[:,None,:,:]#BT r r
+
+    b_A = tl.where((tl.arange(0, BT)[:,None] > tl.arange(0, BT)[None,:])[:,:,None,None], b_A, 0)
+    p_A = tl.make_block_ptr(A + (i_bh*T//BT+i_t)*BT*BT*r*r ,(BT,BT,r,r), (BT*r*r,r*r,r,1), (0,0,0,0), (BT,BT,r,r),(3,2,1,0))
+    tl.store(p_A, (b_A).to(p_A.dtype.element_ty),boundary_check=(0,1,2,3))
+
+    p_g = tl.make_block_ptr(g_cumsum + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+    b_g = tl.load(p_g, boundary_check=(0,))
+    b_g_diff = b_g[:, None] - b_g[None, :]
+    b_g_diff = safe_exp(b_g_diff)
+
+    b_Ag = b_A * ((b_g_diff)[:,:,None,None])#BT BT
+    p_Ag = tl.make_block_ptr(Ag + (i_bh*T//BT+i_t)*BT*BT*r*r ,(BT,BT,r,r), (BT*r*r,r*r,r,1), (0,0,0,0), (BT,BT,r,r),(3,2,1,0))
+    tl.store(p_Ag, (b_Ag).to(p_Ag.dtype.element_ty),boundary_check=(0,1,2,3))
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16)
+    ],
+    key=["BT", "r"],
+)
+@triton.jit
+def solve_tril_16x16_kernel(
+    A,
+    Ad,
+    s_A_bh,
+    s_Ad_bh,
+    T,
+    r:  tl.constexpr,
+    BT: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    offset = (i_t * 16) % BT 
+
+    p_A = tl.make_block_ptr(A + (i_bh)*s_A_bh, (T,BT,r,r),(BT*r*r,r*r,r,1) ,(i_t * 16, offset, 0, 0), (16, 16,r,r), (3,2,1,0))
+    b_A = tl.load(p_A, boundary_check=(0,1,2,3)).to(tl.float32)
+    b_A = -tl.where((tl.arange(0, 16)[:,None] > tl.arange(0, 16)[None,:])[:,:,None,None], b_A, 0)
+
+    for i in range(1, 16):
+        mask = tl.arange(0, 16) == i 
+        b_a = tl.sum(tl.where(mask[:,None,None,None], b_A, 0), 0)
+        q = (tl.sum(b_a[:,None,:,:,None]*b_A[:,:,None,:,:],-2))
+        b_a = b_a + tl.sum(q,0)*((tl.arange(0, 16) < i)[:,None,None])
+        b_A = tl.where(mask[:,None,None,None],b_a,b_A)#按行计算 ，逐步交换结果
+    b_A += ((tl.arange(0, 16)[:, None, None, None] == tl.arange(0, 16)[None, :, None, None])&(tl.arange(0, r)[None, None, :, None] == tl.arange(0, r)[None, None, None, :]))
+    
+    b_A = tl.permute(b_A,(0,2,1,3))
+    b_A = tl.reshape(b_A,(16*r,16*r))#BT*r BT*r
+    p_Ad = tl.make_block_ptr(Ad + (i_bh)*s_Ad_bh,(T*r,16*r),(16*r,1), (i_t * 16 * r, 0), (16*r,16*r), (1,0))
+    tl.store(p_Ad, (b_A).to(p_Ad.dtype.element_ty),boundary_check=(0,1))
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16)
+    ],
+    key=["r"],
+)
+@triton.jit
+def merge_16x16_to_32x32_inverse_kernel(
+        A,
+        Ad,
+        Ai,
+        s_A_bh,
+        s_Ad_bh,
+        T,
+        r: tl.constexpr,
+        BT: tl.constexpr 
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+
+    p_A21 = tl.make_block_ptr(A + (i_bh)*s_A_bh, (T*r,32*r),(32*r,1) ,((i_t * 32 + 16) *r, 0), (16*r, 16*r), (1,0))
+    b_A21 = tl.load(p_A21, boundary_check=(0,1)).to(tl.float32)
+
+    p_Ad11  = tl.make_block_ptr(Ad + (i_bh)*s_Ad_bh,(T*r,16*r),(16*r,1), (i_t * 32 * r, 0), (16*r,16*r), (1,0))
+    p_Ad22  = tl.make_block_ptr(Ad + (i_bh)*s_Ad_bh,(T*r,16*r),(16*r,1), ((i_t *32 +16) * r, 0), (16*r,16*r), (1,0))
+
+    p_Ai11 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,32*r), (32*r, 1), (i_t * 32 * r , 0), (16*r, 16*r), (1, 0))
+    p_Ai22 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,32*r), (32*r, 1), ((i_t * 32 + 16) * r , 16*r), (16*r, 16*r), (1, 0))
+    p_Ai21 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,32*r), (32*r, 1), ((i_t * 32 + 16) * r, 0), (16*r, 16*r), (1, 0))
+
+    Ai11 = tl.load(p_Ad11, boundary_check=(0, 1)).to(tl.float32)
+    Ai22 = tl.load(p_Ad22, boundary_check=(0, 1)).to(tl.float32)
+    Ai21 = -tl.dot(tl.dot(Ai22,b_A21, input_precision='ieee'),Ai11,input_precision='ieee')
+    tl.store(p_Ai11,Ai11.to(p_Ai11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai22,Ai22.to(p_Ai22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai21,Ai21.to(p_Ai21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16)
+    ],
+    key=["r"],
+)
+@triton.jit
+def merge_16x16_to_64x64_inverse_kernel(
+        A,
+        Ad,
+        Ai,
+        s_A_bh,
+        s_Ad_bh,
+        T,
+        r: tl.constexpr,
+        BT: tl.constexpr 
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+
+    p_A21 = tl.make_block_ptr(A + (i_bh)*s_A_bh, (T*r,64*r),(64*r,1) ,((i_t * 64 + 16) *r, 0), (16*r, 16*r), (1,0))
+    p_A31 = tl.make_block_ptr(A + (i_bh)*s_A_bh, (T*r,64*r),(64*r,1) ,((i_t * 64 + 32) *r, 0), (16*r, 16*r), (1,0))
+    p_A32 = tl.make_block_ptr(A + (i_bh)*s_A_bh, (T*r,64*r),(64*r,1) ,((i_t * 64 + 32) *r, 16*r), (16*r, 16*r), (1,0))
+    p_A41 = tl.make_block_ptr(A + (i_bh)*s_A_bh, (T*r,64*r),(64*r,1) ,((i_t * 64 + 48) *r, 0), (16*r, 16*r), (1,0))
+    p_A42 = tl.make_block_ptr(A + (i_bh)*s_A_bh, (T*r,64*r),(64*r,1) ,((i_t * 64 + 48) *r, 16*r), (16*r, 16*r), (1,0))
+    p_A43 = tl.make_block_ptr(A + (i_bh)*s_A_bh, (T*r,64*r),(64*r,1) ,((i_t * 64 + 48) *r, 32*r), (16*r, 16*r), (1,0))
+    
+    b_A21 = tl.load(p_A21, boundary_check=(0,1)).to(tl.float32)
+    b_A31 = tl.load(p_A31, boundary_check=(0,1)).to(tl.float32)
+    b_A32 = tl.load(p_A32, boundary_check=(0,1)).to(tl.float32)
+    b_A41 = tl.load(p_A41, boundary_check=(0,1)).to(tl.float32)
+    b_A42 = tl.load(p_A42, boundary_check=(0,1)).to(tl.float32)
+    b_A43 = tl.load(p_A43, boundary_check=(0,1)).to(tl.float32)
+
+
+    p_Ad11  = tl.make_block_ptr(Ad + (i_bh)*s_Ad_bh,(T*r,16*r),(16*r,1), (i_t * 64 * r, 0), (16*r,16*r), (1,0))
+    p_Ad22  = tl.make_block_ptr(Ad + (i_bh)*s_Ad_bh,(T*r,16*r),(16*r,1), ((i_t * 64 + 16) * r, 0), (16*r,16*r), (1,0))
+    p_Ad33  = tl.make_block_ptr(Ad + (i_bh)*s_Ad_bh,(T*r,16*r),(16*r,1), ((i_t * 64 + 32) * r, 0), (16*r,16*r), (1,0))
+    p_Ad44  = tl.make_block_ptr(Ad + (i_bh)*s_Ad_bh,(T*r,16*r),(16*r,1), ((i_t * 64 + 48) * r, 0), (16*r,16*r), (1,0))
+
+
+    p_Ai11 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,64*r), (64*r, 1), ((i_t * 64 ) *r, 0), (16*r, 16*r), (1, 0))
+    p_Ai22 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,64*r), (64*r, 1), ((i_t * 64 + 16) *r, 16*r), (16*r, 16*r), (1, 0))
+    p_Ai33 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,64*r), (64*r, 1), ((i_t * 64 + 32) *r, 32*r), (16*r, 16*r), (1, 0))
+    p_Ai44 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,64*r), (64*r, 1), ((i_t * 64 + 48) *r, 48*r), (16*r, 16*r), (1, 0))
+    p_Ai21 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,64*r), (64*r, 1), ((i_t * 64 + 16) *r, 0), (16*r, 16*r), (1, 0))
+    p_Ai31 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,64*r), (64*r, 1), ((i_t * 64 + 32) *r, 0), (16*r, 16*r), (1, 0))
+    p_Ai32 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,64*r), (64*r, 1), ((i_t * 64 + 32) *r, 16*r), (16*r, 16*r), (1, 0))
+    p_Ai41 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,64*r), (64*r, 1), ((i_t * 64 + 48) *r ,0), (16*r, 16*r), (1, 0))
+    p_Ai42 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,64*r), (64*r, 1), ((i_t * 64 + 48) *r, 16*r), (16*r, 16*r), (1, 0))
+    p_Ai43 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,64*r), (64*r, 1), ((i_t * 64 + 48) *r, 32*r), (16*r, 16*r), (1, 0))
+
+
+    Ai11 = tl.load(p_Ad11, boundary_check=(0, 1)).to(tl.float32)
+    Ai22 = tl.load(p_Ad22, boundary_check=(0, 1)).to(tl.float32)
+    Ai33 = tl.load(p_Ad33, boundary_check=(0, 1)).to(tl.float32)
+    Ai44 = tl.load(p_Ad44, boundary_check=(0, 1)).to(tl.float32)
+    
+    Ai21 = -tl.dot(tl.dot(Ai22,b_A21, input_precision='ieee'),Ai11,input_precision='ieee')
+    Ai32 = -tl.dot(tl.dot(Ai33,b_A32, input_precision='ieee'),Ai11,input_precision='ieee')
+    Ai43 = -tl.dot(tl.dot(Ai44,b_A43, input_precision='ieee'),Ai11,input_precision='ieee')
+
+    Ai31 = -tl.dot(
+            Ai33,
+            tl.dot(b_A31,Ai11, input_precision='ieee')+
+            tl.dot(b_A32,Ai21, input_precision='ieee'),
+            input_precision='ieee')
+
+    Ai42 = -tl.dot(
+            Ai44,
+            tl.dot(b_A42,Ai22, input_precision='ieee')+
+            tl.dot(b_A43,Ai32, input_precision='ieee'),
+            input_precision='ieee')
+
+    Ai41 = -tl.dot(
+        Ai44,
+        tl.dot(b_A41, Ai11, input_precision='ieee') +
+        tl.dot(b_A42, Ai21, input_precision='ieee') +
+        tl.dot(b_A43, Ai31, input_precision='ieee'),
+        input_precision='ieee'
+    )
+
+    tl.store(p_Ai11,Ai11.to(p_Ai11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai22,Ai22.to(p_Ai22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai33,Ai33.to(p_Ai33.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai44,Ai44.to(p_Ai44.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai21,Ai21.to(p_Ai21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai31,Ai31.to(p_Ai31.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai32,Ai32.to(p_Ai32.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai41,Ai41.to(p_Ai41.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai42,Ai42.to(p_Ai42.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai43,Ai43.to(p_Ai43.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+
+
+
+def gated_chunk_scaled_dot_kkt_fwd(k: torch.Tensor,
+                                   beta: torch.Tensor,
+                                   mask: torch.Tensor,
+                                   g_cumsum:Optional[torch.Tensor] = None,
+                                   BT:int = 32,
+                                   output_dtype: torch.dtype=torch.float32):
+    B, H, T, K = k.shape
+    r = mask.shape[-1] #B H T r r
+    NT = triton.cdiv(T, BT)
+    BK = min(triton.next_power_of_2(K//r), 64)
+    A = torch.empty(B*H*NT,BT*BT,r*r,device=k.device, dtype=output_dtype).contiguous()   
+    Ag = torch.empty(B*H*NT,BT*BT,r*r,device=k.device, dtype=output_dtype).contiguous()                                                                                                                      
+    gated_chunk_scaled_dot_kkt_fwd_kernel[(NT, B*H)](
+        k, beta, g_cumsum, mask, A,Ag,
+        T*K, K, 1,
+        T, K, r, BT, BK
+    )
+    return A,Ag
+
+def solve_tril(A,mask,k,BT,output_dtype=torch.float32):
+    B, H, T, K = k.shape
+    r = mask.shape[-1]
+    NT = triton.cdiv(T, 16)
+    Ad = torch.empty(B,H,NT*16*r,16*r,device=A.device, dtype=torch.float if BT != 16 else output_dtype)
+    solve_tril_16x16_kernel[(NT, B*H)](
+            A,Ad,
+            T*BT*r*r,#s_abh
+            T*16*r*r,#s_adbh
+            T,
+            r, BT
+    )
+    if BT == 16:                                                                                                                       
+        return Ad
+
+    A = rearrange(A,'b (t l) (c r)->b (t c) (l r)',t=BT,c=r).contiguous()#BT*r BT*r  
+    if BT == 32:
+        NT = triton.cdiv(T, BT)
+        Ai = torch.zeros(B,H,NT*BT*r,BT*r,device=A.device, dtype=output_dtype)
+        merge_16x16_to_32x32_inverse_kernel[(NT, B*H)](
+            A,Ad,Ai,
+            T*BT*r*r,#s_a_bh and s_ai_bh
+            T*16*r*r,#s_ad_bh
+            T,r,BT
+        )
+        return Ai
+
+    if BT == 64:
+        NT = triton.cdiv(T, BT)
+        Ai = torch.zeros(B,H,NT*BT*r,BT*r,device=A.device, dtype=output_dtype)
+        merge_16x16_to_64x64_inverse_kernel[(NT, B*H)](
+            A,Ad,Ai,
+            T*BT*r*r,#s_a_bh and s_ai_bh
+            T*16*r*r,#s_ad_bh
+            T,r,BT
+        )
+        return Ai
+
+
+def gated_fwd_recompute_w_u(k, v, beta,mask, Aw,Au,BT):
+    B, H, T, K, V = *k.shape, v.shape[-1]
+    r = mask.shape[-1]
+    u = torch.empty(B,H,r*T,V,device=k.device, dtype=k.dtype)
+    w = torch.empty(B,H,r*T,K,device=k.device, dtype=k.dtype)
+    NT = triton.cdiv(T, BT)
+    BK = min(triton.next_power_of_2(K//r), 64)#32
+    BV = min(triton.next_power_of_2(V), 64)
+    gated_fwd_recompute_w_u_kernel[(NT, B*H)](
+        k, v, beta,mask, w, u, Aw,Au,
+        T*K, K, 1,
+        T*V, V, 1,
+        T, K, V, r,BT, BK, BV
+    )
+    return w, u
+
+
+
+
+# class WYRepresentationPrepration(torch.autograd.Function):
+#     @staticmethod
+#     @contiguous
+#     @autocast_custom_fwd
+#     def forward(ctx, k, v, beta,mask,chunk_size=64):
+#         ctx.BT = chunk_size
+#         w, u, A = fwd_prepare_wy_repr(k, v,beta,mask, ctx.BT)
+#         ctx.save_for_backward(k, v, beta,mask,A)
+#         return w, u
+#     @staticmethod
+#     @contiguous
+#     @autocast_custom_bwd
+#     def backward(ctx, dw, du):
+#         k, v, beta,mask, A = ctx.saved_tensors
+#         BT = ctx.BT
+#         dk, dv, dbeta,dmask = bwd_prepare_wy_repr(k, v, beta,mask, A, dw, du, BT)
+#         return dk, dv, dbeta, dmask, None
+
+# prepare_wy_repr = WYRepresentationPrepration.apply
+
+
+# def naive(k, v, beta,maskij,chunk_size):
+#     l_org = k.shape[2]
+#     l_new = triton.next_power_of_2(l_org)
+#     k = torch.cat([k, torch.zeros_like(k)[:, :, :l_new-l_org, :]], dim=2)
+#     v = torch.cat([v, torch.zeros_like(v)[:, :, :l_new-l_org, :]], dim=2)
+#     beta = torch.cat([beta, torch.zeros_like(beta)[:, :, :l_new-l_org]], dim=2)
+#     k, v = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), (k, v))
+#     beta = rearrange(beta, 'b h (n c) -> b h n c', c=chunk_size)
+    
+#     b,h,nt,BT,dk = k.shape
+#     dv = v.shape[-1]
+#     r = maskij.shape[-1] 
+#     k_beta = k * beta[..., None]
+#     k_beta = rearrange(k_beta,'b h n t (r k)->b h n t r k', r=r)
+#     k_beta = torch.einsum('b h n t r k,l r-> b h n t l r k',k_beta,maskij)
+#     k_beta = rearrange(k_beta,'b h n t l r k->b h n t l (r k)')#l=1 rk=org
+#     v_beta = v * beta[..., None]
+#     v_beta = v_beta
+#     v_beta = v_beta.unsqueeze(-2).expand(-1,-1,-1,-1,r,-1)
+#     ki = rearrange(k,'b h n c (r k)-> b h n r c k',r=r)
+    
+#     attn = (ki @ ki.transpose(-1, -2))
+#     attn = torch.tril(attn, diagonal=-1)#bhnr cc
+#     attn = torch.einsum('b h n r t l,c r->b h n t l c r',attn,maskij)#bhn  rr cc
+#     attn = torch.einsum('b h n t l c r,b h n t->b h n t l c r',attn,beta)
+
+#     o = torch.zeros_like(k_beta)
+#     o2 = torch.zeros_like(v_beta)
+
+#     o[..., 0, :,:] = k_beta[..., 0,:,:].clone()
+#     o2[..., 0,:, :] = v_beta[..., 0,:,:].clone()
+#     for i in range(1, chunk_size):
+#         o_i = (o[..., :i,:,:]).clone()#bhn :t cc  
+#         o[..., i,:,:] =  (-(attn[:,:,:,i, :i,:,:]@o_i).sum(3) + k_beta[..., i,:,:])
+#         o2_i = (o2[..., :i,:,:]).clone()#少一个维度
+#         o2[..., i,:,:] = (-(attn[:,:,:,i, :i,:,:]@o2_i).sum(3) + v_beta[..., i,:,:])
+#     return map(lambda x: rearrange(x, 'b h n c r k -> b h (n c r) k'), (o, o2))
+
+
+# if __name__ == "__main__":
+#     #all compute here
+#     import sys
+#     sys.path.append('/mnt/jfzn/msj/flash-linear-attention-main/legacy/training/fla2-copy')
+#     torch.set_default_dtype(torch.bfloat16)
+#     seq_len = 32
+#     b = 2
+#     h = 2
+#     k = torch.nn.functional.normalize(torch.randn(b, h, seq_len, 128), dim=-1, p=2)#d=128
+#     v = torch.randn(b, h, seq_len, 128)
+#     beta = torch.rand(b, h, seq_len).sigmoid()
+#     require_grad = True
+#     BT = 16
+#     k, v, beta = map(lambda x: x.cuda().requires_grad_(require_grad).contiguous(), (k, v, beta))
+#     r = 4
+#     # mask = torch.tensor([[1,1,0,0],[0.5,1,0.5,0],[0,0.5,1,0.5],[0,0,1,1]]).cuda().contiguous()
+#     mask = torch.randn([r,r])
+#     mask = mask.cuda().requires_grad_(require_grad).contiguous()
+#     # w,u,a0 = fwd_prepare_wy_repr(k,v,beta,mask, 16)
+#     # w2,u2 = fwd_recompute_w_u(k,v,beta,mask,a0,16)
+#     # from einops import rearrange
+
+#     k2 = rearrange(k,'b h (n t) (r k)-> b h n r t k',t = 16,r=r)
+#     b2 = rearrange(beta,'b h (n t)-> b h n t',t = 16)
+#     a1 = (k2*b2.unsqueeze(-2).unsqueeze(-1))@k2.transpose(-1,-2)#bhnrtt
+#     qq = torch.tril(a1,diagonal=-1)
+#     qq = torch.einsum('b h n r t l,c r-> b h n t c l r',qq,mask)
+#     sf = rearrange(qq,'b h n t c l r->b h n (t c) (l r)')
+#     sf = rearrange(sf,'b h n (t c) (l r)->b h n t l c r',c=r ,r =r)#这个
+    
+    
+#     # #长条对角线
+#     i_mask = ((torch.arange(0, BT)[:, None, None, None] == torch.arange(0, BT)[None, :, None, None]) & (torch.arange(0, r)[None, None, :, None] == torch.arange(0, r)[None, None, None, :]))
+#     s = sf+i_mask.unsqueeze(0).unsqueeze(0).unsqueeze(0).cuda()
+#     s = rearrange(s,'b h n a d c r->b h n (a c) (d r)')
+#     s = torch.linalg.inv(s.float()).to(k)#矩阵逆#bhn tr tr
+    
+
+#     # A = chunk_scaled_dot_kkt_fwd(k,beta,mask,BT,output_dtype=torch.float32)#bh nt BT bt r r
+#     # Ad = solve_tril(A,mask,k,BT,output_dtype=torch.float32)
+#     # s = rearrange(s,'b h n a c->(b h) (n a) c')
+#     # print(Ad)
+#     # print(s)
+#     # print((Ad-s).abs().max())
+
+#     w,u,As = fwd_prepare_wy_repr(k, v, beta,mask, 16)
+#     As = rearrange(As,'b h (n t) l->(b h n) t l',t =BT*r)
+#     # print((As-s).abs().max())
+#     # B*H*NT,BT*r,16*r
+#     # k_exp = torch.einsum('b h n r t k,b h n t-> b h n r t k',k2,b2)
+#     # k_exp = torch.einsum('b h n r t k,c r-> b h n r t k c',k_exp,mask)
+#     # k_exp = rearrange(k_exp,'b h n r t k c->b h n (t c) (r k)')
+#     # wc = s_copy@k_exp
+
+#     # v_exp = rearrange(v,'b h (n t) v-> b h n t v',t = BT)
+#     # v_exp = torch.einsum('b h n t v,b h n t-> b h n t v',v_exp,b2)
+#     # v_exp = v_exp.unsqueeze(4).expand(-1,-1,-1,-1,r,-1)
+#     # v_exp = rearrange(v_exp, ' b h n t r v-> b h n (t r) v')
+#     # uc = s_copy@v_exp
+#     # wc,uc = map(lambda x: rearrange(x,"b h n t r->b h (n t) r"), (wc,uc))
+#     # do = torch.rand_like(wc)
+#     # do2 = torch.rand_like(uc)#b h n t t
+#     # o1, o2 = naive(k.clone(), v.clone(), beta.clone(),mask.clone(), BT)#这个代码有问题
+#     # do = torch.rand_like(o1)
+#     # do2 = torch.rand_like(o2)#b h n t t
+#     # if require_grad:
+#     #     o1.backward(do, retain_graph=True)
+#     #     o2.backward(do2, retain_graph=True)
+#     #     k_grad2, v_grad2, beta_grad2,mask_grad2 = k.grad, v.grad, beta.grad, mask.grad
+
+#     # w0,u0,s0 = fwd_prepare_wy_repr(k, v, beta,mask, 16)
+#     # k_grad, v_grad, beta_grad,mask_grad = bwd_prepare_wy_repr(k,v,beta,mask,s0,do,do2,BT)   
+    
+#     # print((o1-w0).abs().max())
+#     # print((o2-u0).abs().max())
+#     # print((k_grad-k_grad2).abs().max())
+#     # print((v_grad-v_grad2).abs().max())
+#     # print((beta_grad-beta_grad2).abs().max())
+#     # print((mask_grad-mask_grad2).abs().max())
+#     # print(mask_grad)
+#     # print(mask_grad2)
+
+

fla2/ops/mask_gated_delta_rule_t/wy_fast_test.py

@@ -0,0 +1,676 @@
+# -*- coding: utf-8 -*-
+import pdb
+import torch
+import triton
+import triton.language as tl
+from einops import rearrange
+# from ...utils import autocast_custom_bwd, autocast_custom_fwd, contiguous
+from fla.utils import autocast_custom_bwd, autocast_custom_fwd, contiguous
+# Inspired by "THE WY REPRESENTATION FOR PRODUCTS OF HOUSEHOLDER MATRICES" https://epubs.siam.org/doi/pdf/10.1137/0908009
+# o: cumprod
+# o2: cumprodsum
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16)
+    ],
+    key=["BT", "BK", "BV"],
+)
+@triton.jit
+def fwd_prepare_wy_repr_kernel(
+    k,
+    v,
+    beta,
+    mask_ij,
+    w,
+    u,
+    A,
+    s_qk_h,
+    s_qk_t,
+    s_qk_d,
+    s_vo_h,
+    s_vo_t,
+    s_vo_d,
+    T,
+    K,
+    V,
+    r:  tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    b_A = tl.zeros([BT,BT,r,r], dtype=tl.float32)#r*BT r*BT
+    dk = K//r
+    p_beta = tl.make_block_ptr(beta + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+    b_beta = tl.load(p_beta, boundary_check=(0,))
+    for i_r in range(r):
+        r_mask = tl.arange(0, r) == i_r 
+        p_mask = mask_ij + tl.arange(0,r)* r + i_r#列读，因而是行数目
+        b_mask = tl.load(p_mask)
+        ij_mask = b_mask[:,None]*r_mask[None,:]#行数
+
+        for i_k in range(tl.cdiv(dk, BK)):#分块k读取计算
+            p_k = tl.make_block_ptr(k + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_r * dk + i_k * BK), (BT, BK), (1, 0))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            b_kb = (b_k * b_beta[:, None]).to(b_k.dtype)
+            dot = tl.dot(b_kb, tl.trans(b_k), allow_tf32=False)
+            b_A += dot[:,:,None,None]*ij_mask[None,None,:,:]
+    b_A = -tl.where((tl.arange(0, BT)[:,None] > tl.arange(0, BT)[None,:])[:,:,None,None], b_A, 0)
+    #先save这个看看
+
+    for i in range(1, BT):#此时矩阵为 BT,r,BT,r
+        mask = tl.arange(0, BT) == i 
+        b_a = tl.sum(tl.where(mask[:,None,None,None], b_A, 0), 0)#get ba BT*r*r
+        q = tl.sum(b_a[:,None,:,:,None]*b_A[:,:,None,:,:],-2)#矩阵乘法解决，get BT,BT*r*r
+        b_a = b_a + tl.sum(q,0)*((tl.arange(0, BT) < i)[:,None,None])#BT*r*r
+        b_A = tl.where(mask[:,None,None,None],b_a,b_A)#按行计算 ，逐步交换结果
+    b_A += ((tl.arange(0, BT)[:, None, None, None] == tl.arange(0, BT)[None, :, None, None])&(tl.arange(0, r)[None, None, :, None] == tl.arange(0, r)[None, None, None, :]))
+    b_A = tl.permute(b_A,(0,2,1,3))
+    b_A = tl.reshape(b_A,(BT*r,BT*r))#BT*r BT*r
+    p_A = tl.make_block_ptr(A + i_bh*T*BT*r*r ,(T*r,BT*r), (BT*r,1), (i_t*BT*r,0), (BT*r,BT*r),(1,0))#旧版本实现需要很多乘法
+    tl.store(p_A, (b_A).to(p_A.dtype.element_ty),boundary_check=(0, 1))
+    #解决矩阵求逆
+    b_A = b_A.to(k.dtype.element_ty)#ok 解决求逆了 #下一步计算结果
+
+    for i_r in range(r):
+        p_mask = mask_ij + tl.arange(0,r)*r+i_r#读取第ir列
+        b_mask = tl.load(p_mask)
+        for i_k in range(tl.cdiv(dk, BK)):
+            p_k = tl.make_block_ptr(k + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_r*dk + i_k * BK), (BT, BK), (1, 0))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            b_kb = (b_k * b_beta[:, None]).to(b_k.dtype)[:,None,:]*b_mask[None,:,None].to(b_k.dtype)#BT*r*d
+            b_kb = tl.reshape(b_kb,(BT*r,BK))
+            b_w = tl.dot(b_A, b_kb, allow_tf32=False)#get BT*r *BK
+            p_w = tl.make_block_ptr(w + i_bh * s_qk_h*r, (T*r, K), (s_qk_t, s_qk_d), (i_t * BT * r, i_r*dk + i_k * BK), (BT*r, BK), (1, 0))
+            tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
+
+    for i_v in range(tl.cdiv(V, BV)):#no need for 任意mask不使用 #无需for 循环 ，这里也不存在mask
+        p_v = tl.make_block_ptr(v + i_bh * s_vo_h, (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        b_vb = (b_v * b_beta[:, None]).to(b_v.dtype)[:,None,:]*tl.full([r],1, dtype=b_v.dtype)[None,:,None]
+        b_vb = tl.reshape(b_vb,(BT*r,BV))
+        b_u = tl.dot(b_A, b_vb, allow_tf32=False)
+        p_u = tl.make_block_ptr(u + i_bh * s_vo_h*r, (T*r, V), (s_vo_t, s_vo_d), (i_t * BT*r, i_v * BV), (BT*r, BV), (1, 0))
+        tl.store(p_u, (b_u).to(p_u.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16)
+    ],
+    key=["BT", "BK", "BV"],
+)
+@triton.jit
+def fwd_recompute_w_u_kernel(
+    k,
+    v,
+    beta,
+    mask_ij,
+    w,
+    u,
+    A,
+    s_qk_h,
+    s_qk_t,
+    s_qk_d,
+    s_vo_h,
+    s_vo_t,
+    s_vo_d,
+    T,
+    K,
+    V,
+    r: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    dk = K//r
+    p_beta = tl.make_block_ptr(beta + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+    b_beta = tl.load(p_beta, boundary_check=(0,))
+    p_A = tl.make_block_ptr(A + i_bh*T*BT*r*r ,(T*r,BT*r), (BT*r,1), (i_t*BT*r,0), (BT*r,BT*r),(1,0))
+    b_A = tl.load(p_A, boundary_check=(0, 1)).to(k.dtype.element_ty)
+    for i_r in range(r):
+        # r_mask = tl.arange(0, r) == i_r # 
+        p_mask = mask_ij + tl.arange(0,r)*r+i_r#读取第ir列
+        b_mask = tl.load(p_mask)
+        for i_k in range(tl.cdiv(dk, BK)):
+            p_k = tl.make_block_ptr(k + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_r*dk + i_k * BK), (BT, BK), (1, 0))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            b_kb = (b_k * b_beta[:, None]).to(b_k.dtype)[:,None,:]*b_mask[None,:,None].to(b_k.dtype)#BT*r*d
+            b_kb = tl.reshape(b_kb,(BT*r,BK))
+            b_w = tl.dot(b_A, b_kb, allow_tf32=False)#get BT*r *BK
+            p_w = tl.make_block_ptr(w + i_bh * s_qk_h*r, (T*r, K), (s_qk_t, s_qk_d), (i_t * BT * r, i_r*dk + i_k * BK), (BT*r, BK), (1, 0))
+            tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
+
+    for i_v in range(tl.cdiv(V, BV)):#no need for 任意mask不使用 #无需for 循环 ，这里也不存在mask
+        p_v = tl.make_block_ptr(v + i_bh * s_vo_h, (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        b_vb = (b_v * b_beta[:, None]).to(b_v.dtype)[:,None,:]*tl.full([r],1, dtype=b_v.dtype)[None,:,None]
+        b_vb = tl.reshape(b_vb,(BT*r,BV))
+        b_u = tl.dot(b_A, b_vb, allow_tf32=False)
+        p_u = tl.make_block_ptr(u + i_bh * s_vo_h*r, (T*r, V), (s_vo_t, s_vo_d), (i_t * BT*r, i_v * BV), (BT*r, BV), (1, 0))
+        tl.store(p_u, (b_u).to(p_u.dtype.element_ty), boundary_check=(0, 1))
+
+#compute this 
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16)
+    ],
+    key=["BT", "BK", "BV"],
+)
+@triton.jit
+def bwd_prepare_wy_repr_kernel(
+    k, v, beta,mask_ij,A,
+    dw, du,
+    dk, dv, dbeta,dmask,
+    s_qk_h,
+    s_qk_t,
+    s_qk_d,
+    s_vo_h,
+    s_vo_t,
+    s_vo_d,
+    T,
+    K,
+    V,
+    r: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    p_A = tl.make_block_ptr(A + i_bh*T*BT*r*r ,(T*r,BT*r), (BT*r,1), (i_t * BT * r,0), (BT*r,BT*r),(1,0))
+    b_A = tl.load(p_A, boundary_check=(0, 1)).to(k.dtype.element_ty)
+    b_dbeta = tl.zeros([BT], dtype=tl.float32)
+    b_dA = tl.zeros([BT*r,BT*r], dtype=tl.float32)
+    p_beta = tl.make_block_ptr(beta + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+    b_beta = tl.load(p_beta, boundary_check=(0,))
+
+    b_dmask = tl.zeros([r,r],dtype=tl.float32)
+    for i_v in range(tl.cdiv(V, BV)):#分块r 
+        p_v = tl.make_block_ptr(v + i_bh * s_vo_h, (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        p_du = tl.make_block_ptr(du + i_bh * s_vo_h * r, (T * r, V), (s_vo_t, s_vo_d), (i_t * BT * r, i_v * BV), (BT * r, BV), (1, 0))#r*BT BV
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        b_v_beta = ((b_v * b_beta[:, None])[:,None,:]*tl.full([r],1, dtype=b_v.dtype)[None,:,None]).to(b_v.dtype)##BT*r*BV
+        b_v_beta = tl.reshape(b_v_beta,(BT*r,BV))
+        b_du = tl.load(p_du, boundary_check=(0, 1))
+        b_dA += tl.dot(b_du, tl.trans(b_v_beta), allow_tf32=False)#BT*r,BT*r
+        b_dv_beta = tl.dot(tl.trans(b_A), b_du, allow_tf32=False)#BT*r,BV
+        b_dv_beta = tl.reshape(b_dv_beta,(BT,r,BV))#
+        sum_dv = tl.sum(b_dv_beta,-2)#这里不一样，结果
+        b_dv = (sum_dv * b_beta[:, None])#？哪一步结果不一样呢
+        b_dbeta += tl.sum(sum_dv * b_v, 1)
+        p_dv = tl.make_block_ptr(dv + i_bh * s_vo_h, (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))
+    block_k = K//r
+    for i_r in range(r):
+        p_mask = mask_ij + tl.arange(0,r)*r + i_r#读取第ir列
+        b_mask = tl.load(p_mask)#第r列
+        rmask = tl.arange(0, r) == i_r #第r列
+        for i_k in range(tl.cdiv(block_k, BK)):
+            p_k = tl.make_block_ptr(k + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_r*block_k + i_k * BK), (BT, BK), (1, 0))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            p_dw = tl.make_block_ptr(dw + i_bh * s_qk_h*r, (T*r, K), (s_qk_t, s_qk_d), (i_t * BT * r, i_r*block_k + i_k * BK), (BT * r, BK), (1, 0))
+            b_k_beta = ((b_k * b_beta[:, None])[:,None,:]*b_mask[None,:,None]).to(b_k.dtype)#BT*r*d
+            b_k_beta = tl.reshape(b_k_beta,(BT*r,BK))
+            b_dw = tl.load(p_dw, boundary_check=(0, 1))
+            b_dA += tl.dot(b_dw, tl.trans(b_k_beta), allow_tf32=False)
+            b_dk_beta = tl.dot(tl.trans(b_A), b_dw, allow_tf32=False)
+            b_dk_beta = tl.reshape(b_dk_beta,(BT,r,BK))
+            sum_dk = tl.sum(b_dk_beta * b_mask[None,:,None],1)
+            b_dk = sum_dk* b_beta[:, None]
+            b_dbeta += tl.sum(sum_dk * b_k, 1)
+
+
+            b_ss = b_dk_beta * b_beta[:,None,None] * b_k[:,None,:]
+            b_ss = tl.reshape(tl.permute(b_ss,(2,0,1)),(BT*BK,r))
+            b_ss = tl.sum(b_ss,0)
+            # b_ss = (tl.sum(tl.sum(b_dk_beta * b_beta[:,None,None] * b_k[:,None,:],0),-1))
+            b_dmask += (b_ss[:,None]*rmask[None,:]).to(tl.float32)
+            p_dk = tl.make_block_ptr(dk + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_r*block_k + i_k * BK), (BT, BK), (1, 0))
+            tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
+    
+    i = tl.arange(0, BT * r)[:, None]
+    j = tl.arange(0, BT * r)[None, :]
+    iB = i // r
+    jB = j // r
+    da_mask = iB > jB
+    b_dA = tl.where(da_mask, b_dA, 0)
+    b_dA = tl.dot(b_dA.to(b_A.dtype), tl.trans(b_A), allow_tf32=False)
+    b_dA = tl.dot(tl.trans(b_A), b_dA.to(b_A.dtype), allow_tf32=False)
+    b_dA = tl.where(da_mask, -b_dA, 0) #等价于 kkt的 dA 很多0，对角处
+    
+
+    b_dA = tl.reshape(b_dA,(BT,r,BT,r))
+    #bt r bt r
+
+
+    for i_r in range(r):#只取ir项 
+        p_mask = mask_ij + tl.arange(0,r)*r+i_r#读取第ir列
+        b_mask = tl.load(p_mask)#第ir列
+        rmask = tl.arange(0, r) == i_r #第ir列
+        g = tl.sum(tl.where(rmask[None,None,None,:], b_dA, 0), -1)#BT r BT #取出第ir列
+        ir_A = tl.sum(g * b_mask[None,:,None],1).to(k.dtype.element_ty)#BT BT
+        #对应的c部分
+
+        for i_k in range(tl.cdiv(block_k, BK)):#ik = 1
+            p_k = tl.make_block_ptr(k + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_r*block_k + i_k * BK), (BT, BK), (1, 0))
+            p_dk = tl.make_block_ptr(dk + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_r*block_k + i_k * BK), (BT, BK), (1, 0))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            b_dk = tl.load(p_dk, boundary_check=(0, 1))
+            b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype)#BT*BK
+
+            b_dk_beta = tl.dot(ir_A, b_k, allow_tf32=False)
+            b_dbeta += tl.sum(b_dk_beta * b_k, 1)
+            b_dk += tl.dot(tl.trans(ir_A), b_k_beta, allow_tf32=False)
+            b_dk += b_dk_beta * b_beta[:, None]
+            tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
+
+            beta_kkt = (tl.dot(b_k_beta,tl.trans(b_k), allow_tf32=False))#BT BT
+
+            beta_y = (beta_kkt[:,None,:]*g)
+            beta_y = tl.reshape(tl.permute(beta_y,(2,0,1)),(BT*BT,r))
+            betas = tl.sum(beta_y,0)
+            b_dmask +=  (betas[:,None]*rmask[None,:]).to(tl.float32)
+
+    p_dbeta = tl.make_block_ptr(dbeta + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+    tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,))
+    
+    p_dmask = tl.make_block_ptr(dmask + (i_bh * (T//BT) + i_t)* r * r , (r,r), (r,1), (0,0), (r,r), (1,0))
+    tl.store(p_dmask, b_dmask.to(p_dmask.dtype.element_ty), boundary_check=(0,1))
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16)
+    ],
+    key=["BT", "BK", "r"],
+)
+@triton.jit
+def chunk_scaled_dot_kkt_fwd_kernel(
+    k,
+    beta,
+    mask_ij,
+    A,
+    s_qk_h,
+    s_qk_t,
+    s_qk_d,
+    T,
+    K,
+    r:  tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    b_A = tl.zeros([BT,BT,r,r], dtype=tl.float32)#r*BT r*BT
+    dk = K//r
+    p_beta = tl.make_block_ptr(beta + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+    b_beta = tl.load(p_beta, boundary_check=(0,))
+    for i_r in range(r):
+        r_mask = tl.arange(0, r) == i_r 
+        p_mask = mask_ij + tl.arange(0,r)* r + i_r#列读，因而是行数目
+        b_mask = tl.load(p_mask)
+        ij_mask = b_mask[:,None]*r_mask[None,:]#行数
+
+        for i_k in range(tl.cdiv(dk, BK)):#分块k读取计算
+            p_k = tl.make_block_ptr(k + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_r * dk + i_k * BK), (BT, BK), (1, 0))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            b_kb = (b_k * b_beta[:, None]).to(b_k.dtype)
+            dot = tl.dot(b_kb, tl.trans(b_k), allow_tf32=False)
+            b_A += dot[:,:,None,None]*ij_mask[None,None,:,:]
+    b_A = tl.where((tl.arange(0, BT)[:,None] > tl.arange(0, BT)[None,:])[:,:,None,None], b_A, 0)
+    p_A = tl.make_block_ptr(A + (i_bh*T//BT+i_t)*BT*BT*r*r ,(BT,BT,r,r), (BT*r*r,r*r,r,1), (0,0,0,0), (BT,BT,r,r),(3,2,1,0))
+    tl.store(p_A, (b_A).to(p_A.dtype.element_ty),boundary_check=(0,1,2,3))
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16)
+    ],
+    key=["BT", "r"],
+)
+@triton.jit
+def solve_tril_16x16_kernel(
+    A,
+    Ad,
+    s_A_bh,
+    s_Ad_bh,
+    T,
+    r:  tl.constexpr,
+    BT: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    offset = (i_t * 16) % BT 
+
+    p_A = tl.make_block_ptr(A + (i_bh)*s_A_bh, (T,BT,r,r),(BT*r*r,r*r,r,1) ,(i_t * 16, offset, 0, 0), (16, 16,r,r), (3,2,1,0))
+    b_A = tl.load(p_A, boundary_check=(0,1,2,3)).to(tl.float32)
+    b_A = -tl.where((tl.arange(0, 16)[:,None] > tl.arange(0, 16)[None,:])[:,:,None,None], b_A, 0)
+
+    for i in range(1, 16):
+        mask = tl.arange(0, 16) == i 
+        b_a = tl.sum(tl.where(mask[:,None,None,None], b_A, 0), 0)
+        q = (tl.sum(b_a[:,None,:,:,None]*b_A[:,:,None,:,:],-2))
+        b_a = b_a + tl.sum(q,0)*((tl.arange(0, 16) < i)[:,None,None])
+        b_A = tl.where(mask[:,None,None,None],b_a,b_A)#按行计算 ，逐步交换结果
+    b_A += ((tl.arange(0, 16)[:, None, None, None] == tl.arange(0, 16)[None, :, None, None])&(tl.arange(0, r)[None, None, :, None] == tl.arange(0, r)[None, None, None, :]))
+    
+    b_A = tl.permute(b_A,(0,2,1,3))
+    b_A = tl.reshape(b_A,(16*r,16*r))#BT*r BT*r
+    p_Ad = tl.make_block_ptr(Ad + (i_bh)*s_Ad_bh,(T*r,16*r),(16*r,1), (i_t * 16 * r, 0), (16*r,16*r), (1,0))
+    tl.store(p_Ad, (b_A).to(p_Ad.dtype.element_ty),boundary_check=(0,1))
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16)
+    ],
+    key=["r"],
+)
+@triton.jit
+def merge_16x16_to_32x32_inverse_kernel(
+        A,
+        Ad,
+        Ai,
+        s_A_bh,
+        s_Ad_bh,
+        T,
+        r: tl.constexpr,
+        BT: tl.constexpr 
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+
+    p_A21 = tl.make_block_ptr(A + (i_bh)*s_A_bh, (T,32,r,r),(32*r*r,r*r,r,1) ,(i_t * 32 + 16, 0, 0, 0), (16, 16,r,r), (3,2,1,0))
+    b_A21 = tl.load(p_A21, boundary_check=(0,1,2,3)).to(tl.float32)
+    b_A21 = tl.permute(b_A21,(0,2,1,3))
+    b_A21 = tl.reshape(b_A21,(16*r,16*r))#BT*r BT*r
+
+    p_Ad11  = tl.make_block_ptr(Ad + (i_bh)*s_Ad_bh,(T*r,16*r),(16*r,1), (i_t * 32 * r, 0), (16*r,16*r), (1,0))
+    p_Ad22  = tl.make_block_ptr(Ad + (i_bh)*s_Ad_bh,(T*r,16*r),(16*r,1), ((i_t *32 +16) * r, 0), (16*r,16*r), (1,0))
+
+
+    p_Ai11 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,32*r), (32*r, 1), (i_t * 32 * r , 0), (16*r, 16*r), (1, 0))
+    p_Ai22 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,32*r), (32*r, 1), ((i_t * 32 + 16) * r , 16*r), (16*r, 16*r), (1, 0))
+    p_Ai21 = tl.make_block_ptr(Ai+ (i_bh)*s_A_bh, (T*r,32*r), (32*r, 1), ((i_t * 32 + 16) * r, 0), (16*r, 16*r), (1, 0))
+
+    Ai11 = tl.load(p_Ad11, boundary_check=(0, 1)).to(tl.float32)
+    Ai22 = tl.load(p_Ad22, boundary_check=(0, 1)).to(tl.float32)
+    Ai21 = -tl.dot(tl.dot(Ai22,b_A21, input_precision='ieee'),Ai11,input_precision='ieee')
+    tl.store(p_Ai11,Ai11.to(p_Ai11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai22,Ai22.to(p_Ai22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai21,Ai21.to(p_Ai21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+
+def chunk_scaled_dot_kkt_fwd(k,beta,mask,BT,output_dtype=torch.float32):
+    B, H, T, K = k.shape
+    r = mask.shape[-1]
+    NT = triton.cdiv(T, BT)
+    BK = min(triton.next_power_of_2(K//r), 64)
+    A = torch.empty(B*H*NT,BT*BT,r*r,device=k.device, dtype=output_dtype).contiguous()                                                                                                                         
+    chunk_scaled_dot_kkt_fwd_kernel[(NT, B*H)](
+        k, beta, mask, A,
+        T*K, K, 1,
+        T, K, r, BT, BK
+    )
+    return A
+
+def solve_tril(A,mask,k,BT,output_dtype=torch.float32):
+    B, H, T, K = k.shape
+    r = mask.shape[-1]
+    NT = triton.cdiv(T, 16)
+    Ad = torch.empty(B,H,NT*16*r,16*r,device=A.device, dtype=torch.float if BT != 16 else output_dtype)
+    solve_tril_16x16_kernel[(NT, B*H)](
+            A,Ad,
+            T*BT*r*r,#s_abh
+            T*16*r*r,#s_adbh
+            T,
+            r, BT
+    )
+    if BT == 16:                                                                                                                       
+        return Ad
+
+    NT = triton.cdiv(T, BT)
+    Ai = torch.zeros(B,H,NT*BT*r,BT*r,device=A.device, dtype=output_dtype)
+    merge_16x16_to_32x32_inverse_kernel[(NT, B*H)](
+        A,Ad,Ai,
+        T*BT*r*r,#s_a_bh and s_ai_bh
+        T*16*r*r,#s_ad_bh
+        T,r,BT
+    )
+    return Ai
+
+
+def fwd_prepare_wy_repr2(k, v, beta,mask, BT):
+    A = chunk_scaled_dot_kkt_fwd(k,beta,mask,BT,torch.float32)
+    A = solve_tril(A=A,mask=mask,k=k,BT=BT,output_dtype=k.dtype)
+    w, u = fwd_recompute_w_u(k, v, beta,mask, A, BT)
+    return w, u, A
+
+def fwd_prepare_wy_repr(k, v, beta,mask, BT):
+    B, H, T, K, V = *k.shape, v.shape[-1]
+    r = mask.shape[-1]
+    u = torch.empty(B,H,r*T,V,device=k.device, dtype=k.dtype)
+    w = torch.empty(B,H,r*T,K,device=k.device, dtype=k.dtype)
+    NT = triton.cdiv(T, BT)
+    BK = min(triton.next_power_of_2(K//r), 64)
+    BV = min(triton.next_power_of_2(V), 64)
+    A = torch.empty(B,H,NT*BT*r,BT*r,device=k.device, dtype=k.dtype)                                                                                                                         
+    fwd_prepare_wy_repr_kernel[(NT, B*H)](
+        k, v, beta, mask, w, u, A,
+        T*K, K, 1,
+        T*V, V, 1,
+        T, K, V, r, BT, BK, BV
+    )
+    return w, u, A
+
+
+def fwd_recompute_w_u(k, v, beta,mask, A, BT):
+    B, H, T, K, V = *k.shape, v.shape[-1]
+    r = mask.shape[-1]
+    u = torch.empty(B,H,r*T,V,device=k.device, dtype=k.dtype)
+    w = torch.empty(B,H,r*T,K,device=k.device, dtype=k.dtype)
+    NT = triton.cdiv(T, BT)
+    BK = min(triton.next_power_of_2(K//r), 64)#32
+    BV = min(triton.next_power_of_2(V), 64)
+    fwd_recompute_w_u_kernel[(NT, B*H)](
+        k, v, beta,mask, w, u, A,
+        T*K, K, 1,
+        T*V, V, 1,
+        T, K, V, r,BT, BK, BV
+    )
+    return w, u
+
+def bwd_prepare_wy_repr(k, v, beta, mask, A, dw, du, BT):
+    B, H, T, K, V = *k.shape, v.shape[-1]
+    r = mask.shape[-1]
+    NT = triton.cdiv(T, BT)
+    BK = min(triton.next_power_of_2(K//r), 64)
+    BV = min(triton.next_power_of_2(V), 64)
+    NT = triton.cdiv(T, BT)
+    dk = torch.empty_like(k)
+    dv = torch.empty_like(v).contiguous()
+    dbeta = torch.zeros_like(beta)
+    dmask = torch.zeros([B*H*NT,r,r],device=k.device,dtype=k.dtype).contiguous()
+    bwd_prepare_wy_repr_kernel[(NT, B*H)](
+        k, v, beta, mask, A,
+        dw, du,
+        dk, dv, dbeta,dmask,
+        T*K, K, 1,
+        T*V, V, 1,
+        T, K, V, r, BT, BK, BV
+    )
+    dmask = dmask.sum(0)
+    return dk, dv, dbeta, dmask
+
+
+class WYRepresentationPrepration(torch.autograd.Function):
+    @staticmethod
+    @contiguous
+    @autocast_custom_fwd
+    def forward(ctx, k, v, beta,mask,chunk_size=64):
+        ctx.BT = chunk_size
+        w, u, A = fwd_prepare_wy_repr(k, v,beta,mask, ctx.BT)
+        ctx.save_for_backward(k, v, beta,mask,A)
+        return w, u
+    @staticmethod
+    @contiguous
+    @autocast_custom_bwd
+    def backward(ctx, dw, du):
+        k, v, beta,mask, A = ctx.saved_tensors
+        BT = ctx.BT
+        dk, dv, dbeta,dmask = bwd_prepare_wy_repr(k, v, beta,mask, A, dw, du, BT)
+        return dk, dv, dbeta, dmask, None
+
+prepare_wy_repr = WYRepresentationPrepration.apply
+
+
+def naive(k, v, beta,maskij,chunk_size):
+    l_org = k.shape[2]
+    l_new = triton.next_power_of_2(l_org)
+    k = torch.cat([k, torch.zeros_like(k)[:, :, :l_new-l_org, :]], dim=2)
+    v = torch.cat([v, torch.zeros_like(v)[:, :, :l_new-l_org, :]], dim=2)
+    beta = torch.cat([beta, torch.zeros_like(beta)[:, :, :l_new-l_org]], dim=2)
+    k, v = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), (k, v))
+    beta = rearrange(beta, 'b h (n c) -> b h n c', c=chunk_size)
+    
+    b,h,nt,BT,dk = k.shape
+    dv = v.shape[-1]
+    r = maskij.shape[-1] 
+    k_beta = k * beta[..., None]
+    k_beta = rearrange(k_beta,'b h n t (r k)->b h n t r k', r=r)
+    k_beta = torch.einsum('b h n t r k,l r-> b h n t l r k',k_beta,maskij)
+    k_beta = rearrange(k_beta,'b h n t l r k->b h n t l (r k)')#l=1 rk=org
+    v_beta = v * beta[..., None]
+    v_beta = v_beta
+    v_beta = v_beta.unsqueeze(-2).expand(-1,-1,-1,-1,r,-1)
+    ki = rearrange(k,'b h n c (r k)-> b h n r c k',r=r)
+    
+    attn = (ki @ ki.transpose(-1, -2))
+    attn = torch.tril(attn, diagonal=-1)#bhnr cc
+    attn = torch.einsum('b h n r t l,c r->b h n t l c r',attn,maskij)#bhn  rr cc
+    attn = torch.einsum('b h n t l c r,b h n t->b h n t l c r',attn,beta)
+
+    o = torch.zeros_like(k_beta)
+    o2 = torch.zeros_like(v_beta)
+
+    o[..., 0, :,:] = k_beta[..., 0,:,:].clone()
+    o2[..., 0,:, :] = v_beta[..., 0,:,:].clone()
+    for i in range(1, chunk_size):
+        o_i = (o[..., :i,:,:]).clone()#bhn :t cc  
+        o[..., i,:,:] =  (-(attn[:,:,:,i, :i,:,:]@o_i).sum(3) + k_beta[..., i,:,:])
+        o2_i = (o2[..., :i,:,:]).clone()#少一个维度
+        o2[..., i,:,:] = (-(attn[:,:,:,i, :i,:,:]@o2_i).sum(3) + v_beta[..., i,:,:])
+    return map(lambda x: rearrange(x, 'b h n c r k -> b h (n c r) k'), (o, o2))
+
+
+if __name__ == "__main__":
+    #all compute here
+    import sys
+    torch.manual_seed(42)  
+    sys.path.append('/mnt/jfzn/msj/flash-linear-attention-main/legacy/training/fla2-copy')
+    torch.set_default_dtype(torch.bfloat16)
+    seq_len = 128
+    b = 2
+    h = 2
+    k = torch.nn.functional.normalize(torch.randn(b, h, seq_len, 128), dim=-1, p=2)#d=128
+    v = torch.randn(b, h, seq_len, 128)
+    beta = torch.rand(b, h, seq_len).sigmoid()
+    require_grad = True
+    BT = 32
+    k, v, beta = map(lambda x: x.cuda().requires_grad_(require_grad).contiguous(), (k, v, beta))
+    r = 4
+    # mask = torch.tensor([[1,1,0,0],[0.5,1,0.5,0],[0,0.5,1,0.5],[0,0,1,1]]).cuda().contiguous()
+    mask = torch.randn([r,r])
+    mask = mask.cuda().requires_grad_(require_grad).contiguous()
+    # w,u,a0 = fwd_prepare_wy_repr(k,v,beta,mask, 16)
+    # w2,u2 = fwd_recompute_w_u(k,v,beta,mask,a0,16)
+    # from einops import rearrange
+
+    k2 = rearrange(k,'b h (n t) (r k)-> b h n r t k',t = BT,r=r)
+    b2 = rearrange(beta,'b h (n t)-> b h n t',t = BT)
+    a1 = (k2*b2.unsqueeze(-2).unsqueeze(-1))@k2.transpose(-1,-2)#bhnrtt
+    qq = torch.tril(a1,diagonal=-1)
+    qq = torch.einsum('b h n r t l,c r-> b h n t c l r',qq,mask)
+    sf = rearrange(qq,'b h n t c l r->b h n (t c) (l r)')
+    sf = rearrange(sf,'b h n (t c) (l r)->b h n t l c r',c=r ,r =r)#这个
+    
+    # #长条对角线
+    i_mask = ((torch.arange(0, BT)[:, None, None, None] == torch.arange(0, BT)[None, :, None, None]) & (torch.arange(0, r)[None, None, :, None] == torch.arange(0, r)[None, None, None, :]))
+    s = sf+i_mask.unsqueeze(0).unsqueeze(0).unsqueeze(0).cuda()
+    s = rearrange(s,'b h n a d c r->b h n (a c) (d r)')
+    s = torch.linalg.inv(s.float()).to(k)#矩阵逆#bhn tr tr
+    
+
+    # A = chunk_scaled_dot_kkt_fwd(k,beta,mask,BT,output_dtype=torch.float32)#bh nt BT bt r r
+    # Ad = solve_tril(A,mask,k,BT,output_dtype=torch.bfloat16)
+    # s = rearrange(s,'b h n a c->(b h n) a c')
+    # print(Ad.shape)
+    # print(s.shape)
+
+    w,u,As = fwd_prepare_wy_repr2(k, v, beta,mask, BT)
+    # w2,u2,Ad2 = fwd_prepare_wy_repr(k, v, beta,mask, BT)
+    
+    # print((w2-w).abs().max())
+    # print((u2-u).abs().max())
+    # print((As-Ad2).abs().max())
+
+    # print((Ad-s).abs().max())
+    # print(Ad-s)
+
+    # print((As-s).abs().max())
+    # print(As-s)
+    # B*H*NT,BT*r,16*r
+    # k_exp = torch.einsum('b h n r t k,b h n t-> b h n r t k',k2,b2)
+    # k_exp = torch.einsum('b h n r t k,c r-> b h n r t k c',k_exp,mask)
+    # k_exp = rearrange(k_exp,'b h n r t k c->b h n (t c) (r k)')
+    # wc = s_copy@k_exp
+
+    # v_exp = rearrange(v,'b h (n t) v-> b h n t v',t = BT)
+    # v_exp = torch.einsum('b h n t v,b h n t-> b h n t v',v_exp,b2)
+    # v_exp = v_exp.unsqueeze(4).expand(-1,-1,-1,-1,r,-1)
+    # v_exp = rearrange(v_exp, ' b h n t r v-> b h n (t r) v')
+    # uc = s_copy@v_exp
+    # wc,uc = map(lambda x: rearrange(x,"b h n t r->b h (n t) r"), (wc,uc))
+    # do = torch.rand_like(wc)
+    # do2 = torch.rand_like(uc)#b h n t t
+    # o1, o2 = naive(k.clone(), v.clone(), beta.clone(),mask.clone(), BT)#这个代码有问题
+    # do = torch.rand_like(o1)
+    # do2 = torch.rand_like(o2)#b h n t t
+    # if require_grad:
+    #     o1.backward(do, retain_graph=True)
+    #     o2.backward(do2, retain_graph=True)
+    #     k_grad2, v_grad2, beta_grad2,mask_grad2 = k.grad, v.grad, beta.grad, mask.grad
+
+    # w0,u0,s0 = fwd_prepare_wy_repr(k, v, beta,mask, 16)
+    # k_grad, v_grad, beta_grad,mask_grad = bwd_prepare_wy_repr(k,v,beta,mask,s0,do,do2,BT)   
+    
+    # print((o1-w0).abs().max())
+    # print((o2-u0).abs().max())
+    # print((k_grad-k_grad2).abs().max())
+    # print((v_grad-v_grad2).abs().max())
+    # print((beta_grad-beta_grad2).abs().max())
+    # print((mask_grad-mask_grad2).abs().max())
+    # print(mask_grad)
+    # print(mask_grad2)
+
+

fla2/ops/rwkv6/chunk.py

@@ -0,0 +1,931 @@
+# -*- coding: utf-8 -*-
+
+# Copyright (c) 2023-2024, Yu Zhang, Songlin Yang
+
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from fla.ops.utils import chunk_global_reversed_cumsum
+from fla.utils import contiguous
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({'BS': 16}, num_warps=2),
+        triton.Config({'BS': 16}, num_warps=4),
+        triton.Config({'BS': 16}, num_warps=8),
+        triton.Config({'BS': 32}, num_warps=2),
+        triton.Config({'BS': 32}, num_warps=4),
+        triton.Config({'BS': 32}, num_warps=8),
+        triton.Config({'BS': 64}, num_warps=2),
+        triton.Config({'BS': 64}, num_warps=4),
+        triton.Config({'BS': 64}, num_warps=8),
+    ],
+    key=['S']
+)
+@triton.jit
+def chunk_rwkv6_fwd_kernel_cum(
+    s,
+    o,
+    o_minus_s,
+    s_s_h,
+    s_s_t,
+    s_s_d,
+    T: tl.constexpr,
+    S: tl.constexpr,
+    BT: tl.constexpr,
+    BS: tl.constexpr
+):
+    i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    o_i = tl.arange(0, BT)
+    m_s = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.)
+
+    p_s = tl.make_block_ptr(s + i_bh * s_s_h, (T, S), (s_s_t, s_s_d), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+    p_o = tl.make_block_ptr(o + i_bh * s_s_h, (T, S), (s_s_t, s_s_d), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+    p_o_minus_s = tl.make_block_ptr(o_minus_s + i_bh * s_s_h, (T, S), (s_s_t, s_s_d), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+    # [BT, BS]
+    b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32)
+    b_o = tl.dot(m_s, b_s, allow_tf32=False)
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_o_minus_s, (b_o - b_s).to(p_o_minus_s.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.jit
+def post_process_grad(
+    q,
+    k,
+    v,
+    u,
+    do,
+    dk,
+    dq,
+    du,
+    scale,
+    s_k_h,
+    s_k_t,
+    s_k_d,
+    s_v_h,
+    s_v_t,
+    s_v_d,
+    H,
+    T: tl.constexpr,
+    BT: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    i_h = i_bh % H
+
+    # Note that BK = tl.next_power_of_2(K), BV = tl.next_power_of_2(V)
+    p_q = tl.make_block_ptr(q + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT, 0), (BT, BK), (1, 0))
+    p_dq = tl.make_block_ptr(dq + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT, 0), (BT, BK), (1, 0))
+    p_k = tl.make_block_ptr(k + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT, 0), (BT, BK), (1, 0))
+    p_dk = tl.make_block_ptr(dk + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT, 0), (BT, BK), (1, 0))
+    p_du = tl.make_block_ptr(du + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT, 0), (BT, BK), (1, 0))
+    p_v = tl.make_block_ptr(v + i_bh * s_v_h, (T, V), (s_v_t, s_v_d), (i_t * BT, 0), (BT, BV), (1, 0))
+    p_do = tl.make_block_ptr(do + i_bh * s_v_h, (T, V), (s_v_t, s_v_d), (i_t * BT, 0), (BT, BV), (1, 0))
+    p_u = tl.make_block_ptr(u + i_h * K, (K,), (1,), (0,), (BK,), (0,))
+
+    b_q = tl.load(p_q, boundary_check=(0, 1))
+    b_k = tl.load(p_k, boundary_check=(0, 1))
+    b_v = tl.load(p_v, boundary_check=(0, 1))
+    b_do = tl.load(p_do, boundary_check=(0, 1))
+    b_u = tl.load(p_u, boundary_check=(0,))
+
+    b_vdo = tl.sum(b_v * b_do, axis=1)
+    b_du = b_vdo[:, None] * b_k * b_q * scale
+    b_dq = b_vdo[:, None] * b_k * b_u[None, :] * scale
+    b_dk = b_vdo[:, None] * b_q * b_u[None, :] * scale
+
+    b_dq += tl.load(p_dq, boundary_check=(0, 1))
+    tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
+
+    b_dk += tl.load(p_dk, boundary_check=(0, 1))
+    tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
+
+    tl.store(p_du, b_du.to(p_du.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.jit
+def chunk_rwkv6_fwd_kernel_h(
+    k,
+    v,
+    g,
+    h,
+    h0,
+    ht,
+    s_k_h,
+    s_k_t,
+    s_k_d,
+    s_v_h,
+    s_v_t,
+    s_v_d,
+    s_h_h,
+    s_h_t,
+    s_h_d,
+    T: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    NT: tl.constexpr,
+    USE_INITIAL_STATE: tl.constexpr,
+    STORE_FINAL_STATE: tl.constexpr
+):
+    i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    b_h = tl.zeros([BK, BV], dtype=tl.float32)
+
+    if USE_INITIAL_STATE:
+        p_h = tl.make_block_ptr(h0 + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+        b_h += tl.load(p_h, boundary_check=(0, 1)).to(tl.float32)
+    for i_t in range(NT):
+        o_t = min(i_t * BT + BT, T)
+
+        p_k = tl.make_block_ptr(k + i_bh * s_k_h, (K, T), (s_k_d, s_k_t), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
+        p_v = tl.make_block_ptr(v + i_bh * s_v_h, (T, V), (s_v_t, s_v_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        p_h = tl.make_block_ptr(h + i_bh * s_h_h + i_t * K * V, (K, V), (s_h_t, s_h_d), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+        p_g = tl.make_block_ptr(g + i_bh * s_k_h, (K, T), (s_k_d, s_k_t), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
+        p_gn = tl.make_block_ptr(g + i_bh * s_k_h, (T * K,), (s_k_d,), ((o_t - 1) * K + i_k * BK,), (BK,), (0,))
+
+        tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1))
+        # [BK, BT]
+        b_k = tl.load(p_k, boundary_check=(0, 1))
+        # [BT, BV]
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        # [BK, BT]
+        b_g = tl.load(p_g, boundary_check=(0, 1))
+        if i_t < NT - 1:
+            # [BK,]
+            b_gn = tl.load(p_gn, boundary_check=(0,))
+        else:
+            b_gn = tl.min(b_g, axis=1)
+        b_h *= tl.exp(b_gn)[:, None]
+        b_k = (b_k * tl.exp(b_gn[:, None] - b_g)).to(b_k.dtype)
+        b_h += tl.dot(b_k, b_v, allow_tf32=False)
+
+    if STORE_FINAL_STATE:
+        p_h = tl.make_block_ptr(ht + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+        tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.jit
+def chunk_rwkv6_fwd_kernel_intra(
+    q,
+    k,
+    g,
+    gs,
+    u,
+    A,
+    s_k_h,
+    s_k_t,
+    s_k_d,
+    scale,
+    H,
+    T: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    BC: tl.constexpr,
+    BK: tl.constexpr,
+    NC: tl.constexpr,
+    DK: tl.constexpr
+):
+    i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC
+    i_h = i_bh % H
+    n_bh = tl.num_programs(2)
+
+    o_k = i_k * BK + tl.arange(0, BK)
+    o_q = i_t * BT + i_i * BC
+    m_k = o_k < K
+
+    if i_i > i_j:
+        p_q = tl.make_block_ptr(q + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
+        p_k = tl.make_block_ptr(k + i_bh * s_k_h, (K, T), (s_k_d, s_k_t), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1))
+        p_gs = tl.make_block_ptr(gs + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
+        p_gk = tl.make_block_ptr(g + i_bh * s_k_h, (K, T), (s_k_d, s_k_t), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1))
+        p_A = tl.make_block_ptr(A + (i_k*n_bh+i_bh)*T*BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0))
+        # [BK,]
+        b_gn = tl.load(g + i_bh * T * K + (o_q - 1) * K + o_k, mask=(m_k & (i_i > 0) & (o_q <= T)), other=0)
+        # [BC, BK]
+        b_q = tl.load(p_q, boundary_check=(0, 1))
+        b_gs = tl.load(p_gs, boundary_check=(0, 1))
+        b_qg = (b_q * tl.exp(b_gs - b_gn[None, :]) * scale).to(b_q.dtype)
+        # [BK, BC]
+        b_k = tl.load(p_k, boundary_check=(0, 1))
+        b_gk = tl.load(p_gk, boundary_check=(0, 1))
+        b_kg = (b_k * tl.exp(b_gn[:, None] - b_gk)).to(b_k.dtype)
+        # [BC, BC]
+        b_A = tl.dot(b_qg, b_kg, allow_tf32=False)
+        tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1))
+    elif i_i == i_j:
+        p_q = tl.make_block_ptr(q + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
+        p_gs = tl.make_block_ptr(gs + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
+        p_k = tl.make_block_ptr(k + i_bh * s_k_h, (T * K,), (s_k_d,), ((i_t * BT + i_j * BC) * K + i_k * BK,), (BK,), (0,))
+        p_q_u = tl.make_block_ptr(q + i_bh * s_k_h, (T*K,), (s_k_d,), ((i_t * BT + i_j * BC) * K + i_k * BK,), (BK,), (0,))
+
+        # [BC, BK]
+        b_q = tl.load(p_q, boundary_check=(0, 1))
+        b_gs = tl.load(p_gs, boundary_check=(0, 1))
+        o_i = tl.arange(0, BC)
+        o_g = i_bh * T * K + (i_t * BT + i_j * BC) * K + o_k
+        o_A = (i_bh + i_k * n_bh) * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_j * BC
+        m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T
+        p_u = tl.make_block_ptr(u + i_h * DK, (DK,), (1,), (i_k * BK), (BK,), (0,))
+        b_u = tl.load(p_u, boundary_check=(0,))
+        for j in range(0, BC):
+            # [BK,]
+            b_k = tl.load(p_k, boundary_check=(0,)).to(tl.float32)
+            b_gk = tl.load(g + o_g + j * K, mask=(m_k & ((i_t * BT + i_j * BC + j) < T)), other=0).to(tl.float32)
+            # [BC,]
+            b_A = tl.sum(b_q * b_k[None, :] * tl.exp(b_gs - b_gk[None, :]) * scale, 1)
+            b_A = tl.where(o_i > j, b_A, 0.)
+            # self
+            b_q_u = tl.load(p_q_u, boundary_check=(0,)).to(tl.float32)
+            b_A_u = tl.sum(b_q_u * b_k * b_u * scale, axis=0)
+            m_u = tl.arange(0, BC) == j
+            b_A = tl.where(m_u, b_A_u, b_A)
+            tl.store(A + o_A + j, b_A.to(A.dtype.element_ty), mask=m_A)
+            p_k = tl.advance(p_k, (K,))
+            p_q_u = tl.advance(p_q_u, (K,))
+
+
+@triton.jit
+def chunk_rwkv6_fwd_kernel_inter(
+    q,
+    v,
+    gs,
+    h,
+    o,
+    A,
+    s_k_h,
+    s_k_t,
+    s_k_d,
+    s_v_h,
+    s_v_t,
+    s_v_d,
+    s_h_h,
+    s_h_t,
+    s_h_d,
+    scale,
+    T: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr
+):
+    i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+
+    b_o = tl.zeros([BT, BV], dtype=tl.float32)
+    for i_k in range(tl.cdiv(K, BK)):
+        p_q = tl.make_block_ptr(q + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+        p_gs = tl.make_block_ptr(gs + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+        p_h = tl.make_block_ptr(h + i_bh * s_h_h + i_t * K * V, (K, V), (s_h_t, s_h_d), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+
+        # [BT, BK]
+        b_q = tl.load(p_q, boundary_check=(0, 1))
+        b_q = (b_q * scale).to(b_q.dtype)
+        # [BT, BK]
+        b_gs = tl.load(p_gs, boundary_check=(0, 1))
+        # [BT, BK]
+        b_qg = (b_q * tl.exp(b_gs)).to(b_q.dtype)
+        # [BK, BV]
+        b_h = tl.load(p_h, boundary_check=(0, 1))
+        # works but dkw, owing to divine benevolence
+        # [BT, BV]
+        if i_k >= 0:
+            b_o += tl.dot(b_qg, b_h, allow_tf32=False)
+    p_v = tl.make_block_ptr(v + i_bh * s_v_h, (T, V), (s_v_t, s_v_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+    p_o = tl.make_block_ptr(o + i_bh * s_v_h, (T, V), (s_v_t, s_v_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+    p_A = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0))
+    # [BT, BV]
+    b_v = tl.load(p_v, boundary_check=(0, 1))
+    # [BT, BT]
+    b_A = tl.load(p_A, boundary_check=(0, 1))
+    b_o += tl.dot(b_A, b_v, allow_tf32=False)
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.jit
+def chunk_rwkv6_bwd_kernel_dh(
+    q,
+    g,
+    gs,
+    do,
+    dh,
+    dh0,
+    s_k_h,
+    s_k_t,
+    s_k_d,
+    s_v_h,
+    s_v_t,
+    s_v_d,
+    s_h_h,
+    s_h_t,
+    s_h_d,
+    scale,
+    T: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    NT: tl.constexpr,
+    USE_INITIAL_STATE: tl.constexpr
+):
+    i_k, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+
+    b_dh = tl.zeros([BK, BV], dtype=tl.float32)
+    for i_t in range(NT - 1, -1, -1):
+        o_t = min(i_t * BT + BT, T)
+
+        p_q = tl.make_block_ptr(q + i_bh * s_k_h, (K, T), (s_k_d, s_k_t), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
+        p_do = tl.make_block_ptr(do + i_bh * s_v_h, (T, V), (s_v_t, s_v_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        p_dh = tl.make_block_ptr(dh + i_bh * s_h_h + i_t * K*V, (K, V), (s_h_t, s_h_d), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+        p_gs = tl.make_block_ptr(gs + i_bh * s_k_h, (K, T), (s_k_d, s_k_t), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
+        p_gn = tl.make_block_ptr(g + i_bh * s_k_h, (T * K,), (s_k_d,), ((o_t - 1) * K + i_k * BK,), (BK,), (0,))
+
+        # [BK, BT]
+        b_q = tl.load(p_q, boundary_check=(0, 1))
+        b_q = (b_q * scale).to(b_q.dtype)
+        # [BT, BV]
+        b_do = tl.load(p_do, boundary_check=(0, 1))
+
+        tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1))
+
+        # [BK,]
+        b_gn = tl.load(p_gn, boundary_check=(0,))
+        # [BK, BV]
+        b_dh *= tl.exp(b_gn)[:, None]
+        # [BK, BT]
+        b_gs = tl.load(p_gs, boundary_check=(0, 1))
+        b_q = (b_q * tl.exp(b_gs)).to(b_q.dtype)
+
+        # [BK, BV]
+        b_dh += tl.dot(b_q, b_do, allow_tf32=False)
+
+    if USE_INITIAL_STATE:
+        p_dh0 = tl.make_block_ptr(dh0 + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+        tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.jit
+def chunk_rwkv6_bwd_kernel_inter(
+    k,
+    v,
+    h,
+    g,
+    gs,
+    A,
+    do,
+    dh,
+    dq,
+    dk,
+    dv,
+    dA,
+    s_k_h,
+    s_k_t,
+    s_k_d,
+    s_v_h,
+    s_v_t,
+    s_v_d,
+    s_h_h,
+    s_h_t,
+    s_h_d,
+    scale,
+    T: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr
+):
+    i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    n_bh = tl.num_programs(2)
+    o_t = min(i_t * BT + BT, T)
+
+    p_k = tl.make_block_ptr(k + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+    p_gk = tl.make_block_ptr(g + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+    p_gq = tl.make_block_ptr(gs + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+    p_gn = tl.make_block_ptr(g + i_bh * s_k_h, (T * K,), (s_k_d,), ((o_t - 1) * K + i_k * BK,), (BK,), (0,))
+    p_A = tl.make_block_ptr(A + i_bh * T * BT, (BT, T), (1, BT), (0, i_t * BT), (BT, BT), (0, 1))
+
+    # [BT, BK]
+    b_k = tl.load(p_k, boundary_check=(0, 1))
+    b_gk = tl.load(p_gk, boundary_check=(0, 1))
+    b_gq = tl.load(p_gq, boundary_check=(0, 1))
+    b_gn = tl.exp(tl.load(p_gn, boundary_check=(0,))[None, :] - b_gk)
+    b_k = (b_k * b_gn).to(b_k.dtype)
+    # [BT, BT]
+    b_A = tl.load(p_A, boundary_check=(0, 1))
+
+    b_dq = tl.zeros([BT, BK], dtype=tl.float32)
+    b_dk = tl.zeros([BT, BK], dtype=tl.float32)
+    b_dA = tl.zeros([BT, BT], dtype=tl.float32)
+    for i_v in range(tl.cdiv(V, BV)):
+        p_v = tl.make_block_ptr(v + i_bh * s_v_h, (T, V), (s_v_t, s_v_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        p_h = tl.make_block_ptr(h + i_bh * s_h_h + i_t * V * K, (V, K), (s_h_d, s_h_t), (i_v * BV, i_k * BK), (BV, BK), (0, 1))
+        p_do = tl.make_block_ptr(do + i_bh * s_v_h, (T, V), (s_v_t, s_v_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        p_dh = tl.make_block_ptr(dh + i_bh * s_h_h + i_t * K*V, (K, V), (s_h_t, s_h_d), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+        p_dv = tl.make_block_ptr(dv + (i_k*n_bh+i_bh) * s_v_h, (T, V), (s_v_t, s_v_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+
+        # [BT, BV]
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        # [BV, BK]
+        b_h = tl.load(p_h, boundary_check=(0, 1))
+        # [BT, BV]
+        b_do = tl.load(p_do, boundary_check=(0, 1))
+        # [BK, BV]
+        b_dh = tl.load(p_dh, boundary_check=(0, 1))
+
+        # [BT, BV]
+        b_dv = tl.dot(b_k, b_dh, allow_tf32=False)
+        if i_k == 0:
+            b_dv += tl.dot(b_A, b_do, allow_tf32=False)
+        b_do = (b_do * scale).to(b_do.dtype)
+        tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))
+        # [BT, BT]
+        b_dA += tl.dot(b_do, tl.trans(b_v), allow_tf32=False)
+        # [BT, BK]
+        b_dq += tl.dot(b_do, b_h, allow_tf32=False)
+        # [BT, BK]
+        b_dk += tl.dot(b_v, tl.trans(b_dh).to(b_v.dtype), allow_tf32=False)
+
+    b_dq = b_dq * tl.exp(b_gq)
+    b_dk = b_dk * b_gn
+
+    p_dq = tl.make_block_ptr(dq + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+    p_dk = tl.make_block_ptr(dk + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+    p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT, ), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0))
+    tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
+
+    o_i = tl.arange(0, BT)
+    m_s = o_i[:, None] > o_i[None, :]
+    # [BT, BT]
+    b_dA = tl.where(m_s, b_dA, 0.).to(b_k.dtype)
+    if i_k == 0:
+        tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.jit
+def chunk_rwkv6_bwd_kernel_intra(
+    q,
+    k,
+    g,
+    gs,
+    dA,
+    dq,
+    dk,
+    s_k_h,
+    s_k_t,
+    s_k_d,
+    T: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    BC: tl.constexpr,
+    BK: tl.constexpr,
+    NC: tl.constexpr
+):
+    i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_t, i_i = i_c // NC, i_c % NC
+
+    o_k = i_k * BK + tl.arange(0, BK)
+    o_q = i_t * BT + i_i * BC
+    m_k = o_k < K
+
+    p_gs = tl.make_block_ptr(gs + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
+    # [BK,]
+    b_gn = tl.load(g + i_bh * T * K + (o_q - 1) * K + o_k, mask=(m_k & (i_i > 0) & (o_q <= T)), other=0)
+    # [BC, BK]
+    b_gs = tl.load(p_gs, boundary_check=(0, 1))
+    b_dq = tl.zeros([BC, BK], dtype=tl.float32)
+    for i_j in range(0, i_i):
+        p_k = tl.make_block_ptr(k + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0))
+        p_gk = tl.make_block_ptr(g + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0))
+        p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0))
+        # [BC, BK]
+        b_k = tl.load(p_k, boundary_check=(0, 1))
+        b_gk = tl.load(p_gk, boundary_check=(0, 1))
+        b_kg = (b_k * tl.exp(b_gn[None, :] - b_gk)).to(b_k.dtype)
+        # [BC, BC]
+        b_dA = tl.load(p_dA, boundary_check=(0, 1))
+        # [BC, BK]
+        b_dq += tl.dot(b_dA, b_kg, allow_tf32=False)
+    b_dq *= tl.exp(b_gs - b_gn[None, :])
+
+    o_i = tl.arange(0, BC)
+    o_dA = i_bh * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_i * BC
+    m_dA = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T
+
+    for j in range(0, BC):
+        p_kj = tl.make_block_ptr(k + i_bh * s_k_h, (T * K,), (1,), ((i_t * BT + i_i*BC+j) * K + i_k * BK,), (BK,), (0,))
+
+        # [BC,]
+        b_dA = tl.load(dA + o_dA + j, mask=m_dA, other=0)
+        # [BK,]
+        b_kj = tl.load(p_kj, boundary_check=(0,)).to(tl.float32)
+        b_gkj = tl.load(g + i_bh * T * K + (o_q + j) * K + o_k, mask=(m_k & ((o_q + j) < T)), other=0)
+        # [BC, BK]
+        m_i = o_i[:, None] > j
+        # [BC, BK]
+        b_dq += tl.where(m_i, b_dA[:, None] * b_kj[None, :] * tl.exp(b_gs - b_gkj[None, :]), 0.)
+
+    p_dq = tl.make_block_ptr(dq + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
+
+    b_dq = b_dq + tl.load(p_dq, boundary_check=(0, 1))
+    tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
+
+    tl.debug_barrier()
+    p_k = tl.make_block_ptr(k + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
+    p_gk = tl.make_block_ptr(g + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
+    p_gn = tl.make_block_ptr(g + i_bh * s_k_h, (T*K,), (s_k_d,), ((i_t * BT + i_i * BC + BC - 1) * K + i_k * BK,), (BK,), (0,))
+    # [BK,]
+    b_gn = tl.load(p_gn, boundary_check=(0,))
+    # [BC, BK]
+    b_k = tl.load(p_k, boundary_check=(0, 1))
+    b_gk = tl.load(p_gk, boundary_check=(0, 1))
+    b_dk = tl.zeros([BC, BK], dtype=tl.float32)
+    for i_j in range(i_i + 1, NC):
+        p_q = tl.make_block_ptr(q + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0))
+        p_gs = tl.make_block_ptr(gs + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0))
+        p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT + i_j * BC, i_i * BC), (BC, BC), (1, 0))
+        # [BC, BK]
+        b_q = tl.load(p_q, boundary_check=(0, 1))
+        b_gs = tl.load(p_gs, boundary_check=(0, 1))
+        b_qg = (b_q * tl.exp(b_gs - b_gn[None, :])).to(b_q.dtype)
+        # [BC, BC]
+        b_dA = tl.load(p_dA, boundary_check=(0, 1))
+        # [BC, BK]
+        b_dk += tl.dot(tl.trans(b_dA), b_qg, allow_tf32=False)
+    b_dk *= tl.exp(b_gn[None, :] - b_gk)
+
+    o_dA = i_bh * T * BT + (i_t * BT + i_i * BC) * BT + i_i * BC + tl.arange(0, BC)
+    for j in range(0, BC):
+        p_qj = tl.make_block_ptr(q + i_bh * s_k_h, (T * K,), (1,), ((i_t * BT + i_i * BC + j) * K + i_k * BK,), (BK,), (0,))
+        p_gqj = tl.make_block_ptr(gs + i_bh * s_k_h, (T * K,), (1,), ((i_t * BT + i_i * BC + j) * K + i_k * BK,), (BK,), (0,))
+        # [BC,]
+        b_dA = tl.load(dA + o_dA + j * BT, mask=(i_t * BT + i_i * BC + j < T), other=0)
+        # [BK,]
+        b_qj = tl.load(p_qj, boundary_check=(0,)).to(tl.float32)
+        b_gqj = tl.load(p_gqj, boundary_check=(0,)).to(tl.float32)
+        # [BC, BK]
+        m_i = o_i[:, None] < j
+        b_dk += tl.where(m_i, b_dA[:, None] * b_qj[None, :] * tl.exp(b_gqj[None, :] - b_gk), 0.)
+
+    p_dk = tl.make_block_ptr(dk + i_bh * s_k_h, (T, K), (s_k_t, s_k_d), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0))
+    b_dk = b_dk + tl.load(p_dk, boundary_check=(0, 1))
+    tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
+
+
+class ChunkRWKV6Function(torch.autograd.Function):
+
+    @staticmethod
+    @contiguous
+    def forward(ctx, r, k, v, g, u, scale, initial_state, output_final_state, checkpoint_level):
+        q = r  # alias
+        B, H, T, K, V = *q.shape, v.shape[-1]
+        BT, BC = 64, 16
+        BK = min(64, triton.next_power_of_2(K))
+        BV = min(64, triton.next_power_of_2(V))
+        NT, NC = triton.cdiv(T, BT), triton.cdiv(BT, BC)
+        NK = triton.cdiv(K, BK)
+        NV = triton.cdiv(V, BV)
+        num_warps = 4 if BK == 64 else 2
+        num_stages = 1
+
+        def fwd_inner(q, k, v, g, B, H, T, K, V, BT, BK, BV, NT, h0=None, ht=None):
+            NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
+            h = q.new_empty(B, H, NT * K, V)
+            grid = (NV, NK, B * H)
+            chunk_rwkv6_fwd_kernel_h[grid](
+                k, v, g, h, h0, ht,
+                k.stride(1), k.stride(2), k.stride(3),
+                v.stride(1), v.stride(2), v.stride(3),
+                h.stride(1), h.stride(2), h.stride(3),
+                T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT,
+                USE_INITIAL_STATE=h0 is not None,
+                STORE_FINAL_STATE=ht is not None,
+                num_warps=num_warps,
+                num_stages=num_stages
+            )
+            return h
+
+        final_state = None
+        if output_final_state:
+            final_state = q.new_empty(B, H, K, V, dtype=torch.float)
+
+        g_org, g, gs = g, torch.empty_like(g, dtype=torch.float), torch.empty_like(g, dtype=torch.float)
+        def grid(meta): return ((triton.cdiv(meta['S'], meta['BS']), NT, B * H))
+        # keep cummulative normalizer in fp32
+        # this kernel is equivalent to
+        # g_org = g_org.view(B, H, NT, BT, -1)
+        # g = g_org.cumsum(-2).view(B, H, T, -1)
+        # gs = g - g_org
+        chunk_rwkv6_fwd_kernel_cum[grid](
+            g_org, g, gs,
+            g.stride(1), g.stride(2), g.stride(3),
+            T=T, S=K, BT=BT
+        )
+        h = fwd_inner(
+            q=q, k=k, v=v, g=g,
+            B=B, H=H, T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT,
+            h0=initial_state if initial_state is not None else None,
+            ht=final_state if final_state is not None else None
+        )
+        A = q.new_zeros(NK, B, H, T, BT)
+        grid = (NK, NT * NC * NC, B * H)
+        chunk_rwkv6_fwd_kernel_intra[grid](
+            q, k, g, gs, u, A,
+            k.stride(1), k.stride(2), k.stride(3),
+            scale,
+            H=H, T=T, K=K, BT=BT, BC=BC, BK=BK, NC=NC, DK=K,
+            num_warps=num_warps,
+            num_stages=num_stages
+        )
+        A = A.sum(0, dtype=A.dtype)
+        o = torch.empty_like(v)
+
+        grid = (NV, NT, B * H)
+        chunk_rwkv6_fwd_kernel_inter[grid](
+            q, v, gs, h, o, A,
+            k.stride(1), k.stride(2), k.stride(3),
+            v.stride(1), v.stride(2), v.stride(3),
+            h.stride(1), h.stride(2), h.stride(3),
+            scale,
+            T=T, K=K, V=V, BT=BT, BK=BK, BV=BV,
+            num_warps=num_warps,
+            num_stages=num_stages
+        )
+
+        if checkpoint_level > 1:
+            del h
+            h, initial_state = None, None
+        del g, gs
+        ctx.save_for_backward(q, k, v, g_org, u, h, initial_state, A)
+        ctx.BT = BT
+        ctx.scale = scale
+        ctx.checkpoint_level = checkpoint_level
+        return o, final_state
+
+    @staticmethod
+    @contiguous
+    def backward(ctx, do, dht=None):
+        q, k, v, g, u, h, initial_state, A = ctx.saved_tensors
+        B, H, T, K, V = *q.shape, v.shape[-1]
+        BT, BC = ctx.BT, 16
+        BK = min(64, triton.next_power_of_2(K))
+        BV = min(64, triton.next_power_of_2(V))
+        NT, NC = triton.cdiv(T, BT), triton.cdiv(BT, BC)
+        NK = triton.cdiv(K, BK)
+        num_warps = 4 if BK == 64 else 2
+        num_stages = 1
+
+        def fwd_inner(q, k, v, g, B, H, T, K, V, BT, BK, BV, NT, h0=None, ht=None):
+            NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
+            h = q.new_empty(B, H, NT * K, V)
+            grid = (NV, NK, B * H)
+            chunk_rwkv6_fwd_kernel_h[grid](
+                k, v, g, h, h0, ht,
+                k.stride(1), k.stride(2), k.stride(3),
+                v.stride(1), v.stride(2), v.stride(3),
+                h.stride(1), h.stride(2), h.stride(3),
+                T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT,
+                USE_INITIAL_STATE=h0 is not None,
+                STORE_FINAL_STATE=ht is not None,
+                num_warps=num_warps,
+                num_stages=num_stages
+            )
+            return h
+
+        def bwd_inner(q, g, gs, h0, do, B, H, T, K, V, BT, BK, BV, NT, scale):
+            NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
+            dh = q.new_empty(B, H, NT * K, V)
+            dh0 = torch.empty_like(h0) if h0 is not None else None
+            grid = (NK, NV, B * H)
+            chunk_rwkv6_bwd_kernel_dh[grid](
+                q, g, gs, do, dh, dh0,
+                q.stride(1), q.stride(2), q.stride(3),
+                do.stride(1), do.stride(2), do.stride(3),
+                dh.stride(1), dh.stride(2), dh.stride(3),
+                scale,
+                T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT,
+                USE_INITIAL_STATE=h0 is not None,
+                num_warps=num_warps,
+                num_stages=num_stages
+            )
+            return dh, dh0
+
+        # recompute cumulative log decays.
+        g_org, g, gs = g, torch.empty_like(g, dtype=torch.float), torch.empty_like(g, dtype=torch.float)
+        def grid(meta): return ((triton.cdiv(meta['S'], meta['BS']), NT, B * H))
+        # keep cummulative normalizer in fp32
+        # this kernel is equivalent to
+        # g = g.view(B, H, NT, BT, -1).cumsum(-2).view(B, H, T, -1)
+        chunk_rwkv6_fwd_kernel_cum[grid](
+            g_org, g, gs,
+            g.stride(1), g.stride(2), g.stride(3),
+            T=T, S=K, BT=BT
+        )
+
+        # rerun the forward pass to get h if checkpoint_level >= 1
+        if ctx.checkpoint_level == 1:
+            h = fwd_inner(
+                q=q, k=k, v=v, g=g,
+                B=B, H=H, T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT,
+                h0=initial_state if initial_state is not None else None,
+                ht=None
+            )
+
+        scale = ctx.scale
+        # g, gs: torch.float32
+        dh, dh0 = bwd_inner(
+            q.to(torch.float), g, gs, initial_state, do.to(torch.float),
+            B=B, H=H, T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT,
+            scale=scale
+        )
+        dh = dh.to(q)
+        if initial_state is not None:
+            dh0 = dh0.to(q)
+        dq = torch.empty_like(q, dtype=torch.float)
+        dk = torch.empty_like(k, dtype=torch.float)
+        dv = v.new_empty(NK, *v.shape)
+        dA = q.new_zeros(B, H, T, BT)
+        grid = (NK, NT, B * H)
+        chunk_rwkv6_bwd_kernel_inter[grid](
+            k, v, h, g, gs, A, do, dh, dq, dk, dv, dA,
+            k.stride(1), k.stride(2), k.stride(3),
+            v.stride(1), v.stride(2), v.stride(3),
+            h.stride(1), h.stride(2), h.stride(3),
+            scale,
+            T=T, K=K, V=V, BT=BT, BK=BK, BV=BV,
+            num_warps=num_warps,
+            num_stages=num_stages
+        )
+        dv = dv.sum(0, dtype=dv.dtype)
+        grid = (NK, NT * NC, B * H)
+        chunk_rwkv6_bwd_kernel_intra[grid](
+            q, k, g, gs, dA, dq, dk,
+            k.stride(1), k.stride(2), k.stride(3),
+            T=T, K=K, BT=BT, BC=BC, BK=BK, NC=NC,
+            num_warps=num_warps,
+            num_stages=num_stages
+        )
+
+        # TODO: fuse?
+        dg = (dq * q)[:, :, 1:] - (dk * k)[:, :, 0:-1]
+        dg = torch.nn.functional.pad(dg, (0, 0, 0, 1, 0, 0, 0, 0), value=0)
+        dg = chunk_global_reversed_cumsum(dg).to(g)
+        # equivalent to the following pytorch code.
+        # du = ((do * v).sum(-1)[..., None] * k * q * scale).sum(-2).to(u)
+        # dq += ((do * v).sum(-1)[..., None] * k * scale * u[:, :, None, :])
+        # dk += ((do * v).sum(-1)[..., None] * q * scale * u[:, :, None, :])
+        BT = 64
+        grid = (triton.cdiv(T, BT), B * H)
+        du = torch.empty_like(g, dtype=torch.float)
+        post_process_grad[grid](
+            q, k, v, u, do, dk, dq, du, scale,
+            q.stride(1), q.stride(2), q.stride(3),
+            v.stride(1), v.stride(2), v.stride(3), H=H,
+            T=T, BT=BT, K=K, V=V, BK=triton.next_power_of_2(K), BV=triton.next_power_of_2(V),
+            num_warps=4
+        )
+        du = du.sum([0, 2])
+        return dq.to(q), dk.to(k), dv.to(v), dg.to(g), du.to(u), None, dh0, None, None
+
+
+def chunk_rwkv6(
+    r: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    u: torch.Tensor,
+    scale: Optional[int] = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+    checkpoint_level: Optional[int] = 0
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    r"""
+    Args:
+        r (torch.Tensor):
+            reception of shape `(B, H, T, K)`. Alias: q, query in linear attention.
+        k (torch.Tensor):
+            keys of shape `(B, H, T, K)`
+        v (torch.Tensor):
+            values of shape `(B, H, T, V)`
+        w (torch.Tensor):
+            data-dependent decays of shape `(B, H, T, K)` in log space! Alias: g.
+        u (torch.Tensor):
+            bonus of shape `(H, K)`
+        scale (Optional[int]):
+            Scale factor for the RWKV6 attention scores.
+            If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
+        initial_state (Optional[torch.Tensor]):
+            Initial state of shape `(B, H, K, V)`. Default: `None`.
+        output_final_state (Optional[bool]):
+            Whether to output the final state of shape `(B, H, K, V)`. Default: `False`.
+        checkpoint_level (Optional[int]):
+            Checkpointing level; higher values will save more memories and do more recomputations during backward.
+            Default: `0`:
+            - Level `0`: store forward hidden states for backprop.
+            - Level `1`: recompute the forward hidden states during backward.
+    """
+    assert checkpoint_level in [0, 1]
+    if scale is None:
+        scale = r.shape[-1] ** -0.5
+    o, final_state = ChunkRWKV6Function.apply(r, k, v, g, u, scale, initial_state, output_final_state, checkpoint_level)
+    return o, final_state
+
+
+if __name__ == "__main__":
+    import torch.nn.functional as F
+
+    from fla.ops.rwkv6.recurrent_fuse import fused_recurrent_rwkv6
+    B = 8
+    H = 4
+    L = 1024
+    K = 100
+    V = 120
+
+    torch.manual_seed(0)
+    dtype = torch.float
+    q = torch.randn(B, H, L, K).cuda().to(dtype).requires_grad_(True)
+    k = torch.randn(B, H, L, K).cuda().to(dtype).requires_grad_(True)
+    v = torch.randn(B, H, L, V).cuda().to(dtype).requires_grad_(True)
+    w = (-torch.randn(B, H, L, K).exp()).cuda().requires_grad_(True)
+    u = torch.randn(H, K).cuda().to(dtype).requires_grad_(True)
+    h0 = torch.randn(B, H, K, V).cuda().to(dtype).requires_grad_(True)
+    do = torch.rand_like(v).cuda()
+    o, ht = fused_recurrent_rwkv6(q, k, v, w, u, initial_state=h0, output_final_state=True)
+    o.backward(do)
+    dq, q.grad = q.grad.clone(), None
+    dk, k.grad = k.grad.clone(), None
+    dv, v.grad = v.grad.clone(), None
+    dw, w.grad = w.grad.clone(), None
+    du, u.grad = u.grad.clone(), None
+    dh0, h0.grad = h0.grad.clone(), None
+    o2, ht2 = chunk_rwkv6(q, k, v, w, u, initial_state=h0, output_final_state=True)
+    o2.backward(do)
+    torch.testing.assert_close(o, o2, rtol=0, atol=1e-4)
+    torch.testing.assert_close(ht, ht2, rtol=0, atol=1e-4)
+    torch.testing.assert_close(q.grad, dq, rtol=0, atol=1e-4)
+    torch.testing.assert_close(k.grad, dk, rtol=0, atol=1e-4)
+    torch.testing.assert_close(v.grad, dv, rtol=0, atol=1e-4)
+    torch.testing.assert_close(w.grad, dw, rtol=0, atol=1e-4)
+    torch.testing.assert_close(u.grad, du, rtol=0, atol=2e-4)
+    torch.testing.assert_close(h0.grad, dh0, rtol=0, atol=2e-4)
+
+    print("All tests passed!")
+
+    @triton.testing.perf_report(
+        triton.testing.Benchmark(
+            # argument names to use as an x-axis for the plot
+            x_names=['T'],
+            # different possible values for `x_name`
+            x_vals=[128 * 2 ** i for i in range(0, 8)],
+            # argument name whose value corresponds to a different line in the plot
+            line_arg='provider',
+            # possible values for `line_arg``
+            line_vals=['recurrent', 'chunk', 'recurrent_bwd', 'chunk_bwd'],
+            # label name for the lines
+            line_names=['recurrent', 'chunk', 'recurrent_bwd', 'chunk_bwd'],
+            # line styles
+            styles=[('green', '-'), ('blue', '--'), ('red', '-.'), ('cyan', ':'), ('yellow', 'dotted'), ('black', 'dashed')],
+            ylabel="Execution Time (ms)",  # label name for the y-axis
+            # name for the plot. Used also as a file name for saving the plot.
+            plot_name="Performance",
+            args={},
+        )
+    )
+    def benchmark(T, provider):
+        device = 'cuda'
+        dtype = torch.bfloat16
+        requires_grad = True
+        B, H, K = 16, 4, 128
+
+        q = torch.randn(B, H, T, K, device=device, requires_grad=requires_grad, dtype=dtype)
+        k = torch.randn(B, H, T, K, device=device, requires_grad=requires_grad, dtype=dtype)
+        v = torch.randn(B, H, T, K, device=device, requires_grad=requires_grad, dtype=dtype)
+        w = F.logsigmoid(torch.randn(B, H, T, K)).to(dtype=dtype, device=device).requires_grad_(True)
+        u = torch.randn(H, K, device=device, requires_grad=requires_grad, dtype=dtype)
+
+        do = torch.ones_like(q, dtype=dtype)
+        quantiles = [0.5, 0.2, 0.8]
+        results = 0, 0, 0
+        if provider == 'recurrent':
+            results = triton.testing.do_bench(lambda: fused_recurrent_rwkv6(q, k, v, w, u), quantiles=quantiles)
+        if provider == 'chunk':
+            results = triton.testing.do_bench(lambda: chunk_rwkv6(q, k, v, w, u), quantiles=quantiles)
+        if provider == 'recurrent_bwd':
+            results = triton.testing.do_bench(lambda: fused_recurrent_rwkv6(q, k, v, w, u)
+                                              [0].backward(do), quantiles=quantiles)
+        if provider == 'chunk_bwd':
+            results = triton.testing.do_bench(lambda: chunk_rwkv6(q, k, v, w, u)[0].backward(do), quantiles=quantiles)
+        return results
+    benchmark.run(print_data=True)

fla2/ops/rwkv6/chunk_naive.py

@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+
+import torch
+from einops import rearrange
+
+
+def naive_chunk_rwkv6(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    w: torch.Tensor,
+    u: torch.Tensor,
+    chunk_size: int = 32
+):
+    assert q.shape[-2] % chunk_size == 0
+    orig_dtype = q.dtype
+    num_chunk = q.shape[-2] // chunk_size
+    u = u.unsqueeze(0)
+
+    q, k, v, w = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size).float(), (q, k, v, w))
+
+    w_cumsum = w.cumsum(-2)
+
+    kw = k * (w_cumsum[..., -1, None, :] - w_cumsum).exp()
+    wkv = kw.transpose(-1, -2) @ v
+
+    wkv_new = torch.zeros_like(wkv)
+
+    for i in range(num_chunk - 1):
+        wkv_new[:, :, i+1] = (wkv_new[:, :, i] * w_cumsum[:, :, i, -1, :, None].exp()) + wkv[:, :, i]
+
+    o_inter = torch.einsum('b h n d p, b h n c d -> b h n c p', wkv_new, (q * (w_cumsum - w).exp()))
+
+    o_intra = torch.zeros_like(o_inter)
+    for i in range(chunk_size):
+        attn = (q[:, :, :, i, None] * k * (w_cumsum[:, :, :, i, None] - w[:, :, :, i, None] - w_cumsum).exp()).sum(-1)
+        mask = (torch.arange(0, chunk_size) < i).to(attn.device)
+        attn.masked_fill_(~mask, 0)
+        intra_inter_o = (attn.unsqueeze(-1) * v).sum(-2)
+        intra_intra_o = (q[:, :, :, i] * u.unsqueeze(2) * k[:, :, :, i]).sum(-1).unsqueeze(-1) * v[:, :, :, i]
+        o_intra[:, :, :, i] = intra_inter_o + intra_intra_o
+    o = o_inter + o_intra
+    return rearrange(o, 'b h n c d -> b h (n c) d').to(orig_dtype)

fla2/ops/rwkv6/recurrent_fuse.py

@@ -0,0 +1,368 @@
+# -*- coding: utf-8 -*-
+
+# Copyright (c) 2024, Songlin Yang
+
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from fla.ops.utils import chunk_global_reversed_cumsum
+from fla.utils import autocast_custom_bwd, autocast_custom_fwd, contiguous
+
+
+@triton.jit
+def fused_recurrent_rwkv6_fwd_kernel(
+    q,  # query [B, H, T, K]
+    k,  # key [B, H, T, K]
+    v,  # value [B, H, T, V]
+    w,  # log gate [B, H, T, K]
+    u,  # bonus [B, H, K]
+    o,  # output [B, H, T, V]
+    # initial hidden state initialization [B, H, K, V]
+    h0,
+    ht,  # final hidden state [B, H, K, V]
+    s_k_h,  # stride size: T * K
+    s_v_h,  # stride size: T * V
+    scale,  # K ** -0.5
+    B: tl.constexpr,
+    H: tl.constexpr,
+    T: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BK: tl.constexpr,  # BLOCK SIZE along the K dimension
+    BV: tl.constexpr,  # BLOCK SIZE along the V dimension
+    USE_INITIAL_STATE: tl.constexpr,  # whether to use initial state
+    STORE_FINAL_STATE: tl.constexpr,  # whether to store final state
+    REVERSE: tl.constexpr,  # whether to do autoregressive modeling in the reverse direction
+):
+    i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_h = i_bh % H
+
+    p_q = q + i_bh * s_k_h + i_k * BK + tl.arange(0, BK) + ((T-1) * K if REVERSE else 0)
+    p_k = k + i_bh * s_k_h + i_k * BK + tl.arange(0, BK) + ((T-1) * K if REVERSE else 0)
+    p_v = v + i_bh * s_v_h + i_v * BV + tl.arange(0, BV) + ((T-1) * V if REVERSE else 0)
+    p_o = o + (i_bh + i_k * B * H) * s_v_h + i_v * BV + tl.arange(0, BV) + ((T-1) * V if REVERSE else 0)
+    p_w = w + i_bh * s_k_h + i_k * BK + tl.arange(0, BK) + ((T-1) * K if REVERSE else 0)
+    p_u = u + i_h * K + tl.arange(0, BK) + i_k * BK
+
+    mask_bk = (i_k * BK + tl.arange(0, BK)) < K
+    mask_bv = (i_v * BV + tl.arange(0, BV)) < V
+    mask_kv = mask_bv[:, None] & mask_bk[None, :]
+
+    b_h = tl.zeros([BV, BK], dtype=tl.float32)
+    if USE_INITIAL_STATE:
+        p_h0 = h0 + i_bh * K * V + (i_k * BK + tl.arange(0, BK)[None, :]) * V + (i_v * BV + tl.arange(0, BV)[:, None])
+        b_h += tl.load(p_h0, mask=mask_kv, other=0).to(tl.float32)
+
+    b_u = tl.load(p_u, mask=mask_bk, other=0).to(tl.float32)
+    for _ in range(0, T):
+        b_k = tl.load(p_k, mask=mask_bk, other=0).to(tl.float32)
+        b_v = tl.load(p_v, mask=mask_bv, other=0).to(tl.float32)
+        b_q = tl.load(p_q, mask=mask_bk, other=0).to(tl.float32) * scale
+        b_w = tl.load(p_w, mask=mask_bk, other=0).to(tl.float32)
+        b_w = tl.exp(b_w)
+        b_kv = b_k[None, :] * b_v[:, None]
+        b_o = (b_h + b_kv * b_u[None, :]) * b_q[None, :]
+        b_o = tl.sum(b_o, axis=1)
+        b_h = b_h * b_w[None, :]
+        b_h += b_kv
+        tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_bv)
+        p_q += -K if REVERSE else K
+        p_k += -K if REVERSE else K
+        p_o += -V if REVERSE else V
+        p_v += -V if REVERSE else V
+        p_w += -K if REVERSE else K
+
+    if STORE_FINAL_STATE:
+        p_ht = ht + i_bh * K * V + (i_k * BK + tl.arange(0, BK)[None, :]) * V + (i_v * BV + tl.arange(0, BV)[:, None])
+        tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_kv)
+
+
+# Similar to Algorithm1 of https://arxiv.org/abs/2006.16236
+@triton.jit
+def fused_recurrent_rwkv6_bwd_kernel_dq(
+    # B: B, H: H, T: T, D: d_head
+    # NV: number of split in the V dimension. NK: number of split in the K dimension
+    k,  # key [B, H, T, V]
+    v,  # value [B, H, T, V]
+    w,  # log gate [B, H, T, K]
+    u,  # bonus [B, H, K]
+
+    do,  # gradient of output [B, H, T, V]
+    dq,  # gradient of query [NV, B, H, T, K]
+    dq_aux,  # gradient of query_aux [NV, B, H, T, K]
+
+    # initial hidden state initialization [B, H, K, V]
+    h0,
+
+    s_k_h,  # stride size: T * K
+    s_v_h,  # stride size: T * V
+
+    scale,  # K ** -0.5
+    B: tl.constexpr,  # B
+    H: tl.constexpr,  # H
+    T: tl.constexpr,  # T
+    K: tl.constexpr,  # K
+    V: tl.constexpr,  # V
+    BK: tl.constexpr,  # BLOCK SIZE along the K dimension
+    BV: tl.constexpr,  # BLOCK SIZE along the V dimension
+    USE_INITIAL_STATE: tl.constexpr,  # whether to use initial state
+    REVERSE: tl.constexpr,  # whether to do autoregressive modeling in the reverse direction
+):
+    i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_h = i_bh % H
+    p_k = k + i_bh * s_k_h + i_k * BK + tl.arange(0, BK) + ((T-1) * K if REVERSE else 0)
+    p_v = v + i_bh * s_v_h + i_v * BV + tl.arange(0, BV) + ((T-1) * V if REVERSE else 0)
+    p_do = do + i_bh * s_v_h + i_v * BV + tl.arange(0, BV) + ((T-1) * V if REVERSE else 0)
+    p_dq = dq + (i_bh + i_v * B * H) * s_k_h + i_k * BK + tl.arange(0, BK) + ((T-1) * K if REVERSE else 0)
+    p_dq_aux = dq_aux + (i_bh + i_v * B * H) * s_k_h + i_k * BK + tl.arange(0, BK) + ((T-1) * K if REVERSE else 0)
+    p_w = w + i_bh * s_k_h + i_k * BK + tl.arange(0, BK) + ((T-1) * K if REVERSE else 0)
+    p_u = u + i_h * K + tl.arange(0, BK) + i_k * BK
+
+    mask_bk = i_k * BK + tl.arange(0, BK) < K
+    mask_bv = i_v * BV + tl.arange(0, BV) < V
+    mask_kv = mask_bv[:, None] & mask_bk[None, :]
+    b_u = tl.load(p_u, mask=mask_bk, other=0).to(tl.float32)
+    b_h = tl.zeros([BV, BK], dtype=tl.float32)
+
+    if USE_INITIAL_STATE:
+        p_h0 = h0 + i_bh * K * V + (i_k * BK + tl.arange(0, BK)[None, :]) * V + (i_v * BV + tl.arange(0, BV)[:, None])
+        b_h += tl.load(p_h0, mask=mask_kv, other=0).to(tl.float32)
+
+    for _ in range(0, T):
+        b_k = tl.load(p_k, mask=mask_bk, other=0).to(tl.float32)
+        b_v = tl.load(p_v, mask=mask_bv, other=0).to(tl.float32)
+        b_kv = b_k[None, :] * b_v[:, None]
+        b_do = tl.load(p_do, mask=mask_bv, other=0).to(tl.float32)
+        b_w = tl.load(p_w, mask=mask_bk, other=0).to(tl.float32)
+        b_w = tl.exp(b_w)
+        h_q = b_h * b_do[:, None]
+        b_dq = tl.sum(h_q + b_kv * b_u[None, :] * b_do[:, None], axis=0)
+        b_dq *= scale
+        b_dq_aux = tl.sum(h_q, axis=0)
+        b_h = b_h * b_w[None, :]
+        b_h += b_kv
+        tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=mask_bk)
+        tl.store(p_dq_aux, b_dq_aux.to(p_dq_aux.dtype.element_ty), mask=mask_bk)
+        p_k += -K if REVERSE else K
+        p_do += -V if REVERSE else V
+        p_v += -V if REVERSE else V
+        p_w += -K if REVERSE else K
+        p_dq += -K if REVERSE else K
+        p_dq_aux += -K if REVERSE else K
+
+
+@triton.jit
+def fused_recurrent_rwkv6_bwd_kernel_dkv(
+    # B: B, H: H, T: T, D: d_head
+    # NV: number of split in the V dimension. NK: number of split in the K dimension
+    q,  # query [B, H, T, K]
+    k,  # key [B, H, T, V]
+    v,  # value [B, H, T, V]
+    w,  # log gate [B, H, T, K]
+    u,  # bonus [B, H, K]
+
+    do,  # gradient of output [B, H, T, V]
+    dk,
+    dk_aux,
+    dv,
+    dh0,
+
+    # initial hidden state initialization [B, H, K, V]
+    s_k_h,  # stride size: T * K
+    s_v_h,  # stride size: T * V
+
+    scale,  # K ** -0.5
+    B: tl.constexpr,  # B
+    H: tl.constexpr,  # H
+    T: tl.constexpr,  # T
+    K: tl.constexpr,  # K
+    V: tl.constexpr,  # V
+    BK: tl.constexpr,  # BLOCK SIZE along the K dimension
+    BV: tl.constexpr,  # BLOCK SIZE along the V dimension
+    USE_INITIAL_STATE: tl.constexpr,  # whether to use initial state
+    REVERSE: tl.constexpr,  # whether to do autoregressive modeling in the reverse direction
+):
+    i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_h = i_bh % H
+    p_q = q + i_bh * s_k_h + i_k * BK + tl.arange(0, BK) + ((T - 1) * K if not REVERSE else 0)
+    p_k = k + i_bh * s_k_h + i_k * BK + tl.arange(0, BK) + ((T - 1) * K if not REVERSE else 0)
+    p_do = do + i_bh * s_v_h + i_v * BV + tl.arange(0, BV) + ((T - 1) * V if not REVERSE else 0)
+    p_v = v + i_bh * s_v_h + i_v * BV + tl.arange(0, BV) + ((T - 1) * V if not REVERSE else 0)
+    p_dk = dk + (i_bh + i_v * B * H) * s_k_h + i_k * BK + tl.arange(0, BK) + ((T - 1) * K if not REVERSE else 0)
+    p_dk_aux = dk_aux + (i_bh + i_v * B * H) * s_k_h + i_k * BK + tl.arange(0, BK) + ((T - 1) * K if not REVERSE else 0)
+    p_dv = dv + (i_bh + i_k * B * H) * s_v_h + i_v * BV + tl.arange(0, BV) + ((T - 1) * V if not REVERSE else 0)
+    p_w = w + i_bh * s_k_h + i_k * BK + tl.arange(0, BK) + ((T - 1) * K if not REVERSE else 0)
+    b_dh = tl.zeros([BK, BV], dtype=tl.float32)
+    mask_bk = i_k * BK + tl.arange(0, BK) < K
+    mask_bv = i_v * BV + tl.arange(0, BV) < V
+    mask_kv = mask_bk[:, None] & mask_bv[None, :]
+
+    p_u = u + i_h * K + tl.arange(0, BK) + i_k * BK
+    b_u = tl.load(p_u, mask=mask_bk, other=0).to(tl.float32)
+
+    for _ in range(T-1, -1, -1):
+        b_q = tl.load(p_q, mask=mask_bk, other=0).to(tl.float32) * scale
+        b_k = tl.load(p_k, mask=mask_bk, other=0).to(tl.float32)
+        b_v = tl.load(p_v, mask=mask_bv, other=0).to(tl.float32)
+        b_w = tl.load(p_w, mask=mask_bk, other=0).to(tl.float32)
+        b_do = tl.load(p_do, mask=mask_bv, other=0).to(tl.float32)
+        b_dkv = b_q[:, None] * b_do[None, :]
+        b_dk = tl.sum(b_dh * b_v[None, :], axis=1)
+        tl.store(p_dk_aux, b_dk.to(p_dk_aux.dtype.element_ty), mask=mask_bk)
+        b_dk += tl.sum(b_dkv * b_u[:, None] * b_v[None, :], axis=1)
+        b_dv = tl.sum((b_dh + (b_dkv * b_u[:, None])) * b_k[:, None], axis=0)
+
+        tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=mask_bk)
+        tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=mask_bv)
+        b_dh *= tl.exp(b_w)[:, None]
+        b_dh += b_dkv
+
+        p_q += K if REVERSE else -K
+        p_k += K if REVERSE else -K
+        p_v += V if REVERSE else -V
+        p_w += K if REVERSE else -K
+        p_do += V if REVERSE else -V
+        p_dk += K if REVERSE else -K
+        p_dk_aux += K if REVERSE else -K
+        p_dv += V if REVERSE else -V
+
+    if USE_INITIAL_STATE:
+        p_dh0 = dh0 + i_bh * K * V + (i_k * BK + tl.arange(0, BK)[:, None]) * V + (i_v * BV + tl.arange(0, BV)[None, :])
+        tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=mask_kv)
+
+
+class FusedRecurrentRWKV6Function(torch.autograd.Function):
+
+    @staticmethod
+    @contiguous
+    @autocast_custom_fwd
+    def forward(ctx, r, k, v, w, u, scale=None, initial_state=None, output_final_state=False, reverse=False):
+        q = r
+        B, H, T, K, V = *q.shape, v.shape[-1]
+
+        BK, BV = min(triton.next_power_of_2(K), 32), min(triton.next_power_of_2(V), 32)
+        NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
+        num_stages = 1
+        num_warps = 1
+
+        final_state = q.new_empty(B, H, K, V) if output_final_state else None
+
+        o = q.new_empty(NK, B, H, T, V, dtype=torch.float32)
+        grid = (NV, NK, B * H)
+        fused_recurrent_rwkv6_fwd_kernel[grid](
+            q, k, v, w, u, o, initial_state, final_state,
+            k.stride(1),
+            v.stride(1),
+            scale,
+            B=B, H=H, T=T, K=K, V=V, BK=BK, BV=BV,
+            USE_INITIAL_STATE=initial_state is not None,
+            STORE_FINAL_STATE=final_state is not None,
+            REVERSE=reverse,
+            num_warps=num_warps,
+            num_stages=num_stages
+        )
+
+        o = o.sum(0)
+        ctx.save_for_backward(q, k, v, w, u, initial_state)
+        ctx.scale = scale
+        ctx.reverse = reverse
+        return o.to(q.dtype), final_state
+
+    @staticmethod
+    @contiguous
+    @autocast_custom_bwd
+    def backward(ctx, do, dht=None):
+        q, k, v, w, u, initial_state = ctx.saved_tensors
+        B, H, T, K, V = *q.shape, v.shape[-1]
+        scale = ctx.scale
+
+        BK, BV = min(triton.next_power_of_2(K), 16), min(triton.next_power_of_2(V), 64)
+        NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
+        num_stages = 1
+        num_warps = 1
+        dq = q.new_empty(NV, B, H, T, K, dtype=torch.float32)
+        dq_aux = torch.empty_like(dq)
+        grid = (NV, NK, B * H)
+
+        fused_recurrent_rwkv6_bwd_kernel_dq[grid](
+            k, v, w, u, do, dq, dq_aux, initial_state,
+            q.stride(1),
+            v.stride(1),
+            scale,
+            B=B, H=H, T=T, K=K, V=V, BK=BK, BV=BV,
+            USE_INITIAL_STATE=initial_state is not None,
+            REVERSE=ctx.reverse,
+            num_warps=num_warps,
+            num_stages=num_stages
+        )
+        dq = dq.sum(0).to(q)
+        dq_aux = dq_aux.sum(0)
+
+        BK, BV = min(triton.next_power_of_2(K), 32), min(triton.next_power_of_2(V), 32)
+        NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
+
+        dk = q.new_empty(NV, B, H, T, K, dtype=torch.float32)
+        dk_aux = q.new_empty(NV, B, H, T, K, dtype=torch.float32)
+        dv = q.new_empty(NK, B, H, T, V, dtype=torch.float32)
+        dh0 = initial_state.new_empty(B, H, K, V) if initial_state is not None else None
+        grid = (NV, NK, B * H)
+        fused_recurrent_rwkv6_bwd_kernel_dkv[grid](
+            q, k, v, w, u, do, dk, dk_aux, dv, dh0,
+            q.stride(1),
+            v.stride(1),
+            scale,
+            B=B, H=H, T=T, K=K, V=V, BK=BK, BV=BV,
+            num_warps=num_warps,
+            num_stages=num_stages,
+            USE_INITIAL_STATE=initial_state is not None,
+            REVERSE=ctx.reverse,
+        )
+        dk = dk.sum(0).to(k)
+        dv = dv.sum(0).to(v)
+        dk_aux = dk_aux.sum(0)
+
+        dw = (dq_aux * q * scale)[:, :, 1:] - (dk_aux * k)[:, :, 0:-1]
+        dw = torch.nn.functional.pad(dw, (0, 0, 0, 1, 0, 0, 0, 0), value=0)
+        dw = chunk_global_reversed_cumsum(dw).to(w)
+
+        du = ((do * v).sum(-1)[..., None] * k * q * scale).sum([0, -2]).to(u)
+        return dq, dk, dv, dw, du, None, dh0, None, None
+
+
+def fused_recurrent_rwkv6(
+    r: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    w: torch.Tensor,
+    u: torch.Tensor,
+    scale: float = -1,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    r"""
+    Args:
+        r (torch.Tensor):
+            reception of shape `(B, H, T, K)`. Alias: q, query in linear attention.
+        k (torch.Tensor):
+            keys of shape `(B, H, T, K)`
+        v (torch.Tensor):
+            values of shape `(B, H, T, V)`
+        w (torch.Tensor):
+            data-dependent decays of shape `(B, H, T, K)` in log space! Alias: g.
+        u (torch.Tensor):
+            bonus of shape `(H, K)`
+        scale (Optional[int]):
+            Scale factor for the RWKV6 attention scores.
+            If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
+        initial_state (Optional[torch.Tensor]):
+            Initial state of shape `(B, H, K, V)`. Default: `None`.
+        output_final_state (Optional[bool]):
+            Whether to output the final state of shape `(B, H, K, V)`. Default: `False`.
+    """
+    if scale == -1:
+        scale = r.shape[-1] ** -0.5
+    o, final_state = FusedRecurrentRWKV6Function.apply(r, k, v, w, u, scale, initial_state, output_final_state)
+    return o, final_state

fla2/ops/rwkv6/recurrent_naive.py

@@ -0,0 +1,103 @@
+# -*- coding: utf-8 -*-
+
+from typing import Optional
+
+import torch
+
+
+def naive_recurrent_rwkv6(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    w: torch.Tensor,
+    u: torch.Tensor,
+    scale: Optional[float] = None,
+    initial_state: Optional[torch.Tensor] = None,
+    output_final_state: Optional[bool] = False
+):
+    orig_dtype = q.dtype
+    B, H, T, K, V = *q.shape, v.shape[-1]
+    q, k, v, w, u = map(lambda x: x.float(), (q, k, v, w, u))
+    h = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device)
+    o = torch.zeros_like(v)
+
+    if scale is None:
+        scale = K ** -0.5
+
+    if initial_state is not None:
+        h += initial_state
+
+    for i in range(T):
+        q_i = q[:, :, i, :] * scale
+        k_i = k[:, :, i]
+        v_i = v[:, :, i, :]
+        w_i = w[:, :, i].exp()
+        kv_i = k_i[..., None] * v_i[..., None, :]
+        o_i = (h + u[None, ..., None] * kv_i) * q_i[..., None]
+        o[:, :, i] = o_i.sum(-2)
+        h = h * w_i[..., None] + kv_i
+    ht = h if output_final_state else None
+    return o.to(orig_dtype), ht
+
+
+@torch.no_grad
+@torch.jit.script
+def naive_recurrent_rwkv6_bwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    w: torch.Tensor,
+    u: torch.Tensor,
+    o: torch.Tensor,
+    do: torch.Tensor,
+    initial_state: Optional[torch.Tensor] = None
+):
+    q, k, v, w, u, o, do = (x.to(dtype=torch.float32) for x in (q, k, v, w, u, o, do))
+    B, H, T, K, V = q.shape[0], q.shape[1], q.shape[2], q.shape[3], v.shape[-1]
+    h = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device)
+    dq = torch.zeros_like(q)
+    dq_aux = torch.zeros_like(q)
+
+    if initial_state is not None:
+        h += initial_state
+
+    for i in range(T):
+        k_i = k[:, :, i]
+        v_i = v[:, :, i]
+        w_i = w[:, :, i].exp()
+        kv_i = k_i[..., None] * v_i[..., None, :]
+        h_i = (h + u[None, ..., None] * kv_i)
+        dq_i = (do[:, :, i, None, :] * h_i).sum(-1)
+        dq_aux_i = (do[:, :, i, None, :] * h).sum(-1)
+        dq[:, :, i] = dq_i
+        dq_aux[:, :, i] = dq_aux_i
+        h = h * w_i[..., None] + kv_i
+
+    du = torch.zeros_like(u)
+    dh = torch.zeros_like(h)
+    dk = torch.zeros_like(k)
+    dk_aux = torch.zeros_like(k)
+    dv = torch.zeros_like(v)
+
+    for i in range(T - 1, -1, -1):
+        d_kv_i = do[:, :, i, None, :] * q[:, :, i, :, None]
+        k_i = k[:, :, i]
+        v_i = v[:, :, i]
+        du_i = (d_kv_i * k_i[..., None] * v_i[..., None, :]).sum(-1)
+        du += du_i.sum(0)
+        dk_i = (dh * v_i[..., None, :]).sum(-1)
+        dk_aux[:, :, i] = dk_i
+        dk_i += (d_kv_i * u[None, ..., None] * v_i[..., None, :]).sum(-1)
+        dv_i = (d_kv_i * u[None, ..., None] * k_i[..., None]).sum(-2)
+        dv_i += (dh * k_i[..., None]).sum(-2)
+
+        dk[:, :, i] = dk_i
+        dv[:, :, i] = dv_i
+        dh = dh * w[:, :, i, :, None].exp() + d_kv_i
+
+    # dw = q * dq_aux - k * dk_aux
+    dw = torch.zeros_like(w)
+    for i in range(T - 2, -1, -1):
+        dw[:, :, i] = dw[:, :, i+1] + dq_aux[:, :, i+1] * q[:, :, i+1] - dk_aux[:, :, i] * k[:, :, i]
+
+    return dq, dk, dv, dw, du, dh

fla2/ops/simple_gla/README.md

@@ -0,0 +1,5 @@
+- Simple GLA
+
+Gating mechanism in https://arxiv.org/abs/2103.02143. Compared to GLA, the gating is head-wise instead of elementwise. As a result, we can adapt the RetNet kernel for training using matmul w/o numerical instability. It is faster than GLA but has less expressive power. I will use it as a baseline for the GLA.
+
+$S_{t+1} = g_{t+1} \odot S_{t} + K_{t+1} V_{t+1}^{\top}$ where $g$ is a scalar.

fla2/ops/simple_gla/__init__.py

@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+
+from .chunk import chunk_simple_gla
+
+__all__ = [
+    'chunk_simple_gla'
+]

fla2/ops/simple_gla/chunk.py

@@ -0,0 +1,299 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023, Yu Zhang, Songlin Yang
+
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+from fla.utils import autocast_custom_bwd, autocast_custom_fwd, contiguous
+from fla.ops.utils import chunk_local_cumsum, chunk_global_reversed_cumsum
+from fla.ops.common.chunk_h import chunk_fwd_h_fn, chunk_bwd_dh_fn
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=4),
+    ],
+    key=["BT", "BK", "BV"],
+)
+@triton.jit
+def chunk_simple_gla_fwd_kernel_o(
+    q,
+    k,
+    v,
+    h,
+    g,
+    o,
+    s_qk_h,
+    s_qk_t,
+    s_qk_d,
+    s_vo_h,
+    s_vo_t,
+    s_vo_d,
+    s_h_h,
+    s_h_t,
+    scale,
+    T: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr
+):
+    i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+
+    o_i = tl.arange(0, BT)
+    m_s = o_i[:, None] >= o_i[None, :]
+
+    b_o = tl.zeros([BT, BV], dtype=tl.float32)
+    b_s = tl.zeros([BT, BT], dtype=tl.float32)
+    for i_k in range(tl.cdiv(K, BK)):
+        p_q = tl.make_block_ptr(q + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+        p_k = tl.make_block_ptr(k + i_bh * s_qk_h, (K, T), (s_qk_d, s_qk_t), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
+        p_h = tl.make_block_ptr(h + i_bh * s_h_h + i_t * K * V, (K, V), (s_h_t, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+        # [BT, BK]
+        b_q = tl.load(p_q, boundary_check=(0, 1))
+        # [BK, BT]
+        b_k = tl.load(p_k, boundary_check=(0, 1))
+        # [BK, BV]
+        b_h = tl.load(p_h, boundary_check=(0, 1))
+        b_o += tl.dot(b_q, b_h, allow_tf32=False)
+        b_s += tl.dot(b_q, b_k, allow_tf32=False)
+
+    p_g = tl.make_block_ptr(g + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+    b_g = tl.load(p_g, boundary_check=(0,))
+    b_o = b_o * tl.exp(b_g)[:, None]
+    b_s = b_s * tl.exp(b_g[:, None] - b_g[None, :])
+    b_s = tl.where(m_s, b_s, 0)
+
+    p_v = tl.make_block_ptr(v + i_bh * s_vo_h, (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+    b_v = tl.load(p_v, boundary_check=(0, 1))
+    b_o = (b_o + tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False)) * scale
+    p_o = tl.make_block_ptr(o + i_bh * s_vo_h, (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8)
+    ],
+    key=["BT", "BK", "BV"],
+)
+@triton.jit
+def chunk_simple_gla_bwd_kernel_dqkvg(
+    q,
+    k,
+    v,
+    h,
+    g,
+    do,
+    dh,
+    dq,
+    dk,
+    dv,
+    dg,
+    s_qk_h,
+    s_qk_t,
+    s_qk_d,
+    s_vo_h,
+    s_vo_t,
+    s_vo_d,
+    s_h_h,
+    s_h_t,
+    scale,
+    T: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    NT: tl.constexpr
+):
+    i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    n_bh = tl.num_programs(2)
+    o_i = tl.arange(0, BT)
+
+    p_q = tl.make_block_ptr(q + i_bh * s_qk_h, (K, T), (s_qk_d, s_qk_t), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
+    p_k = tl.make_block_ptr(k + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+
+    b_q = tl.load(p_q, boundary_check=(0, 1))
+    b_k = tl.load(p_k, boundary_check=(0, 1))
+    b_s = tl.dot(b_k, b_q, allow_tf32=False)
+    p_g = tl.make_block_ptr(g + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+    b_g = tl.load(p_g, boundary_check=(0,))
+    if i_t < NT - 1:
+        b_g_last = tl.load(g + i_bh * T + i_t * BT + BT - 1)
+    else:
+        b_g_last = tl.load(g + i_bh * T + T - 1)
+    mask = tl.exp(b_g[None, :] - b_g[:, None])
+    mask = tl.where(o_i[:, None] <= o_i[None, :], mask * scale, 0)
+    b_s = b_s * mask
+    b_dq = tl.zeros([BT, BK], dtype=tl.float32)
+    b_dk = tl.zeros([BT, BK], dtype=tl.float32)
+    b_ds = tl.zeros([BT, BT], dtype=tl.float32)
+    for i_v in range(tl.cdiv(V, BV)):
+        p_v = tl.make_block_ptr(v + i_bh * s_vo_h, (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        p_h = tl.make_block_ptr(h + i_bh * s_h_h, (V, NT * K), (1, s_h_t), (i_v * BV, i_t * K + i_k * BK), (BV, BK), (0, 1))
+        p_do = tl.make_block_ptr(do + i_bh * s_vo_h, (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        p_dh = tl.make_block_ptr(dh + i_bh * s_h_h, (NT * K, V), (s_h_t, 1), (i_t * K + i_k * BK, i_v * BV), (BK, BV), (1, 0))
+        p_dv = tl.make_block_ptr(dv + (i_k*n_bh+i_bh)*s_vo_h, (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        # [BT, BV]
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        b_do = tl.load(p_do, boundary_check=(0, 1))
+        # [BV, BK]
+        b_h = tl.load(p_h, boundary_check=(0, 1))
+        # [BK, BV]
+        b_dh = tl.load(p_dh, boundary_check=(0, 1))
+        # [BT, BT]
+        b_ds += tl.dot(b_do, tl.trans(b_v), allow_tf32=False)
+        # [BT, BK]
+        b_dq += tl.dot(b_do, b_h, allow_tf32=False) * scale
+        b_dk += tl.dot(b_v, tl.trans(b_dh), allow_tf32=False)
+        # [BT, BV]
+        b_dv = tl.dot(b_k, b_dh, allow_tf32=False) * tl.exp(-b_g + b_g_last)[:, None]
+        b_dv += tl.dot(b_s.to(b_q.dtype), b_do, allow_tf32=False)
+        tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))
+
+    b_dq = b_dq * tl.exp(b_g)[:, None]
+    b_dk = b_dk * tl.exp(-b_g + b_g_last)[:, None]
+    b_ds = b_ds * tl.trans(mask)
+    b_ds = b_ds.to(b_k.dtype)
+    # [BT, BK]
+    b_dq += tl.dot(b_ds, b_k, allow_tf32=False)
+    b_dk += tl.trans(tl.dot(b_q, b_ds, allow_tf32=False))
+    p_dq = tl.make_block_ptr(dq + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+    p_dk = tl.make_block_ptr(dk + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+    tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
+    
+    tl.debug_barrier()
+    b_ds = None 
+    b_s = None
+    b_q = None
+    p_q = tl.make_block_ptr(q + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+    b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) 
+    b_dg = tl.sum(b_dq * b_q - b_dk * b_k.to(tl.float32), axis=1)
+    p_dg = tl.make_block_ptr(dg + (i_k*n_bh + i_bh) * T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+    tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,))
+
+    
+def chunk_fwd_o_fn(h, q, k, v, g, BT, scale):
+    B, H, T, K, V = *k.shape, v.shape[-1]
+    o = torch.empty_like(v)
+    BK = min(triton.next_power_of_2(K), 64)
+    BV = min(triton.next_power_of_2(V), 64)
+    NV = triton.cdiv(V, BV)
+    NT = triton.cdiv(T, BT)
+    grid = (NV, NT, B * H)
+    chunk_simple_gla_fwd_kernel_o[grid](
+        q, k, v, h, g, o,
+        q.stride(1), q.stride(2), q.stride(3),
+        v.stride(1), v.stride(2), v.stride(3),
+        h.stride(1), h.stride(2),
+        scale,
+        T=T, K=K, V=V, BT=BT, BK=BK, BV=BV
+    )
+    return o
+
+
+def chunk_bwd_dqkvg_fn(do, q, k, v, g, h, dh, scale):
+    B, H, T, K, V = *k.shape, v.shape[-1]
+    BT = 64
+    BK = min(triton.next_power_of_2(K), 64)
+    BV = min(triton.next_power_of_2(V), 64)
+    NT, NK = triton.cdiv(T, BT), triton.cdiv(K, BK)
+    grid = (NK, NT, B * H)
+    dq = torch.empty_like(q)
+    dk = torch.empty_like(k)
+    dv = v.new_empty(NK, *v.shape)
+    dg = torch.empty(NK, B, H, T, dtype=torch.float32, device=g.device)
+    chunk_simple_gla_bwd_kernel_dqkvg[grid](
+        q, k, v, h, g, do, dh, dq, dk, dv, dg,
+        q.stride(1), q.stride(2), q.stride(3), 
+        v.stride(1), v.stride(2), v.stride(3),
+        dh.stride(1), dh.stride(2),
+        scale,
+        T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT
+    )
+    dv = dv.sum(0)
+    dg = dg.sum(0)
+    dg = chunk_global_reversed_cumsum(dg)
+    return dq, dk, dv, dg
+
+
+
+
+class SimpleGLAFunction(torch.autograd.Function):
+    @staticmethod
+    @contiguous
+    @autocast_custom_fwd
+    def forward(ctx, q, k, v, g, scale, initial_state, output_final_state, checkpoint_level=1):
+        B, H, T, K, V = *q.shape, v.shape[-1]
+        BT = 64
+        g = chunk_local_cumsum(g, BT)
+        h, final_state = chunk_fwd_h_fn(k=k, v=v, g=g, gk=None, gv=None, BT=BT, h0=initial_state, output_final_state=output_final_state)
+        o = chunk_fwd_o_fn(h, q, k, v, g, BT, scale)        
+        if checkpoint_level == 1:
+            h = None
+        ctx.save_for_backward(q, k, v, h, g, initial_state)
+        ctx.scale = scale
+        ctx.BT = BT
+        return o.to(q.dtype), final_state
+    
+    @staticmethod
+    @contiguous
+    @autocast_custom_bwd
+    def backward(ctx, do, dht):
+        BT, scale = ctx.BT, ctx.scale
+        q, k, v, h, g, initial_state = ctx.saved_tensors
+        if h is None:
+            h, final_state = chunk_fwd_h_fn(k=k, v=v, g=g, gk=None, gv=None, BT=BT, h0=initial_state, output_final_state=False)
+        dh, dh0 = chunk_bwd_dh_fn(q=q, k=k, v=v, g=g, gk=None, gv=None, do=do, h0=initial_state, dht=dht, BT=BT, scale=scale)
+        dq, dk, dv, dg = chunk_bwd_dqkvg_fn(do, q, k, v, g, h, dh, scale)
+        return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), dg.to(g.dtype), None, dh0, None, None
+
+
+
+def chunk_simple_gla(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,  # log decay
+    scale: Optional[float] = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+    checkpoint_level: int = 1
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    r"""
+    Args:
+        q (torch.Tensor):
+            queries of shape `(B, H, T, K)`
+        k (torch.Tensor):
+            keys of shape `(B, H, T, K)`
+        v (torch.Tensor):
+            values of shape `(B, H, T, V)`
+        g (torch.Tensor):
+            Forget gates of shape `(B, H, T)` applied to keys.
+            Compared to GLA, the gating is head-wise instead of elementwise.
+        scale (Optional[int]):
+            Scale factor for the attention scores.
+            If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
+        initial_state (Optional[torch.Tensor]):
+            Initial state of shape `(B, H, K, V)`. Default: `None`.
+        output_final_state (Optional[bool]):
+            Whether to output the final state of shape `(B, H, K, V)`. Default: `False`.
+        checkpoint_level (Optional[int]):
+            Checkpointing level; higher values will save more memories and do more recomputations during backward.
+            Default: `1` (recommended):
+            - Level `0`: no memory saved, no recomputation.
+            - Level `1`: recompute the chunk-level hidden state `h` during backward pass.
+    """
+    assert checkpoint_level in [0, 1], "checkpoint_level must be 0, 1"
+    assert q.dim() == k.dim() == v.dim() == 4, "q, k, v must have 4 dimensions (b, h, l, d)"
+    assert q.dtype == k.dtype == v.dtype, "q, k, v must have the same dtype"
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+    g = g.float()
+    o, final_state = SimpleGLAFunction.apply(q, k, v, g, scale, initial_state, output_final_state, checkpoint_level)
+    return o, final_state

fla2/ops/simple_gla/naive.py

@@ -0,0 +1,81 @@
+# -*- coding: utf-8 -*-
+
+import torch
+from einops import rearrange
+
+
+def torch_simple_gla(q, k, v, g, chunk_size=64, scale=None):
+    if scale is None:
+        scale = (q.shape[-1] ** -0.5)
+    q = rearrange(q, 'b h (n c) d -> b h n c d', c=chunk_size) * scale
+    k = rearrange(k, 'b h (n c) d -> b h n c d', c=chunk_size)
+    v = rearrange(v, 'b h (n c) d -> b h n c d', c=chunk_size)
+    g = rearrange(g, 'b h (n c) -> b h n c', c=chunk_size)
+    g = g.cumsum(-1)
+    kv = k.transpose(-1, -2) @ (v * (-g + g[:, :, :, -1, None]).exp()[..., None])
+    S = torch.zeros_like(kv)
+
+    for i in range(1, g.shape[-2]):
+        S[:, :, i] = S[:, :, i-1].clone() * g[:, :, i-1, -1, None, None].exp() + kv[:, :, i-1]
+
+    inter = (q * g[..., None].exp()) @ S
+    attn = q @ k.transpose(-1, -2)
+    attn = attn * (g[..., None] - g[..., None, :]).exp()
+    attn = attn.masked_fill(torch.triu(torch.ones(chunk_size, chunk_size, dtype=bool, device=q.device), diagonal=1), 0)
+    intra = attn @ v
+    o = inter + intra
+    return rearrange(o, 'b h n c d -> b h (n c) d')
+
+
+def torch_simple_gla_recurrent(q, k, v, g, initial_state=None, scale=None):
+    B, H, T, DK = q.shape
+    if scale is None:
+        scale = DK ** -0.5
+    q = q * scale
+    _, _, _, DV = v.shape
+    if initial_state is None:
+        S = torch.zeros(B, H, DK, DV).to(q)
+    else:
+        S = initial_state
+    o = torch.zeros(B, H, T, DV).to(q)
+    for i in range(T):
+        gate = g[:, :, i].exp()
+        key = k[:, :, i]
+        value = v[:, :, i]
+        kv = key.unsqueeze(-1) * value.unsqueeze(-2)
+        S = S.clone() * gate.unsqueeze(-1).unsqueeze(-1) + kv
+        q_i = q[:, :, i, :]
+        o_i = (q_i.unsqueeze(-1) * S).sum(-2)
+        o[:, :, i] = o_i
+    return o, S
+
+if __name__ == '__main__':
+    torch.set_default_dtype(torch.bfloat16)
+    B = 4
+    H = 4
+    L = 100
+    DK = 32
+    DV = 32
+    q = torch.randn(B, H, L, DK)
+    k = torch.randn(B, H, L, DK)
+    v = torch.randn(B, H, L, DV)
+    g = torch.nn.functional.logsigmoid(torch.randn(B, H, L))
+    q, k, v, g = map(lambda x: x.cuda().requires_grad_(True), [q, k, v, g])
+    from fla.ops.simple_gla import chunk_simple_gla, fused_recurrent_simple_gla
+
+    o, _ = fused_recurrent_simple_gla(q, k, v, g)
+    do = torch.randn_like(o)
+    o.backward(do)
+    q_grad, k_grad, v_grad, g_grad = q.grad, k.grad, v.grad, g.grad
+    q.grad, k.grad, v.grad, g.grad = None, None, None, None
+    o2, _ = chunk_simple_gla(q, k, v, g)
+    o2.backward(do)
+    q_grad2, k_grad2, v_grad2, g_grad2 = q.grad, k.grad, v.grad, g.grad
+
+    print((o-o2).abs().max())
+    print((q_grad-q_grad2).abs().max())
+    print((k_grad-k_grad2).abs().max())
+    print((v_grad-v_grad2).abs().max())
+    print((g_grad-g_grad2).abs().max())
+
+

fla2/ops/simple_gla/recurrent_fuse.py

@@ -0,0 +1,21 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023, Yu Zhang, Songlin Yang
+
+from typing import Tuple, Optional
+import torch
+from fla.ops.common.fused_recurrent import fused_recurrent 
+
+def fused_recurrent_simple_gla(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    scale: Optional[float] = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+    reverse: bool = False
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if scale is None:
+        scale = q.shape[-1] ** -0.5
+    o, final_state = fused_recurrent(q, k, v, g, None, None, scale, initial_state, output_final_state, reverse)
+    return o, final_state

fla3/layers/__init__.py

@@ -0,0 +1,51 @@
+# # -*- coding: utf-8 -*-
+# # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# from .abc import ABCAttention
+# from .attn import Attention
+# from .based import BasedLinearAttention
+# from .bitattn import BitAttention
+# from .delta_net import DeltaNet
+# from .forgetting_attn import ForgettingAttention
+# from .gated_deltanet import GatedDeltaNet
+# from .gated_deltaproduct import GatedDeltaProduct
+# from .gla import GatedLinearAttention
+# from .gsa import GatedSlotAttention
+# from .hgrn import HGRNAttention
+# from .hgrn2 import HGRN2Attention
+# from .lightnet import LightNetAttention
+# from .linear_attn import LinearAttention
+# from .mamba import Mamba
+# from .mamba2 import Mamba2
+# from .multiscale_retention import MultiScaleRetention
+# from .nsa import NativeSparseAttention
+# from .path_attn import PaTHAttention
+# from .rebased import ReBasedLinearAttention
+# from .rwkv6 import RWKV6Attention
+# from .rwkv7 import RWKV7Attention
+
+# __all__ = [
+#     'ABCAttention',
+#     'Attention',
+#     'BasedLinearAttention',
+#     'BitAttention',
+#     'DeltaNet',
+#     'ForgettingAttention',
+#     'GatedDeltaNet',
+#     'GatedDeltaProduct',
+#     'GatedLinearAttention',
+#     'GatedSlotAttention',
+#     'HGRNAttention',
+#     'HGRN2Attention',
+#     'LightNetAttention',
+#     'LinearAttention',
+#     'Mamba',
+#     'Mamba2',
+#     'MultiScaleRetention',
+#     'NativeSparseAttention',
+#     'ReBasedLinearAttention',
+#     'RWKV6Attention',
+#     'RWKV7Attention',
+#     'PaTHAttention'
+# ]
+from .emdeltanet import emdeltanet