# -*- coding: utf-8 -*- import torch from einops import rearrange from typing import Optional import time import torch import triton import triton.language as tl from einops import rearrange from fla.utils import autocast_custom_bwd, autocast_custom_fwd, contiguous from fla.ops.utils import chunk_local_cumsum from fla.ops import chunk_gated_delta_rule @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 = 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_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 = 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)) 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): # gated_chunk_scaled_dot_kkt_fwd(k=k,beta=beta,g_cumsum=g,mask=mask,BT=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() 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 #finish @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_chunk_delta_rule_fwd_kernel_h( k, v,#u d,#w v_new, g, h, initial_state, # initial state of the chunk [B, H, D_head_K, D_head_V] final_state, # final state of the chunk [B, H, D_head_K, D_head_V] H: tl.constexpr, T: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, NT: tl.constexpr, r: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr ): i_k, i_v, 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_h0 = tl.make_block_ptr(initial_state + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) for i_t in range(NT): p_h = tl.make_block_ptr(h + i_bh * NT * K * V + i_t * 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)) #这里save是对的 b_h_cumsum = tl.zeros([r, BK//r, BV], dtype=tl.float32) for i_r in range(r): for i_c in range(tl.cdiv(BT, BC)):#BK 大，通过BC 分块 r_mask = tl.arange(0,r) == i_r p_k = tl.make_block_ptr(k + i_bh * K * T, (T, K), (K, 1), (i_t * BT + i_c * BC, i_k * BK + i_r * BK//r), (BC,BK//r), (1, 0))#读取对应 p_d = tl.make_block_ptr((d + i_bh * T * r * K),(T, r, K ),(r * K, K, 1), (i_t * BT + i_c * BC, i_r, i_k * BK), (BC,1,BK),(2,1,0)) p_v = tl.make_block_ptr((v + i_bh * T * r * V),(T, r, V ),(r * V, V, 1), (i_t * BT + i_c * BC, i_r, i_v * BV), (BC,1,BV),(2,1,0)) p_v_new = tl.make_block_ptr(v_new + (i_bh * r + i_r)* T * V, (T , V), (V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC , BV), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1))#BK//r,BC b_d = tl.load(p_d, boundary_check=(0, 1, 2))#BK b_v = tl.load(p_v, boundary_check=(0, 1, 2)) b_v = tl.reshape(b_v,(BC,BV)) b_d = tl.reshape(b_d,(BC,BK)) b_v -= tl.dot(b_d, b_h.to(tl.bfloat16), allow_tf32=False)#ok #到这相等的 这里BC tl.store(p_v_new, b_v.to(p_v_new.dtype.element_ty), boundary_check=(0, 1))#至少到这里第一步结果相同 last_idx = min((i_t + 1) * BT, T) - 1 b_g_last = tl.load(g + i_bh*T + last_idx) b_g_last = tl.exp(b_g_last) b_h = b_g_last * b_h bkv = tl.where(r_mask[:,None,None],tl.dot(tl.trans(b_k),b_v.to(b_k.dtype),allow_tf32=False)[None,:,:],0) b_h_cumsum += bkv.to(b_h_cumsum.dtype) b_h += tl.reshape(b_h_cumsum,(BK,BV)) if STORE_FINAL_STATE: p_ht = tl.make_block_ptr(final_state + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) #finish @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_chunk_linear_attn_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, H: tl.constexpr, T: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, r : tl.constexpr ): i_v, i_t, i_bhr = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_bh = i_bhr//r i_r = i_bhr % r rk = K//r 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//r, BK)):#这里需要注意拆分#这里K//BK = r #问题是不同r_block读取了同一份qk，有影响吗 p_q = tl.make_block_ptr(q + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_r * rk + i_k * BK), (BT, BK), (1, 0)) 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 * rk + 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), (V, 1), (i_r * rk + i_k * BK, i_v * BV), (BK, BV), (1, 0)) b_q = tl.load(p_q, boundary_check=(0, 1)) b_k = tl.trans(tl.load(p_k, boundary_check=(0, 1))) 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_g_diff = b_g[:, None] - b_g[None, :] b_s = b_s * safe_exp(b_g_diff)#BT BT o_i = tl.arange(0, BT) m_s = o_i[:, None] >= o_i[None, :] b_s = tl.where(m_s, b_s, 0)#置为0 Bs = 0 p_v = tl.make_block_ptr(v + i_bhr * T * V, (T, V), (V, 1), (i_t * BT , i_v * BV), (BT, BV), (1, 0)) b_v = tl.load(p_v, boundary_check=(0, 1)) b_o = b_o * scale + (tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False)) * scale p_o = tl.make_block_ptr(o + i_bhr * T * V, (T, V), (V,1), (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=1), triton.Config({}, num_warps=2), triton.Config({}, num_warps=4), triton.Config({}, num_warps=8), triton.Config({}, num_warps=16) ], key=["BT", "BK"], ) @triton.jit def preprocess_qkw(q, k, w, g, q_new, k_new, w_new, T, H, K, r:tl.constexpr, BT:tl.constexpr, BK:tl.constexpr, USE_Q:tl.constexpr, ): i_k,i_bh,i_t = tl.program_id(0), tl.program_id(1), tl.program_id(2) p_k = tl.make_block_ptr(k + i_bh*T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) p_w = tl.make_block_ptr(w + i_bh*T*K*r,(T,r*K),(r * K, 1),(i_t * BT, i_k * r * BK) ,(BT,r*BK),(1,0)) p_g = tl.make_block_ptr(g+i_bh*T,(T,),(1,),(i_t*BT,),(BT,),(0,)) p_k_new = tl.make_block_ptr(k_new + i_bh*T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) p_w_new = tl.make_block_ptr(w_new +i_bh*T*K*r,(T,r*K),(r * K, 1),(i_t * BT, i_k * r * BK) ,(BT,r*BK),(1,0)) last_idx = min((i_t + 1) * BT, T) - 1 b_g_last = tl.load(g + i_bh*T + last_idx).to(tl.float32) #read BT 位置 b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) b_w = tl.load(p_w, boundary_check=(0, 1)).to(tl.float32) b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) b_d_last = tl.exp((b_g_last - b_g)) b_d_begin = tl.exp(b_g) b_k = b_k * b_d_last[:, None] b_w = b_w * b_d_begin[:, None] tl.store(p_k_new, b_k.to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) tl.store(p_w_new, b_w.to(p_w_new.dtype.element_ty), boundary_check=(0, 1)) if USE_Q: p_q = tl.make_block_ptr(q + i_bh*T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) p_q_new = tl.make_block_ptr(q_new + i_bh*T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) b_q = b_q * b_d_begin[:, None] tl.store(p_q_new, b_q.to(p_q_new.dtype.element_ty), boundary_check=(0, 1)) #finish def gated_chunk_fwd_h_fn(k, w, u, g, BT, initial_state, final_state): # k, w, u, g, BT, initial_state, final_state B, H, T, K, V = *k.shape,u.shape[-1] _,_,rT,_ = w.shape r = rT//T BK = triton.next_power_of_2(K)#直接划分好 assert BK <= 256, "current kernel does not support head dimension larger than 256." BV = 16 if BK > 128 else 32 BV = 64 if BK <= 64 else BV BC = 16 if BK > 128 else 32 BC = 64 if BK <= 64 else BC BC = min(BT, BC) NT, NK, NV = triton.cdiv(T, BT), triton.cdiv(K, BK), triton.cdiv(V, BV) assert NK == 1 h = k.new_empty(B, H, NT * K, V) grid = (NK,B*H,NT) k_new = torch.empty_like(k) w_new = torch.empty_like(w) preprocess_qkw[grid]( q=None, k=k, w=w, g=g, q_new=None, k_new=k_new, w_new=w_new, T=T, H=H, K=K, r=r, BT=BT, BK=BK, USE_Q=False, ) grid = (NK, NV, B * H) v_new = torch.empty(B,H,r,T,V,dtype=u.dtype,device=u.device)#做了v_new的r_first gated_chunk_delta_rule_fwd_kernel_h[grid](#r没有for循环 k_new,u,w_new, v_new,g,h, initial_state, final_state, H=H, T=T, K=K, V=V, BT=BT, BC=BC, BK=BK, BV=BV, NT=NT,r=r, USE_INITIAL_STATE=initial_state is not None, STORE_FINAL_STATE=final_state is not None, ) return h, v_new #finish def gated_chunk_fwd_o_fn(q, k, v_new,h,g,BT): B,H,r,T,V,K = *v_new.shape,q.shape[-1] BK = triton.next_power_of_2(K//r) o = torch.empty_like(v_new)#there_fore,bhr nT,bv BK = min(triton.next_power_of_2(K//r), 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 * r) #h shape b h nk v gated_chunk_linear_attn_fwd_kernel_o[grid]( q, k, v_new, h, g, o, T*K, K, 1 , r*T*V,T*V,V, NT*K*V,V, scale=K**-0.5, H=H, T=T, K=K, V=V, BT=BT, BK=BK, BV=BV,r = r, ) o = o.sum(dim=2)#沿着r维度求和 return o @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_prepare_dv_kernel( q, k, g, do, dv, s_qk_h, s_qk_t, s_qk_d, s_vo_h, s_vo_t, s_vo_d, T, K, V, scale, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, r: tl.constexpr, ): i_t, i_bhr = tl.program_id(0), tl.program_id(1)#或许也可以r并行 i_bh = i_bhr//r i_r = i_bhr % r b_A = tl.zeros([BT, BT], dtype=tl.float32) block_r = K//r for i_k in range(tl.cdiv(block_r, 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_r * block_r + i_k * BK), (BT, BK), (1, 0)) 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_r + i_k * BK), (BT, BK), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_q = tl.trans(tl.load(p_q, boundary_check=(0, 1))) b_A += 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,)) b_A = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], b_A* safe_exp(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty) for i_v in range(tl.cdiv(V, BV)): 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)) b_do = tl.load(p_do, boundary_check=(0, 1)) p_dv = tl.make_block_ptr(dv + i_bhr * s_vo_h , (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) b_dv = tl.dot(b_A, b_do, allow_tf32=False) tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) #finish def gated_fwd_prepare_dv(q, k, g, do, r,BT): B, H, T, K, V = *k.shape, do.shape[-1] dv = torch.empty(B,H,r,T,V,device = do.device, dtype= do.dtype)#没法like NT = triton.cdiv(T, BT) BK = min(triton.next_power_of_2(K//r),64) BV = min(triton.next_power_of_2(V), 64) gated_fwd_prepare_dv_kernel[(NT, B*H*r)]( q, k, g , do, dv, T*K, K, 1, T*V, V, 1, T, K, V, K**-0.5, BT, BK, BV, r ) return dv #finish @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_chunk_delta_rule_bwd_kernel_dhu( q, k, d, g, do, dh, dv, dv2, s_qk_h, s_qk_t, s_qk_d, s_h_h, scale, H: tl.constexpr, T: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, NT: tl.constexpr, r: tl.constexpr, KR: 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):# 向前偏移了一位,计算流程是对的 p_dh = tl.make_block_ptr(dh + i_bh * s_h_h + i_t * K * V , (K, V), (V, 1), (i_k * BK , i_v * BV), (BK, BV), (1, 0)) tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) 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_d = tl.make_block_ptr(d + i_bh * (T * K * r), (K,T*r), (1, K), (i_k * BK, i_t * BT * r), (BK, BT * r), (0, 1)) p_do = tl.make_block_ptr(do + i_bh * T * V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) last_idx = min((i_t + 1) * BT, T) - 1 b_glast = tl.load(g + i_bh * T + last_idx) b_glast = tl.exp(b_glast) b_q = (tl.load(p_q, boundary_check=(0, 1))) b_q = (b_q * scale).to(b_q.dtype) b_do = tl.load(p_do, boundary_check=(0, 1)) b_d = (tl.load(p_d,boundary_check=(0, 1))) p_dv = tl.make_block_ptr(dv + i_bh * r * T * V, (T*r, V), (V , 1), (i_t * BT * r, i_v * BV), (BT*r, BV), (1, 0))#load r b_dv = tl.load(p_dv, boundary_check=(0, 1))#BT*r Bv b_dhtrans = tl.reshape(b_dh,(r,KR,BV)) for i_r in range(r): rmask = tl.arange(0, r) == i_r #第ir列 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*KR + i_k * BK), (BT, KR), (1, 0))# b_k = tl.load(p_k, boundary_check=(0, 1)) b_dhr = tl.sum(tl.where(rmask[:,None,None],b_dhtrans,0), 0) dv_sum = tl.dot(b_k,b_dhr.to(b_k.dtype),allow_tf32=False) b_dv += tl.reshape((dv_sum[:,None,:]*rmask[None,:,None]).to(b_dv.dtype),(BT*r,BV)) p_dv2 = tl.make_block_ptr(dv2 + i_bh * r * T * V, (T*r, V), (V , 1), (i_t * BT * r, i_v * BV), (BT*r, BV), (1, 0)) tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) b_dh *= b_glast b_dh += tl.dot(b_q, b_do.to(b_q.dtype), allow_tf32=False)-tl.dot(b_d,b_dv.to(b_q.dtype),allow_tf32=False) def gated_chunk_bwd_dhu_fn(q, k, w, g,h0, do, dv, BT): B,H,r,T,V,K = *dv.shape,q.shape[-1] BK = triton.next_power_of_2(K) assert BK <= 256, "current kernel does not support head dimension being larger than 256." BV = 16 if BK > 128 else 32 BV = 64 if BK <= 64 else BV NT, NK, NV = triton.cdiv(T, BT), triton.cdiv(K, BK), triton.cdiv(V, BV)#感觉可以放并行度 assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' dh = q.new_empty(B, H, NT * K,V)#一样的#need 求和 得一起算 q_new = torch.empty_like(q) k_new = torch.empty_like(k) w_new = torch.empty_like(w) # grid = (NK,) grid = (NK,B*H,NT) preprocess_qkw[grid]( q=q, k=k, w=w, g=g, q_new=q_new, k_new=k_new, w_new=w_new, T=T, H=H, K=K, r=r, BT=BT, BK=BK, USE_Q=True, ) grid = (NK, NV, B * H) dv = rearrange(dv,'b h r t v-> b h (t r) v').contiguous() dv2 = torch.empty_like(dv)#一样的 #bhr T V gated_chunk_delta_rule_bwd_kernel_dhu[grid]( q_new, k_new, w_new, g, do, dh, dv, dv2, T*K,K,1, NT*K*V, K**-0.5, H=H, T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT,r=r,KR = K//r, ) return dh, dv2 @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_chunk_delta_rule_bwd_kernel_dqkw( q, k, v, w, g, h, do, dh, dq, dk, dv, dw, 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, s_g_k, scale, H: tl.constexpr, T: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, NT: tl.constexpr, r: tl.constexpr, ): i_k, i_t, i_bhr = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_r = i_bhr%r i_bh = i_bhr//r o_i = tl.arange(0, BT) p_q = tl.make_block_ptr(q + i_bh * s_qk_h, (K, T), (1, K), (i_r*K//r + 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_r*K//r + i_k * BK), (BT, BK), (1, 0)) b_dq = tl.zeros([BT, BK], dtype=tl.float32) b_dk = tl.zeros([BT, BK], dtype=tl.float32) b_dw = tl.zeros([BT*r,BK], dtype=tl.float32) b_ds = tl.zeros([BT, BT], dtype=tl.float32) b_dg_last = tl.zeros([1,],dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): p_v = tl.make_block_ptr(v + i_bhr * s_vo_h, (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) 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_dv = tl.make_block_ptr(dv + 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)) 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_r * K // r + i_k * BK), (BV, BK), (0, 1)) p_dh = tl.make_block_ptr(dh + i_bh * s_h_h, (V,NT * K), (1,s_h_t), (i_v * BV,i_t * K + i_r * K // r + i_k * BK), (BV, BK), (0, 1)) b_v = tl.load(p_v, boundary_check=(0, 1)) b_do = tl.load(p_do, boundary_check=(0, 1)) b_h = (tl.load(p_h, boundary_check=(0, 1)))#BV BK b_dh =(tl.load(p_dh, boundary_check=(0, 1))) b_dg_last += tl.sum(b_h * b_dh) b_ds += tl.dot(b_do, tl.trans(b_v), allow_tf32=False)#ok b_dq += tl.dot(b_do, b_h, allow_tf32=False)#d_do 全， bh应该包含 i_Kbufen b_dk += tl.dot(b_v, b_dh, allow_tf32=False)#用来计算dk,yes 行独立没问题 b_dv = (tl.load(p_dv, boundary_check=(0, 1)))#BT*r BV b_dw += (tl.dot(b_dv.to(b_v.dtype),b_h.to(b_v.dtype))) #get BT*r BK b_q = tl.load(p_q, boundary_check=(0, 1)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_dg = tl.zeros([BT,], dtype=tl.float32) 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_glast = tl.load(g +i_bh*T + (min(i_t * BT + BT, T) - 1)) b_dg_last *= tl.exp(b_glast) p_w = tl.make_block_ptr(w + i_bh * T*r*K, (T*r, K), (K,1), (i_t * BT * r,i_r*K//r + i_k * BK), (BT*r ,BK), (1, 0)) b_w = tl.load(p_w,boundary_check=(0,1))#BT * r ,BK b_dw = b_dw * tl.reshape(tl.broadcast_to(tl.reshape(tl.exp(b_g),(BT,1)),(BT,r)),(BT*r))[:,None] b_dg -= tl.sum(tl.reshape(b_w*b_dw,(BT,r*BK)),-1) b_dq = b_dq*scale*tl.exp(b_g)[:,None] b_dg += tl.sum(b_dq*tl.trans(b_q),1)#BT*BK b_dk = b_dk * safe_exp(b_glast-b_g)[:,None] b_dg -= tl.sum(b_dk*b_k,1)#BT*BK b_dg_last += tl.sum(b_dk*b_k) b_ds = tl.where(o_i[:, None] >= o_i[None, :], b_ds* safe_exp(b_g[:, None] - b_g[None, :]) * scale, 0) b_ds2 = b_ds*(tl.dot(tl.trans(b_q),tl.trans(b_k))) b_dg += tl.sum(b_ds2,axis=1) b_dg -= tl.sum(b_ds2,axis=0) b_ds = b_ds.to(b_k.dtype) 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_r*K//r + 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*K//r + i_k * BK), (BT, BK), (1, 0)) p_dw = tl.make_block_ptr(dw + i_bh * T*r*K, (T*r, K), (K,1), (i_t * BT * r,i_r*K//r + i_k * BK), (BT*r ,BK), (1, 0)) p_dg = tl.make_block_ptr(dg + i_k * s_g_k + i_bh * T,(T,),(1,),(i_t*BT,),(BT,),(0,)) b_dg = tl.where(o_i 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)) p_A = 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_A = tl.load(p_A, boundary_check=(0, 1)).to(k.dtype.element_ty) b_dA2 = tl.zeros([BT*r,BT*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_dA2 += 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)) b_dA2 = tl.where(da_mask, b_dA2, 0) b_dA2 = tl.dot(b_dA2.to(b_A.dtype), tl.trans(b_A), allow_tf32=False) b_dA2 = tl.dot(tl.trans(b_A), b_dA2.to(b_A.dtype), allow_tf32=False) b_dA2 = tl.where(da_mask, -b_dA2, 0) #等价于 kkt的 dA 很多0，对角处 b_dA2 = tl.reshape(b_dA2,(BT,r,BT,r)) 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_dA2 *= safe_exp(b_g[:,None]-b_g[None,:])[:,None,:,None] b_dA += b_dA2 b_dA2 = tl.permute(b_dA2,(0,2,1,3))#Bt bt r r b_A = tl.zeros([BT,BT,r,r], dtype=tl.float32) 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 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 b_A += beta_kkt[:,:,None,None] * (rmask[:,None] * b_mask[None,:])[None,None,:,:] betas = tl.sum(tl.sum(beta_kkt[:,None,:]*g,-1),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)) b_dA2 *= b_A #BT BT r r b_dA2 = tl.sum(tl.reshape(b_dA2,(BT,BT,r*r)),-1) b_dg = tl.sum(b_dA2,1)-tl.sum(b_dA2,0) p_dg = tl.make_block_ptr(dg+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 gated_bwd_prepare_wy_repr(k, v, beta, mask,g, Aw,Au, 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) dg = torch.empty_like(g) dmask = torch.zeros([B*H*NT,r,r],device=k.device,dtype=k.dtype).contiguous() assert BK <= K//r gated_bwd_prepare_wy_repr_kernel[(NT, B*H)]( k, v, beta, mask, g, Aw,Au, dw, du, dk, dv, dbeta,dmask,dg, T*K, K, 1, T*V, V, 1, T, K, V, r, BT, BK, BV ) dmask = dmask.sum(0) return dk, dv, dbeta, dmask,dg class gated_ChunkDeltaRuleFunction(torch.autograd.Function): @staticmethod @contiguous @autocast_custom_fwd def forward(ctx, q, k, v, beta,g,mask,BT, initial_state, output_final_state=False, checkpoint_level=1): B,H,L,K = q.shape g = chunk_local_cumsum(g,BT,head_first=True,output_dtype=torch.float) Aw,Au = gated_chunk_scaled_dot_kkt_fwd(k=k,beta=beta,g_cumsum=g,mask=mask,BT=BT,output_dtype=torch.float32) Aw = solve_tril(A=Aw,mask=mask,k=k,BT=BT,output_dtype=k.dtype) Au = solve_tril(A=Au,mask=mask,k=k,BT=BT,output_dtype=k.dtype) #到这里应该没啥问题 r = mask.shape[-1] w, u = gated_fwd_recompute_w_u(k, v, beta, mask,Aw,Au,BT)# final_state = None if output_final_state: final_state = q.new_empty(q.shape[0], q.shape[1], q.shape[-1], v.shape[-1], dtype=torch.float32, requires_grad=False)#这部分不需要修正 h, v_new = gated_chunk_fwd_h_fn(k, w, u, g, BT, initial_state, final_state)#need change' #final_state almost 一致 o = gated_chunk_fwd_o_fn(q, k, v_new, h, g, BT)#need change if checkpoint_level == 1: h, v_new = None, None #这里重新计算了？ ctx.save_for_backward(q, k, v, beta,g, mask, Aw, Au , h, v_new, initial_state) ctx.BT = BT return o.to(q.dtype), final_state @staticmethod @contiguous @autocast_custom_bwd def backward(ctx, do, d_ht=None): q, k, v, beta, g, mask , Aw,Au, h, v_new, initial_state = ctx.saved_tensors BT = ctx.BT r = mask.shape[-1] start = time.time() w, u = gated_fwd_recompute_w_u(k, v, beta, mask, Aw,Au,BT)#跳过 end = time.time() print('recompute_wu:',end-start) if h is None: h, v_new = gated_chunk_fwd_h_fn(k, w, u, g, BT, initial_state, None) start = time.time() #从这里开始重新书写计算代码 dv = gated_fwd_prepare_dv(q, k, g, do, r, BT)#qk do v_new#因此这个dv应该是一个w的shape finish end = time.time() print('pre:',end-start) #dv BHR T V start = time.time() dh, dv = gated_chunk_bwd_dhu_fn(q, k, w,g,initial_state,do, dv, BT)#new_dv dh #final for wyper dv end = time.time() print('chunk_bwd_dhu_fn:',end-start) start = time.time() dq, dk, dw , dg = gated_chunk_bwd_dqkw_fn(q, k, v_new, w, g, h, dv, do, dh, BT)#这一步也巨慢 end = time.time() print('chunk_bwd_dqkw_fn:',end-start) #仅仅两个dg位置可能出错，别的不会 start = time.time() dk2, dv, dbeta,dmask,dg2 = gated_bwd_prepare_wy_repr(k, v, beta, mask,g, Aw,Au, dw, dv, BT) dk.add_(dk2) dg.add_(dg2) end = time.time() print('bwd_prepare_wy_repr:',end-start) #仅仅两个dg位置可能出错，别的不会 dg = chunk_local_cumsum(dg, BT, reverse=True,head_first=True,output_dtype=torch.float) return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), dbeta.to(beta.dtype),dg,dmask.to(mask.dtype),None, None, None def mask_gated_chunk_delta_rule( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, beta: torch.Tensor, g: torch.Tensor, mask: torch.Tensor,#use for mask org_tensor BT: int, initial_state: torch.Tensor = None, output_final_state: bool = False ): assert q.dtype == k.dtype == v.dtype assert q.dtype != torch.float32, "FusedChunkDeltaRuleFunction does not support float32. Please use bfloat16." o, final_state = gated_ChunkDeltaRuleFunction.apply(q, k, v, beta,g,mask, BT, initial_state, output_final_state) return o, final_state def delta_rule_recurrence(q, k, v, beta,g, mask): b, h, l, d_k = q.shape d_v = v.shape[-1] r = mask.shape[-1] o = torch.zeros_like(v) S = torch.zeros(b, h, d_k, d_v,device=k.device,dtype=torch.float32) q = q * (d_k ** -0.5) if beta.ndim < v.ndim: beta = beta[..., None] for i in range(l): _k = k[:, :, i] _q = q[:, :, i] _v = v[:, :, i].clone() beta_i = beta[:, :, i] _v = _v * beta_i kkt = torch.einsum('b h d,b h v->b h d v',_k*beta_i,_k) kkt = rearrange(kkt,' b h (r d) (l v)-> b h r d l v',r= r,l=r) kkt = torch.einsum('b h r d l v,r l->b h r d l v',kkt,mask.to(kkt)) kkt = rearrange(kkt,'b h r d l v-> b h (r d) (l v)') iplr = torch.eye(d_k).to(q)-kkt iplr = torch.einsum(' b h q k ,b h->b h q k',iplr,g[:,:,i]) S = torch.einsum(' b h q k ,b h k v->b h q v',iplr.float(),S.clone()) + _k.unsqueeze(-1).float() * _v.unsqueeze(-2).float() S = S.float() o[:, :, i] = torch.einsum('bhd,bhdm->bhm', _q.float(), S).to(torch.bfloat16) return o,S @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(#需要解决这几个代码速度的问题，可以考虑分成3个部分，分别参与运算，类似fla3版本通过 拆分进行 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) #get BT BT 16 16 ####在内部尝试一下进行分割16 BT 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): 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","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*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 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 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 #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) betas = tl.sum(tl.sum(beta_kkt[:,None,:]*g,-1),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)) def fwd_prepare_wy_repr(k, v, beta,mask, BT): A = chunk_scaled_dot_kkt_fwd(k,beta,mask,BT,torch.float32) print('done,A') A = solve_tril(A=A,mask=mask,k = k ,BT=BT,output_dtype=k.dtype) print('done,A') w, u = fwd_recompute_w_u(k, v, beta,mask, A, BT) 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() assert BK ==K//r 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 @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_dv_kernel( q, k, do, dv, s_qk_h, s_qk_t, s_qk_d, s_vo_h, s_vo_t, s_vo_d, T, K, V, scale, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, r: tl.constexpr, ): i_t, i_bhr = tl.program_id(0), tl.program_id(1)#或许也可以r并行 i_bh = i_bhr//r i_r = i_bhr % r b_A = tl.zeros([BT, BT], dtype=tl.float32) block_r = K//r for i_k in range(tl.cdiv(block_r, 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_r * block_r + i_k * BK), (BT, BK), (1, 0)) 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_r + i_k * BK), (BT, BK), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_q = tl.trans(tl.load(p_q, boundary_check=(0, 1))) b_q = (b_q * scale).to(b_k.dtype) b_A += tl.dot(b_k, b_q, allow_tf32=False) b_A = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], b_A, 0).to(do.dtype.element_ty) for i_v in range(tl.cdiv(V, BV)): 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)) b_do = tl.load(p_do, boundary_check=(0, 1)) p_dv = tl.make_block_ptr(dv + i_bhr * s_vo_h , (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) b_dv = tl.dot(b_A, b_do, allow_tf32=False) tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) #finish def fwd_prepare_dv(q, k, do, r,BT): B, H, T, K, V = *k.shape, do.shape[-1] dv = torch.empty(B,H,r,T,V,device = do.device, dtype= do.dtype)#没法like NT = triton.cdiv(T, BT) BK = min(triton.next_power_of_2(K//r),64) BV = min(triton.next_power_of_2(V), 64) fwd_prepare_dv_kernel[(NT, B*H*r)]( q, k, do, dv, T*K, K, 1, T*V, V, 1, T, K, V, K**-0.5, BT, BK, BV, r ) return dv #finish @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 chunk_delta_rule_fwd_kernel_h( k, v,#u d,#w v_new, h, initial_state, # initial state of the chunk [B, H, D_head_K, D_head_V] final_state, # final state of the chunk [B, H, D_head_K, D_head_V] H: tl.constexpr, T: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, NT: tl.constexpr, r: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr ): i_k, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)#assert ik=1 all use b_h = tl.zeros([BK, BV], dtype=tl.float32)#读取一横行 if USE_INITIAL_STATE: p_h0 = tl.make_block_ptr(initial_state + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) for i_t in range(NT): p_h = tl.make_block_ptr(h + i_bh * NT * K * V + i_t * 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)) #这里save是对的 b_h_cumsum = tl.zeros([r, BK//r, BV], dtype=tl.float32) for i_r in range(r): for i_c in range(tl.cdiv(BT, BC)):#BK 大，通过BC 分块 r_mask = tl.arange(0,r) == i_r p_k = tl.make_block_ptr(k + i_bh * K * T, (T, K), (K, 1), (i_t * BT + i_c * BC, i_k * BK + i_r * BK//r), (BC,BK//r), (1, 0))#读取对应 p_d = tl.make_block_ptr((d + i_bh * T * r * K),(T, r, K ),(r * K, K, 1), (i_t * BT + i_c * BC, i_r, i_k * BK), (BC,1,BK),(2,1,0)) p_v = tl.make_block_ptr((v + i_bh * T * r * V),(T, r, V ),(r * V, V, 1), (i_t * BT + i_c * BC, i_r, i_v * BV), (BC,1,BV),(2,1,0)) p_v_new = tl.make_block_ptr(v_new + (i_bh * r + i_r)* T * V, (T , V), (V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC , BV), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1))#BK//r,BC b_d = tl.load(p_d, boundary_check=(0, 1, 2))#BK b_v = tl.load(p_v, boundary_check=(0, 1, 2))#BC b_v = tl.reshape(b_v,(BC,BV)) b_d = tl.reshape(b_d,(BC,BK)) b_v -= tl.dot(b_d, b_h.to(b_k.dtype), allow_tf32=False)#ok #到这相等的 这里BC tl.store(p_v_new, b_v.to(p_v_new.dtype.element_ty), boundary_check=(0, 1))#至少到这里第一步结果相同 bkv = tl.where(r_mask[:,None,None],tl.dot(tl.trans(b_k),b_v.to(b_k.dtype),allow_tf32=False)[None,:,:],0) b_h_cumsum += bkv.to(b_h_cumsum.dtype) b_h += tl.reshape(b_h_cumsum,(BK,BV)) if STORE_FINAL_STATE: p_ht = tl.make_block_ptr(final_state + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) #finish @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 chunk_linear_attn_fwd_kernel_o( q, k, v, h, 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, H: tl.constexpr, T: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, r : tl.constexpr ): i_v, i_t, i_bhr = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_bh = i_bhr//r i_r = i_bhr % r rk = K//r 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//r, BK)):#这里需要注意拆分#这里K//BK = r #问题是不同r_block读取了同一份qk，有影响吗 p_q = tl.make_block_ptr(q + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT, i_r * rk + i_k * BK), (BT, BK), (1, 0)) 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 * rk + 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), (V, 1), (i_r * rk + i_k * BK, i_v * BV), (BK, BV), (1, 0)) b_q = tl.load(p_q, boundary_check=(0, 1)) b_q = (b_q * scale).to(b_q.dtype) b_k = tl.trans(tl.load(p_k, boundary_check=(0, 1))) 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) b_s = tl.where(m_s, b_s, 0)#置为0 Bs = 0 p_v = tl.make_block_ptr(v + i_bhr * T * V, (T, V), (V, 1), (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)) p_o = tl.make_block_ptr(o + i_bhr * T * V, (T, V), (V,1), (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)) #finish @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 chunk_delta_rule_bwd_kernel_dhu( q, k, d, do, dh, dv, dv2, s_qk_h, s_qk_t, s_qk_d, s_h_h, scale, H: tl.constexpr, T: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, NT: tl.constexpr, r: tl.constexpr, KR: 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):# 向前偏移了一位,计算流程是对的 p_dh = tl.make_block_ptr(dh + i_bh * s_h_h + i_t * K * V , (K, V), (V, 1), (i_k * BK , i_v * BV), (BK, BV), (1, 0)) tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) b_dh_tmp = tl.zeros([BK, BV], dtype=tl.float32) #全列 for i_c in range(tl.cdiv(BT, BC) - 1, -1, -1): 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 + i_c * BC), (BK, BC), (0, 1))#全读取 p_d = tl.make_block_ptr(d + i_bh * (T * K * r), (K,T*r), (1, K), (i_k * BK, i_t * BT * r + i_c * BC *r), (BK, BC * r), (0, 1)) p_do = tl.make_block_ptr(do + i_bh * T * V, (T, V), (V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) b_q = (tl.load(p_q, boundary_check=(0, 1))) b_q = (b_q * scale).to(b_q.dtype) b_do = tl.load(p_do, boundary_check=(0, 1)) b_d = (tl.load(p_d,boundary_check=(0, 1))) p_dv = tl.make_block_ptr(dv + i_bh * r * T * V, (T*r, V), (V , 1), (i_t * BT * r + i_c * BC * r, i_v * BV), (BC*r, BV), (1, 0))#load r b_dv = tl.load(p_dv, boundary_check=(0, 1))#BT*r Bv b_dhtrans = tl.reshape(b_dh,(r,KR,BV)) for i_r in range(r): rmask = tl.arange(0, r) == i_r #第ir列 p_k = tl.make_block_ptr(k + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d), (i_t * BT + i_c * BC, i_r*KR + i_k * BK), (BC, KR), (1, 0))# b_k = tl.load(p_k, boundary_check=(0, 1)) #BC KR b_dhr = tl.sum(tl.where(rmask[:,None,None],b_dhtrans,0), 0)# KR BV dv_sum = tl.dot(b_k,b_dhr.to(b_k.dtype),allow_tf32=False)#get BC*BV b_dv += tl.reshape((dv_sum[:,None,:]*rmask[None,:,None]).to(b_dv.dtype),(BC*r,BV)) p_dv2 = tl.make_block_ptr(dv2 + i_bh * r * T * V, (T*r, V), (V , 1), (i_t * BT * r + i_c * BC * r, i_v * BV), (BC*r, BV), (1, 0)) tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) b_dh_tmp += tl.dot(b_q, b_do.to(b_q.dtype), allow_tf32=False) b_dh_tmp -= tl.dot(b_d,b_dv.to(b_q.dtype),allow_tf32=False) b_dh += b_dh_tmp @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 chunk_delta_rule_bwd_kernel_dqkw( q, k, v, w, h, do, dh, dq, dk, dv, dw, 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, H: tl.constexpr, T: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, NT: tl.constexpr, r: tl.constexpr, ): i_k, i_t, i_bhr = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_r = i_bhr%r i_bh = i_bhr//r o_i = tl.arange(0, BT) p_q = tl.make_block_ptr(q + i_bh * s_qk_h, (K, T), (1, K), (i_r*K//r + 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_r*K//r + i_k * BK), (BT, BK), (1, 0)) b_dq = tl.zeros([BT, BK], dtype=tl.float32) b_dk = tl.zeros([BT, BK], dtype=tl.float32) b_dw = tl.zeros([BT*r,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_bhr * s_vo_h, (T, V), (s_vo_t, s_vo_d), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) 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_dv = tl.make_block_ptr(dv + 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)) 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_r * K // r + i_k * BK), (BV, BK), (0, 1)) p_dh = tl.make_block_ptr(dh + i_bh * s_h_h, (V,NT * K), (1,s_h_t), (i_v * BV,i_t * K + i_r * K // r + i_k * BK), (BV, BK), (0, 1)) b_v = tl.load(p_v, boundary_check=(0, 1)) b_do = tl.load(p_do, boundary_check=(0, 1)) b_h = (tl.load(p_h, boundary_check=(0, 1)))#BV BK b_dh =(tl.load(p_dh, boundary_check=(0, 1))) b_ds += tl.dot(b_do, tl.trans(b_v), allow_tf32=False)#ok b_dq += tl.dot(b_do, b_h, allow_tf32=False)#d_do 全， bh应该包含 i_Kbufen b_dk += tl.dot(b_v, b_dh, allow_tf32=False)#用来计算dk,yes 行独立没问题 b_dv = (tl.load(p_dv, boundary_check=(0, 1)))#BT*r BV b_dw += (tl.dot(b_dv.to(b_v.dtype),b_h.to(b_v.dtype))) #get BT*r BK b_q = tl.load(p_q, boundary_check=(0, 1)) b_q = (b_q * scale).to(b_q.dtype) b_k = tl.load(p_k, boundary_check=(0, 1)) b_ds = tl.where(o_i[:, None] >= o_i[None, :], b_ds, 0).to(b_q.dtype)#BT*BT b_dq += tl.dot(b_ds, b_k, allow_tf32=False) b_dq *= scale 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_r*K//r + 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*K//r + i_k * BK), (BT, BK), (1, 0)) p_dw = tl.make_block_ptr(dw + i_bh * T*r*K, (T*r, K), (K,1), (i_t * BT * r,i_r*K//r + i_k * BK), (BT*r ,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.store(p_dw, ((-b_dw.to(p_dw.dtype.element_ty))), boundary_check=(0, 1)) #finish def chunk_fwd_h_fn(k, w, u, BT, initial_state, final_state): B, H, T, K, V = *k.shape,u.shape[-1] _,_,rT,_ = w.shape r = rT//T BK = triton.next_power_of_2(K)#直接划分好 assert BK <= 256, "current kernel does not support head dimension larger than 256." BV = 16 if BK > 128 else 32 BV = 64 if BK <= 64 else BV BC = 16 if BK > 128 else 32 BC = 64 if BK <= 64 else BC BC = min(BT, BC) NT, NK, NV = triton.cdiv(T, BT), triton.cdiv(K, BK), triton.cdiv(V, BV) assert NK == 1 h = k.new_empty(B, H, NT * K, V) grid = (NK, NV, B * H) v_new = torch.empty(B,H,r,T,V,dtype=u.dtype,device=u.device)#做了v_new的r_first chunk_delta_rule_fwd_kernel_h[grid](#r没有for循环 k, u, w, v_new, h, initial_state, final_state, H=H, T=T, K=K, V=V, BT=BT, BC=BC, BK=BK, BV=BV, NT=NT,r=r, USE_INITIAL_STATE=initial_state is not None, STORE_FINAL_STATE=final_state is not None, ) return h, v_new #finish def chunk_bwd_dhu_fn(q, k, w, do, dv, BT): B,H,r,T,V,K = *dv.shape,q.shape[-1] BK = triton.next_power_of_2(K) assert BK <= 256, "current kernel does not support head dimension being larger than 256." BV = 16 if BK > 128 else 32 BV = 64 if BK <= 64 else BV BC = 16 if BK > 128 else 32 BC = 64 if BK <= 64 else BC BC = min(BT, BC) NT, NK, NV = triton.cdiv(T, BT), triton.cdiv(K, BK), triton.cdiv(V, BV)#感觉可以放并行度 assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' dh = q.new_empty(B , H, NT * K,V)#一样的#need 求和 得一起算 grid = (NK, NV, B * H) dv = rearrange(dv,'b h r t v-> b h (t r) v').contiguous() dv2 = torch.empty_like(dv)#一样的 #bhr T V chunk_delta_rule_bwd_kernel_dhu[grid]( q, k, w, do, dh, dv, dv2, T*K,K,1, NT*K*V, K**-0.5, H=H, T=T, K=K, V=V, BT=BT, BC=BC, BK=BK, BV=BV, NT=NT,r=r,KR = K//r, ) return dh, dv2 #finish def chunk_fwd_o_fn(q, k, v_new, h, BT): B,H,r,T,V,K = *v_new.shape,q.shape[-1] BK = triton.next_power_of_2(K//r) o = torch.empty_like(v_new)#there_fore,bhr nT,bv BK = min(triton.next_power_of_2(K//r), 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 * r) #h shape b h nk v chunk_linear_attn_fwd_kernel_o[grid]( q, k, v_new, h, o, T*K, K, 1 , r*T*V,T*V,V, NT*K*V,V, scale=K**-0.5, H=H, T=T, K=K, V=V, BT=BT, BK=BK, BV=BV,r = r, ) o = o.sum(dim=2)#沿着r维度求和 return o def chunk_bwd_dqkw_fn(q, k, v_new, w, h, du, do, dh, BT): B, H, T, K, V = *q.shape, v_new.shape[-1] _,_,RT,_ = w.shape r = RT // T #最后一个函数，计算dw,dq,dk BK = triton.next_power_of_2(K//r)#需要更细粒度的划分，确保不会使得 不同位置的划到一起 BK = min(triton.next_power_of_2(K//r), 64) BV = min(triton.next_power_of_2(V), 64) NK = triton.cdiv(K//r, BK) NT = triton.cdiv(T, BT) grid = (NK, NT, B * H * r)#通过NK控制位置 dq = torch.empty_like(q) dk = torch.empty_like(k)#k_org dw = torch.empty_like(w)#bh nt k chunk_delta_rule_bwd_kernel_dqkw[grid]( q, k, v_new, w, h, do, dh, dq, dk, du, dw, T*K,K,1, T*V, V, 1, NT*K*V,V, scale=K ** -0.5, H=H, T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT,r = r ) return dq.to(q.dtype), dk.to(k.dtype), dw.to(w.dtype) class ChunkDeltaRuleFunction(torch.autograd.Function): #前向写完了 @staticmethod @contiguous @autocast_custom_fwd def forward(ctx, q, k, v, beta,mask,BT, initial_state, output_final_state, checkpoint_level=1): start = time.time() w, u, A = fwd_prepare_wy_repr(k, v,beta, mask, BT)#compute for A matrix #compute all final_state = None if output_final_state: final_state = q.new_empty(q.shape[0], q.shape[1], q.shape[-1], v.shape[-1], dtype=torch.float32, requires_grad=False)#这部分不需要修正 end = time.time() print('compute_A:',end-start) start = time.time() h, v_new = chunk_fwd_h_fn(k, w, u, BT, initial_state, final_state)#need change' end = time.time() print('compute_h_s:',end-start) start = time.time() o = chunk_fwd_o_fn(q, k, v_new, h, BT)#need change end = time.time() print('compute_o:',end-start) if checkpoint_level == 1: h, v_new = None, None ctx.save_for_backward(q, k, v, beta,mask, A, h, v_new, initial_state) ctx.BT = BT return o.to(q.dtype), final_state @staticmethod @contiguous @autocast_custom_bwd def backward(ctx, do, d_ht=None): q, k, v, beta,mask , A, h, v_new, initial_state = ctx.saved_tensors BT = ctx.BT r = mask.shape[-1] start = time.time() w, u = fwd_recompute_w_u(k, v, beta, mask, A, BT)#跳过 end = time.time() print('recompute_wu:',end-start) # checkpont_level=1, recomputation. if h is None: h, v_new = chunk_fwd_h_fn(k, w, u, BT, initial_state, None) #v_new b h r T V start = time.time() dv = fwd_prepare_dv(q, k, do, r, BT)#qk do v_new#因此这个dv应该是一个w的shape finish end = time.time() print('pre:',end-start) #dv BHR T V start = time.time() dh, dv = chunk_bwd_dhu_fn(q, k, w, do, dv, BT)#new_dv dh #final for wyper dv end = time.time() print('chunk_bwd_dhu_fn:',end-start) start = time.time() dq, dk, dw = chunk_bwd_dqkw_fn(q, k, v_new, w, h, dv, do, dh, BT)#这一步也巨慢 end = time.time() print('chunk_bwd_dqkw_fn:',end-start) start = time.time() dk2, dv, dbeta,dmask = bwd_prepare_wy_repr(k, v, beta, mask, A, dw, dv, BT) dk.add_(dk2) end = time.time() print('bwd_prepare_wy_repr:',end-start) return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), dbeta.to(beta.dtype), dmask.to(mask.dtype), None, None, None def mask_chunk_delta_rule( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, beta: torch.Tensor, mask: torch.Tensor,#use for mask org_tensor BT: int, initial_state: torch.Tensor = None, output_final_state: bool = False ): assert q.dtype == k.dtype == v.dtype assert q.dtype != torch.float32, "FusedChunkDeltaRuleFunction does not support float32. Please use bfloat16." o, final_state = ChunkDeltaRuleFunction.apply(q, k, v, beta,mask, BT, initial_state, output_final_state) return o, final_state if __name__ =="__main__": import sys import time torch.set_default_dtype(torch.bfloat16) torch.manual_seed(42) # for i in range(200): B = 16 H = 4 L = 128 DK = 256 DV = 256 r = 4 q = (torch.randn(B, H, L, DK)).cuda().requires_grad_(True) k = (torch.randn(B, H, L, DK)).cuda() k = torch.nn.functional.normalize(k, dim=-1, p=2).requires_grad_(True) v = (torch.randn(B, H, L, DV)).cuda().requires_grad_(True) beta = torch.randn(B, H, L).cuda().sigmoid().requires_grad_(True) # mask = torch.randn([r,r]) # mask = mask.cuda().requires_grad_(True).contiguous() mask = torch.ones([1]) mask = mask[:,None] mask = mask.cuda().requires_grad_(True).contiguous() g = torch.nn.functional.logsigmoid(torch.randn(B, H, L).cuda()).requires_grad_(True) g_exp = (torch.exp(g)) do = torch.randn(B, H, L, DV).cuda() # o1,ss = delta_rule_recurrence(q,k,v,beta,g_exp,mask) # o1.backward(do, retain_graph=True) # q_grad, q.grad = q.grad, None # k_grad, k.grad = k.grad, None # v_grad, v.grad = v.grad, None # mask_grad, mask.grad = mask.grad, None # beta_grad, beta.grad = beta.grad, None # g_grad, g.grad = g.grad, None qh,kh,vh,betah,gh = map(lambda x: rearrange(x, 'b h l ... -> b l h ...'), (q, k, v, beta, g)) o,f_state = chunk_gated_delta_rule(qh,kh,vh,gh,betah,use_qk_l2norm_in_kernel=False) # o,f_state = mask_gated_chunk_delta_rule(q, k, v,beta,g,mask,BT=32,output_final_state=True) # o2,f_state = mask_chunk_delta_rule(q, k, v,beta,mask,BT=32) o = rearrange(o,'b l h d->b h l d') o.backward(do,retain_graph=True) q_grad0, q.grad = q.grad, None k_grad0, k.grad = k.grad, None v_grad0, v.grad = v.grad, None beta_grad0, beta.grad = beta.grad, None # mask_grad0, mask.grad = mask.grad, None g_grad0, g.grad = g.grad, None o2,f_state = mask_gated_chunk_delta_rule(q, k, v,beta,g,mask,BT=32,output_final_state=True) o2.backward(do,retain_graph=True) q_grad2, q.grad = q.grad, None k_grad2, k.grad = k.grad, None v_grad2, v.grad = v.grad, None beta_grad2, beta.grad = beta.grad, None # mask_grad0, mask.grad = mask.grad, None g_grad2, g.grad = g.grad, None print((o2-o).abs().max()) print((q_grad2-q_grad0).abs().max()) print((k_grad2-k_grad0).abs().max())#计算结果差距大 差距到1 print((v_grad2-v_grad0).abs().max()) print((beta_grad2-beta_grad0).abs().max()) # print((mask_grad-mask_grad0).abs().max()) print((g_grad2-g_grad0).abs().max())#这个结果是对的 mask = 1 计算也是对的，那估计就是某个地方mask错了 # print(g_grad0) # print(g_grad)#计算不相等 # print('naive:',mask_grad) # print('triton:',mask_grad0)