msj19/gated_deltaproduct
05
1# -*- coding: utf-8 -*-2 3import torch4from einops import rearrange5from typing import Optional6 7import time8import torch9import triton10import triton.language as tl11from einops import rearrange12from fla.utils import autocast_custom_bwd, autocast_custom_fwd, contiguous13from fla.ops.utils import chunk_local_cumsum14 15from fla.ops import chunk_gated_delta_rule16 17@triton.jit18def safe_exp(x):19 return tl.exp(tl.where(x <= 0, x, float('-inf')))20 21 22 23@triton.autotune(24 configs=[25 triton.Config({}, num_warps=1),26 triton.Config({}, num_warps=2),27 triton.Config({}, num_warps=4),28 triton.Config({}, num_warps=8),29 triton.Config({}, num_warps=16)30 ],31 key=["BT", "BK", "BV"],32)33@triton.jit34def gated_fwd_recompute_w_u_kernel(35 k,36 v,37 beta,38 mask_ij,39 w,40 u,41 Aw,42 Au,43 s_qk_h,44 s_qk_t,45 s_qk_d,46 s_vo_h,47 s_vo_t,48 s_vo_d,49 T,50 K,51 V,52 r: tl.constexpr,53 BT: tl.constexpr,54 BK: tl.constexpr,55 BV: tl.constexpr56):57 i_t, i_bh = tl.program_id(0), tl.program_id(1)58 dk = K//r59 p_beta = tl.make_block_ptr(beta + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))60 b_beta = tl.load(p_beta, boundary_check=(0,))61 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))62 b_Aw = tl.load(p_Aw, boundary_check=(0, 1)).to(k.dtype.element_ty)63 for i_r in range(r):64 p_mask = mask_ij + tl.arange(0,r)*r+i_r#读取第ir列65 b_mask = tl.load(p_mask)66 for i_k in range(tl.cdiv(dk, BK)):67 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))68 b_k = tl.load(p_k, boundary_check=(0, 1))69 b_kb = (b_k * b_beta[:, None]).to(b_k.dtype)[:,None,:]*b_mask[None,:,None].to(b_k.dtype)#BT*r*d70 b_kb = tl.reshape(b_kb,(BT*r,BK))71 b_w = tl.dot(b_Aw, b_kb, allow_tf32=False)#get BT*r *BK72 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))73 tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))74 tl.debug_barrier()75 b_Aw = None76 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))77 b_Au = tl.load(p_Au, boundary_check=(0, 1)).to(k.dtype.element_ty)78 79 for i_v in range(tl.cdiv(V, BV)):#no need for 任意mask不使用 #无需for 循环 ，这里也不存在mask80 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))81 b_v = tl.load(p_v, boundary_check=(0, 1))82 b_vb = (b_v * b_beta[:, None]).to(b_v.dtype)[:,None,:]*tl.full([r],1, dtype=b_v.dtype)[None,:,None]83 b_vb = tl.reshape(b_vb,(BT*r,BV))84 b_u = tl.dot(b_Au, b_vb, allow_tf32=False)85 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))86 tl.store(p_u, (b_u).to(p_u.dtype.element_ty), boundary_check=(0, 1))87 88 89@triton.autotune(90 configs=[91 triton.Config({}, num_warps=1),92 triton.Config({}, num_warps=2),93 triton.Config({}, num_warps=4),94 triton.Config({}, num_warps=8),95 triton.Config({}, num_warps=16)96 ],97 key=["BT", "BK","r"],98)99@triton.jit100def gated_chunk_scaled_dot_kkt_fwd_kernel( 101 k,102 beta,103 g_cumsum,104 mask_ij,105 A,106 Ag,107 s_qk_h,108 s_qk_t,109 s_qk_d,110 T,111 K,112 r: tl.constexpr,113 BT: tl.constexpr,114 BK: tl.constexpr,115):116 i_t, i_bh = tl.program_id(0), tl.program_id(1)117 b_A = tl.zeros([BT,BT,r,r], dtype=tl.float32)#r*BT r*BT118 dk = K//r119 p_beta = tl.make_block_ptr(beta + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))120 b_beta = tl.load(p_beta, boundary_check=(0,))121 for i_r in range(r):122 r_mask = tl.arange(0, r) == i_r 123 p_mask = mask_ij + tl.arange(0,r)* r + i_r#列读，因而是行数目124 b_mask = tl.load(p_mask)125 ij_mask = b_mask[:,None]*r_mask[None,:]#行数126 127 for i_k in range(tl.cdiv(dk, BK)):#分块k读取计算128 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))129 b_k = tl.load(p_k, boundary_check=(0, 1))130 b_kb = (b_k * b_beta[:, None]).to(b_k.dtype)131 dot = tl.dot(b_kb, tl.trans(b_k), allow_tf32=False)132 b_A += dot[:,:,None,None]*ij_mask[None,None,:,:]133 b_A = tl.where((tl.arange(0, BT)[:,None] > tl.arange(0, BT)[None,:])[:,:,None,None], b_A, 0)134 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))135 tl.store(p_A, (b_A).to(p_A.dtype.element_ty),boundary_check=(0,1,2,3))136 137 p_g = tl.make_block_ptr(g_cumsum + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))138 b_g = tl.load(p_g, boundary_check=(0,))139 b_g_diff = b_g[:, None] - b_g[None, :]140 b_g_diff = safe_exp(b_g_diff)141 142 b_Ag = b_A * ((b_g_diff)[:,:,None,None])#BT BT143 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))144 tl.store(p_Ag, (b_Ag).to(p_Ag.dtype.element_ty),boundary_check=(0,1,2,3))145 146 147@triton.autotune(148 configs=[149 triton.Config({}, num_warps=1),150 triton.Config({}, num_warps=2),151 triton.Config({}, num_warps=4),152 triton.Config({}, num_warps=8),153 triton.Config({}, num_warps=16)154 ],155 key=["BT", "r"],156)157@triton.jit158def solve_tril_16x16_kernel(159 A,160 Ad,161 s_A_bh,162 s_Ad_bh,163 T,164 r: tl.constexpr,165 BT: tl.constexpr,166):167 i_t, i_bh = tl.program_id(0), tl.program_id(1)168 offset = (i_t * 16) % BT 169 170 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))171 b_A = tl.load(p_A, boundary_check=(0,1,2,3)).to(tl.float32)172 b_A = -tl.where((tl.arange(0, 16)[:,None] > tl.arange(0, 16)[None,:])[:,:,None,None], b_A, 0)173 174 for i in range(1, 16):175 mask = tl.arange(0, 16) == i 176 b_a = tl.sum(tl.where(mask[:,None,None,None], b_A, 0), 0)177 q = (tl.sum(b_a[:,None,:,:,None]*b_A[:,:,None,:,:],-2))178 b_a = b_a + tl.sum(q,0)*((tl.arange(0, 16) < i)[:,None,None])179 b_A = tl.where(mask[:,None,None,None],b_a,b_A)#按行计算 ，逐步交换结果180 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, :]))181 182 b_A = tl.permute(b_A,(0,2,1,3))183 b_A = tl.reshape(b_A,(16*r,16*r))#BT*r BT*r184 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))185 tl.store(p_Ad, (b_A).to(p_Ad.dtype.element_ty),boundary_check=(0,1))186 187@triton.autotune(188 configs=[189 triton.Config({}, num_warps=1),190 triton.Config({}, num_warps=2),191 triton.Config({}, num_warps=4),192 triton.Config({}, num_warps=8),193 triton.Config({}, num_warps=16)194 ],195 key=["r"],196)197@triton.jit198def merge_16x16_to_32x32_inverse_kernel(199 A,200 Ad,201 Ai,202 s_A_bh,203 s_Ad_bh,204 T,205 r: tl.constexpr,206 BT: tl.constexpr 207):208 i_t, i_bh = tl.program_id(0), tl.program_id(1)209 210 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))211 b_A21 = tl.load(p_A21, boundary_check=(0,1)).to(tl.float32)212 213 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))214 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))215 216 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))217 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))218 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))219 220 Ai11 = tl.load(p_Ad11, boundary_check=(0, 1)).to(tl.float32)221 Ai22 = tl.load(p_Ad22, boundary_check=(0, 1)).to(tl.float32)222 Ai21 = -tl.dot(tl.dot(Ai22,b_A21, input_precision='ieee'),Ai11,input_precision='ieee')223 tl.store(p_Ai11,Ai11.to(p_Ai11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))224 tl.store(p_Ai22,Ai22.to(p_Ai22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))225 tl.store(p_Ai21,Ai21.to(p_Ai21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))226 227 228@triton.autotune(229 configs=[230 triton.Config({}, num_warps=1),231 triton.Config({}, num_warps=2),232 triton.Config({}, num_warps=4),233 triton.Config({}, num_warps=8),234 triton.Config({}, num_warps=16)235 ],236 key=["r"],237)238@triton.jit239def merge_16x16_to_64x64_inverse_kernel(240 A,241 Ad,242 Ai,243 s_A_bh,244 s_Ad_bh,245 T,246 r: tl.constexpr,247 BT: tl.constexpr 248):249 i_t, i_bh = tl.program_id(0), tl.program_id(1)250 251 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))252 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))253 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))254 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))255 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))256 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))257 258 b_A21 = tl.load(p_A21, boundary_check=(0,1)).to(tl.float32)259 b_A31 = tl.load(p_A31, boundary_check=(0,1)).to(tl.float32)260 b_A32 = tl.load(p_A32, boundary_check=(0,1)).to(tl.float32)261 b_A41 = tl.load(p_A41, boundary_check=(0,1)).to(tl.float32)262 b_A42 = tl.load(p_A42, boundary_check=(0,1)).to(tl.float32)263 b_A43 = tl.load(p_A43, boundary_check=(0,1)).to(tl.float32)264 265 266 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))267 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))268 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))269 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))270 271 272 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))273 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))274 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))275 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))276 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))277 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))278 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))279 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))280 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))281 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))282 283 284 Ai11 = tl.load(p_Ad11, boundary_check=(0, 1)).to(tl.float32)285 Ai22 = tl.load(p_Ad22, boundary_check=(0, 1)).to(tl.float32)286 Ai33 = tl.load(p_Ad33, boundary_check=(0, 1)).to(tl.float32)287 Ai44 = tl.load(p_Ad44, boundary_check=(0, 1)).to(tl.float32)288 289 Ai21 = -tl.dot(tl.dot(Ai22,b_A21, input_precision='ieee'),Ai11,input_precision='ieee')290 Ai32 = -tl.dot(tl.dot(Ai33,b_A32, input_precision='ieee'),Ai11,input_precision='ieee')291 Ai43 = -tl.dot(tl.dot(Ai44,b_A43, input_precision='ieee'),Ai11,input_precision='ieee')292 293 Ai31 = -tl.dot(294 Ai33,295 tl.dot(b_A31,Ai11, input_precision='ieee')+296 tl.dot(b_A32,Ai21, input_precision='ieee'),297 input_precision='ieee')298 299 Ai42 = -tl.dot(300 Ai44,301 tl.dot(b_A42,Ai22, input_precision='ieee')+302 tl.dot(b_A43,Ai32, input_precision='ieee'),303 input_precision='ieee')304 305 Ai41 = -tl.dot(306 Ai44,307 tl.dot(b_A41, Ai11, input_precision='ieee') +308 tl.dot(b_A42, Ai21, input_precision='ieee') +309 tl.dot(b_A43, Ai31, input_precision='ieee'),310 input_precision='ieee'311 )312 313 tl.store(p_Ai11,Ai11.to(p_Ai11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))314 tl.store(p_Ai22,Ai22.to(p_Ai22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))315 tl.store(p_Ai33,Ai33.to(p_Ai33.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))316 tl.store(p_Ai44,Ai44.to(p_Ai44.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))317 tl.store(p_Ai21,Ai21.to(p_Ai21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))318 tl.store(p_Ai31,Ai31.to(p_Ai31.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))319 tl.store(p_Ai32,Ai32.to(p_Ai32.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))320 tl.store(p_Ai41,Ai41.to(p_Ai41.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))321 tl.store(p_Ai42,Ai42.to(p_Ai42.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))322 tl.store(p_Ai43,Ai43.to(p_Ai43.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))323 324 325 326def gated_chunk_scaled_dot_kkt_fwd(k: torch.Tensor,327 beta: torch.Tensor,328 mask: torch.Tensor,329 g_cumsum:Optional[torch.Tensor] = None,330 BT:int = 32,331 output_dtype: torch.dtype=torch.float32):332 # gated_chunk_scaled_dot_kkt_fwd(k=k,beta=beta,g_cumsum=g,mask=mask,BT=BT,output_dtype=torch.float32)333 B, H, T, K = k.shape334 r = mask.shape[-1]335 NT = triton.cdiv(T, BT)336 BK = min(triton.next_power_of_2(K//r), 64)337 A = torch.empty(B*H*NT,BT*BT,r*r,device=k.device, dtype=output_dtype).contiguous() 338 Ag = torch.empty(B*H*NT,BT*BT,r*r,device=k.device, dtype=output_dtype).contiguous() 339 gated_chunk_scaled_dot_kkt_fwd_kernel[(NT, B*H)](340 k, beta, g_cumsum, mask, A,Ag,341 T*K, K, 1,342 T, K, r, BT, BK343 )344 return A,Ag345 346def solve_tril(A,mask,k,BT,output_dtype=torch.float32):347 B, H, T, K = k.shape348 r = mask.shape[-1]349 NT = triton.cdiv(T, 16)350 Ad = torch.empty(B,H,NT*16*r,16*r,device=A.device, dtype=torch.float if BT != 16 else output_dtype)351 solve_tril_16x16_kernel[(NT, B*H)](352 A,Ad,353 T*BT*r*r,#s_abh354 T*16*r*r,#s_adbh355 T,356 r, BT357 )358 if BT == 16: 359 return Ad360 361 A = rearrange(A,'b (t l) (c r)->b (t c) (l r)',t=BT,c=r).contiguous()#BT*r BT*r 362 if BT == 32:363 NT = triton.cdiv(T, BT)364 Ai = torch.zeros(B,H,NT*BT*r,BT*r,device=A.device, dtype=output_dtype)365 merge_16x16_to_32x32_inverse_kernel[(NT, B*H)](366 A,Ad,Ai,367 T*BT*r*r,#s_a_bh and s_ai_bh368 T*16*r*r,#s_ad_bh369 T,r,BT370 )371 return Ai372 373 if BT == 64:374 NT = triton.cdiv(T, BT)375 Ai = torch.zeros(B,H,NT*BT*r,BT*r,device=A.device, dtype=output_dtype)376 merge_16x16_to_64x64_inverse_kernel[(NT, B*H)](377 A,Ad,Ai,378 T*BT*r*r,#s_a_bh and s_ai_bh379 T*16*r*r,#s_ad_bh380 T,r,BT381 )382 return Ai383 384 385def gated_fwd_recompute_w_u(k, v, beta,mask, Aw,Au,BT):386 B, H, T, K, V = *k.shape, v.shape[-1]387 r = mask.shape[-1]388 u = torch.empty(B,H,r*T,V,device=k.device, dtype=k.dtype)389 w = torch.empty(B,H,r*T,K,device=k.device, dtype=k.dtype)390 NT = triton.cdiv(T, BT)391 BK = min(triton.next_power_of_2(K//r), 64)#32392 BV = min(triton.next_power_of_2(V), 64)393 gated_fwd_recompute_w_u_kernel[(NT, B*H)](394 k, v, beta,mask, w, u, Aw,Au,395 T*K, K, 1,396 T*V, V, 1,397 T, K, V, r,BT, BK, BV398 )399 return w, u400 401 402 403 404#finish405@triton.autotune(406 configs=[407 triton.Config({}, num_warps=1),408 triton.Config({}, num_warps=2),409 triton.Config({}, num_warps=4),410 triton.Config({}, num_warps=8),411 triton.Config({}, num_warps=16)412 ],413 key=["BT", "BK", "BV"],414)415@triton.jit416def gated_chunk_delta_rule_fwd_kernel_h(417 k,418 v,#u419 d,#w420 v_new,421 g,422 h,423 initial_state, # initial state of the chunk [B, H, D_head_K, D_head_V]424 final_state, # final state of the chunk [B, H, D_head_K, D_head_V]425 H: tl.constexpr,426 T: tl.constexpr,427 K: tl.constexpr,428 V: tl.constexpr,429 BT: tl.constexpr,430 BC: tl.constexpr,431 BK: tl.constexpr,432 BV: tl.constexpr,433 NT: tl.constexpr,434 r: tl.constexpr,435 USE_INITIAL_STATE: tl.constexpr,436 STORE_FINAL_STATE: tl.constexpr437):438 i_k, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)439 b_h = tl.zeros([BK, BV], dtype=tl.float32)#读取一横行440 if USE_INITIAL_STATE:441 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))442 b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32)443 444 for i_t in range(NT):445 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))446 tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1))447 #这里save是对的448 b_h_cumsum = tl.zeros([r, BK//r, BV], dtype=tl.float32)449 for i_r in range(r):450 for i_c in range(tl.cdiv(BT, BC)):#BK 大，通过BC 分块451 r_mask = tl.arange(0,r) == i_r452 p_k = tl.make_block_ptr(k + i_bh * K * T, (T, K), (K, 1),453 (i_t * BT + i_c * BC, i_k * BK + i_r * BK//r), (BC,BK//r), (1, 0))#读取对应454 p_d = tl.make_block_ptr((d + i_bh * T * r * K),(T, r, K ),(r * K, K, 1),455 (i_t * BT + i_c * BC, i_r, i_k * BK), (BC,1,BK),(2,1,0))456 p_v = tl.make_block_ptr((v + i_bh * T * r * V),(T, r, V ),(r * V, V, 1),457 (i_t * BT + i_c * BC, i_r, i_v * BV), (BC,1,BV),(2,1,0))458 p_v_new = tl.make_block_ptr(v_new + (i_bh * r + i_r)* T * V, (T , V), (V, 1),459 (i_t * BT + i_c * BC, i_v * BV), (BC , BV), (1, 0)) 460 b_k = tl.load(p_k, boundary_check=(0, 1))#BK//r,BC461 b_d = tl.load(p_d, boundary_check=(0, 1, 2))#BK462 b_v = tl.load(p_v, boundary_check=(0, 1, 2))463 b_v = tl.reshape(b_v,(BC,BV))464 b_d = tl.reshape(b_d,(BC,BK))465 b_v -= tl.dot(b_d, b_h.to(tl.bfloat16), allow_tf32=False)#ok #到这相等的 这里BC466 tl.store(p_v_new, b_v.to(p_v_new.dtype.element_ty), boundary_check=(0, 1))#至少到这里第一步结果相同467 468 last_idx = min((i_t + 1) * BT, T) - 1469 b_g_last = tl.load(g + i_bh*T + last_idx)470 b_g_last = tl.exp(b_g_last)471 b_h = b_g_last * b_h472 473 bkv = tl.where(r_mask[:,None,None],tl.dot(tl.trans(b_k),b_v.to(b_k.dtype),allow_tf32=False)[None,:,:],0)474 b_h_cumsum += bkv.to(b_h_cumsum.dtype)475 b_h += tl.reshape(b_h_cumsum,(BK,BV))476 477 if STORE_FINAL_STATE:478 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))479 tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1))480 481#finish482@triton.autotune(483 configs=[484 triton.Config({}, num_warps=1),485 triton.Config({}, num_warps=2),486 triton.Config({}, num_warps=4),487 triton.Config({}, num_warps=8),488 triton.Config({}, num_warps=16)489 ],490 key=["BT", "BK", "BV"],491)492@triton.jit493def gated_chunk_linear_attn_fwd_kernel_o(494 q,495 k,496 v,497 h,498 g,499 o,500 s_qk_h,501 s_qk_t,502 s_qk_d,503 s_vo_h,504 s_vo_t,505 s_vo_d,506 s_h_h,507 s_h_t,508 scale,509 H: tl.constexpr,510 T: tl.constexpr,511 K: tl.constexpr,512 V: tl.constexpr,513 BT: tl.constexpr,514 BK: tl.constexpr,515 BV: tl.constexpr,516 r : tl.constexpr517):518 i_v, i_t, i_bhr = tl.program_id(0), tl.program_id(1), tl.program_id(2)519 i_bh = i_bhr//r520 i_r = i_bhr % r521 rk = K//r522 b_o = tl.zeros([BT, BV], dtype=tl.float32)523 b_s = tl.zeros([BT, BT], dtype=tl.float32)524 for i_k in range(tl.cdiv(K//r, BK)):#这里需要注意拆分#这里K//BK = r525 #问题是不同r_block读取了同一份qk，有影响吗526 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))527 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))528 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))529 b_q = tl.load(p_q, boundary_check=(0, 1))530 b_k = tl.trans(tl.load(p_k, boundary_check=(0, 1)))531 b_h = tl.load(p_h, boundary_check=(0, 1))532 b_o += tl.dot(b_q, b_h)#, allow_tf32=False)533 b_s += tl.dot(b_q, b_k)#, allow_tf32=False)534 535 p_g = tl.make_block_ptr(g + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))536 b_g = tl.load(p_g, boundary_check=(0,))537 b_o = b_o * tl.exp(b_g)[:,None]538 539 b_g_diff = b_g[:, None] - b_g[None, :]540 b_s = b_s * safe_exp(b_g_diff)#BT BT541 542 o_i = tl.arange(0, BT)543 m_s = o_i[:, None] >= o_i[None, :]544 b_s = tl.where(m_s, b_s, 0)#置为0 Bs = 0545 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))546 b_v = tl.load(p_v, boundary_check=(0, 1))547 b_o = b_o * scale + (tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False)) * scale548 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))549 tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))550 551 552 553 554@triton.autotune(555 configs=[556 triton.Config({}, num_warps=1),557 triton.Config({}, num_warps=2),558 triton.Config({}, num_warps=4),559 triton.Config({}, num_warps=8),560 triton.Config({}, num_warps=16)561 ],562 key=["BT", "BK"],563)564@triton.jit565def preprocess_qkw(q,566 k,567 w,568 g,569 q_new,570 k_new,571 w_new,572 T,573 H,574 K,575 r:tl.constexpr,576 BT:tl.constexpr,577 BK:tl.constexpr,578 USE_Q:tl.constexpr,579 ):580 i_k,i_bh,i_t = tl.program_id(0), tl.program_id(1), tl.program_id(2)581 582 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))583 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))584 585 p_g = tl.make_block_ptr(g+i_bh*T,(T,),(1,),(i_t*BT,),(BT,),(0,))586 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))587 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))588 589 last_idx = min((i_t + 1) * BT, T) - 1590 b_g_last = tl.load(g + i_bh*T + last_idx).to(tl.float32) #read BT 位置591 592 b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32)593 b_w = tl.load(p_w, boundary_check=(0, 1)).to(tl.float32)594 b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32)595 b_d_last = tl.exp((b_g_last - b_g))596 b_d_begin = tl.exp(b_g)597 b_k = b_k * b_d_last[:, None]598 b_w = b_w * b_d_begin[:, None]599 tl.store(p_k_new, b_k.to(p_k_new.dtype.element_ty), boundary_check=(0, 1))600 tl.store(p_w_new, b_w.to(p_w_new.dtype.element_ty), boundary_check=(0, 1))601 602 603 if USE_Q:604 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))605 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))606 b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32)607 b_q = b_q * b_d_begin[:, None]608 tl.store(p_q_new, b_q.to(p_q_new.dtype.element_ty), boundary_check=(0, 1))609 610 611#finish612def gated_chunk_fwd_h_fn(k, w, u, g, BT, initial_state, final_state):613 # k, w, u, g, BT, initial_state, final_state614 B, H, T, K, V = *k.shape,u.shape[-1]615 _,_,rT,_ = w.shape616 r = rT//T617 BK = triton.next_power_of_2(K)#直接划分好618 assert BK <= 256, "current kernel does not support head dimension larger than 256."619 BV = 16 if BK > 128 else 32620 BV = 64 if BK <= 64 else BV621 BC = 16 if BK > 128 else 32622 BC = 64 if BK <= 64 else BC623 BC = min(BT, BC)624 NT, NK, NV = triton.cdiv(T, BT), triton.cdiv(K, BK), triton.cdiv(V, BV)625 assert NK == 1626 h = k.new_empty(B, H, NT * K, V)627 628 grid = (NK,B*H,NT)629 k_new = torch.empty_like(k)630 w_new = torch.empty_like(w)631 preprocess_qkw[grid](632 q=None,633 k=k,634 w=w,635 g=g,636 q_new=None,637 k_new=k_new,638 w_new=w_new,639 T=T,640 H=H,641 K=K,642 r=r,643 BT=BT,644 BK=BK,645 USE_Q=False,646 )647 grid = (NK, NV, B * H)648 v_new = torch.empty(B,H,r,T,V,dtype=u.dtype,device=u.device)#做了v_new的r_first649 650 gated_chunk_delta_rule_fwd_kernel_h[grid](#r没有for循环651 k_new,u,w_new, 652 v_new,g,h,653 initial_state,654 final_state,655 H=H, T=T, K=K, V=V, BT=BT, BC=BC, BK=BK, BV=BV, NT=NT,r=r, 656 USE_INITIAL_STATE=initial_state is not None,657 STORE_FINAL_STATE=final_state is not None,658 )659 return h, v_new660 661 662#finish663def gated_chunk_fwd_o_fn(q, k, v_new,h,g,BT):664 B,H,r,T,V,K = *v_new.shape,q.shape[-1]665 BK = triton.next_power_of_2(K//r)666 o = torch.empty_like(v_new)#there_fore,bhr nT,bv667 BK = min(triton.next_power_of_2(K//r), 64)668 BV = min(triton.next_power_of_2(V), 64)669 NV = triton.cdiv(V, BV)670 NT = triton.cdiv(T, BT)671 grid = (NV, NT, B * H * r)672 #h shape b h nk v673 gated_chunk_linear_attn_fwd_kernel_o[grid](674 q, k, v_new, h, g, o,675 T*K, K, 1 ,676 r*T*V,T*V,V,677 NT*K*V,V,678 scale=K**-0.5,679 H=H, T=T, K=K, V=V, BT=BT, BK=BK, BV=BV,r = r,680 )681 o = o.sum(dim=2)#沿着r维度求和682 return o683 684 685@triton.autotune(686 configs=[687 triton.Config({}, num_warps=1),688 triton.Config({}, num_warps=2),689 triton.Config({}, num_warps=4),690 triton.Config({}, num_warps=8),691 triton.Config({}, num_warps=16)692 ],693 key=["BT", "BK", "BV"],694)695@triton.jit696def gated_fwd_prepare_dv_kernel(697 q,698 k,699 g,700 do,701 dv,702 s_qk_h,703 s_qk_t,704 s_qk_d,705 s_vo_h,706 s_vo_t,707 s_vo_d,708 T,709 K,710 V,711 scale,712 BT: tl.constexpr,713 BK: tl.constexpr,714 BV: tl.constexpr,715 r: tl.constexpr,716):717 i_t, i_bhr = tl.program_id(0), tl.program_id(1)#或许也可以r并行718 i_bh = i_bhr//r719 i_r = i_bhr % r720 b_A = tl.zeros([BT, BT], dtype=tl.float32)721 block_r = K//r722 for i_k in range(tl.cdiv(block_r, BK)):723 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))724 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))725 b_k = tl.load(p_k, boundary_check=(0, 1))726 b_q = tl.trans(tl.load(p_q, boundary_check=(0, 1)))727 b_A += tl.dot(b_k, b_q, allow_tf32=False)728 729 p_g = tl.make_block_ptr(g + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))730 b_g = tl.load(p_g, boundary_check=(0,))731 732 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)733 for i_v in range(tl.cdiv(V, BV)):734 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))735 b_do = tl.load(p_do, boundary_check=(0, 1))736 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))737 b_dv = tl.dot(b_A, b_do, allow_tf32=False)738 tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))739 740#finish741def gated_fwd_prepare_dv(q, k, g, do, r,BT):742 B, H, T, K, V = *k.shape, do.shape[-1]743 dv = torch.empty(B,H,r,T,V,device = do.device, dtype= do.dtype)#没法like744 NT = triton.cdiv(T, BT)745 BK = min(triton.next_power_of_2(K//r),64)746 BV = min(triton.next_power_of_2(V), 64)747 gated_fwd_prepare_dv_kernel[(NT, B*H*r)](748 q, k, g , do, dv,749 T*K, K, 1,750 T*V, V, 1,751 T, K, V, K**-0.5, BT, BK, BV, r752 )753 return dv754 755 756 757#finish758@triton.autotune(759 configs=[760 triton.Config({}, num_warps=1),761 triton.Config({}, num_warps=2),762 triton.Config({}, num_warps=4),763 triton.Config({}, num_warps=8),764 triton.Config({}, num_warps=16)765 ],766 key=["BT", "BK", "BV"],767)768@triton.jit769def gated_chunk_delta_rule_bwd_kernel_dhu(770 q,771 k,772 d,773 g,774 do,775 dh,776 dv,777 dv2,778 s_qk_h,779 s_qk_t,780 s_qk_d,781 s_h_h,782 scale,783 H: tl.constexpr,784 T: tl.constexpr,785 K: tl.constexpr,786 V: tl.constexpr,787 BT: tl.constexpr,788 BK: tl.constexpr,789 BV: tl.constexpr,790 NT: tl.constexpr,791 r: tl.constexpr,792 KR: tl.constexpr,793):794 i_k, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)795 b_dh = tl.zeros([BK, BV], dtype=tl.float32)#这个不变 读取所有796 for i_t in range(NT - 1, -1, -1):# 向前偏移了一位,计算流程是对的797 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))798 tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) 799 800 p_q = tl.make_block_ptr(q + i_bh * s_qk_h, (K, T), (s_qk_d, s_qk_t),801 (i_k * BK, i_t * BT), (BK, BT), (0, 1))#全读取802 p_d = tl.make_block_ptr(d + i_bh * (T * K * r), (K,T*r), (1, K),803 (i_k * BK, i_t * BT * r), (BK, BT * r), (0, 1))804 p_do = tl.make_block_ptr(do + i_bh * T * V, (T, V), (V, 1),805 (i_t * BT, i_v * BV), (BT, BV), (1, 0))806 807 last_idx = min((i_t + 1) * BT, T) - 1808 b_glast = tl.load(g + i_bh * T + last_idx)809 b_glast = tl.exp(b_glast)810 811 b_q = (tl.load(p_q, boundary_check=(0, 1)))812 b_q = (b_q * scale).to(b_q.dtype)813 814 b_do = tl.load(p_do, boundary_check=(0, 1))815 b_d = (tl.load(p_d,boundary_check=(0, 1))) 816 p_dv = tl.make_block_ptr(dv + i_bh * r * T * V, (T*r, V), (V , 1),817 (i_t * BT * r, i_v * BV), (BT*r, BV), (1, 0))#load r818 b_dv = tl.load(p_dv, boundary_check=(0, 1))#BT*r Bv 819 b_dhtrans = tl.reshape(b_dh,(r,KR,BV))820 for i_r in range(r):821 rmask = tl.arange(0, r) == i_r #第ir列822 p_k = tl.make_block_ptr(k + i_bh * s_qk_h, (T, K), (s_qk_t, s_qk_d),823 (i_t * BT , i_r*KR + i_k * BK), (BT, KR), (1, 0))# 824 b_k = tl.load(p_k, boundary_check=(0, 1))825 b_dhr = tl.sum(tl.where(rmask[:,None,None],b_dhtrans,0), 0)826 dv_sum = tl.dot(b_k,b_dhr.to(b_k.dtype),allow_tf32=False)827 b_dv += tl.reshape((dv_sum[:,None,:]*rmask[None,:,None]).to(b_dv.dtype),(BT*r,BV))828 829 p_dv2 = tl.make_block_ptr(dv2 + i_bh * r * T * V, (T*r, V), (V , 1),830 (i_t * BT * r, i_v * BV), (BT*r, BV), (1, 0))831 tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))832 833 b_dh *= b_glast834 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)835 836 837 838def gated_chunk_bwd_dhu_fn(q, k, w, g,h0, do, dv, BT):839 B,H,r,T,V,K = *dv.shape,q.shape[-1]840 BK = triton.next_power_of_2(K)841 assert BK <= 256, "current kernel does not support head dimension being larger than 256."842 BV = 16 if BK > 128 else 32843 BV = 64 if BK <= 64 else BV844 845 NT, NK, NV = triton.cdiv(T, BT), triton.cdiv(K, BK), triton.cdiv(V, BV)#感觉可以放并行度846 assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization'847 848 dh = q.new_empty(B, H, NT * K,V)#一样的#need 求和 得一起算849 q_new = torch.empty_like(q)850 k_new = torch.empty_like(k)851 w_new = torch.empty_like(w)852 # grid = (NK,)853 grid = (NK,B*H,NT)854 preprocess_qkw[grid](855 q=q,856 k=k,857 w=w,858 g=g,859 q_new=q_new,860 k_new=k_new,861 w_new=w_new,862 T=T,863 H=H,864 K=K,865 r=r,866 BT=BT,867 BK=BK,868 USE_Q=True,869 )870 871 872 grid = (NK, NV, B * H)873 dv = rearrange(dv,'b h r t v-> b h (t r) v').contiguous()874 dv2 = torch.empty_like(dv)#一样的 #bhr T V875 gated_chunk_delta_rule_bwd_kernel_dhu[grid](876 q_new, k_new, w_new, g, do, dh, dv, dv2,877 T*K,K,1,878 NT*K*V,879 K**-0.5,880 H=H, T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT,r=r,KR = K//r,881 )882 return dh, dv2883 884@triton.autotune(885 configs=[886 triton.Config({}, num_warps=1),887 triton.Config({}, num_warps=2),888 triton.Config({}, num_warps=4),889 triton.Config({}, num_warps=8),890 triton.Config({}, num_warps=16)891 ],892 key=["BT", "BK", "BV"],893)894@triton.jit895def gated_chunk_delta_rule_bwd_kernel_dqkw(896 q,897 k,898 v,899 w,900 g,901 h,902 do,903 dh,904 dq,905 dk,906 dv,907 dw,908 dg,909 s_qk_h,910 s_qk_t,911 s_qk_d,912 s_vo_h,913 s_vo_t,914 s_vo_d,915 s_h_h,916 s_h_t,917 s_g_k,918 scale,919 H: tl.constexpr,920 T: tl.constexpr,921 K: tl.constexpr,922 V: tl.constexpr,923 BT: tl.constexpr,924 BK: tl.constexpr,925 BV: tl.constexpr,926 NT: tl.constexpr,927 r: tl.constexpr,928):929 i_k, i_t, i_bhr = tl.program_id(0), tl.program_id(1), tl.program_id(2)930 i_r = i_bhr%r931 i_bh = i_bhr//r932 o_i = tl.arange(0, BT)933 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))934 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))935 b_dq = tl.zeros([BT, BK], dtype=tl.float32)936 b_dk = tl.zeros([BT, BK], dtype=tl.float32)937 b_dw = tl.zeros([BT*r,BK], dtype=tl.float32)938 b_ds = tl.zeros([BT, BT], dtype=tl.float32)939 b_dg_last = tl.zeros([1,],dtype=tl.float32)940 941 for i_v in range(tl.cdiv(V, BV)):942 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))943 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))944 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))945 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))946 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))947 948 b_v = tl.load(p_v, boundary_check=(0, 1))949 b_do = tl.load(p_do, boundary_check=(0, 1))950 b_h = (tl.load(p_h, boundary_check=(0, 1)))#BV BK951 b_dh =(tl.load(p_dh, boundary_check=(0, 1)))952 953 b_dg_last += tl.sum(b_h * b_dh)954 955 b_ds += tl.dot(b_do, tl.trans(b_v), allow_tf32=False)#ok 956 b_dq += tl.dot(b_do, b_h, allow_tf32=False)#d_do 全， bh应该包含 i_Kbufen957 b_dk += tl.dot(b_v, b_dh, allow_tf32=False)#用来计算dk,yes 行独立没问题958 b_dv = (tl.load(p_dv, boundary_check=(0, 1)))#BT*r BV959 b_dw += (tl.dot(b_dv.to(b_v.dtype),b_h.to(b_v.dtype))) #get BT*r BK960 961 b_q = tl.load(p_q, boundary_check=(0, 1))962 b_k = tl.load(p_k, boundary_check=(0, 1))963 964 b_dg = tl.zeros([BT,], dtype=tl.float32)965 p_g = tl.make_block_ptr(g + i_bh * T ,(T,),(1,),(i_t*BT,),(BT,),(0,))966 b_g = tl.load(p_g,boundary_check=(0,))967 b_glast = tl.load(g +i_bh*T + (min(i_t * BT + BT, T) - 1))968 b_dg_last *= tl.exp(b_glast)969 970 971 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))972 b_w = tl.load(p_w,boundary_check=(0,1))#BT * r ,BK973 b_dw = b_dw * tl.reshape(tl.broadcast_to(tl.reshape(tl.exp(b_g),(BT,1)),(BT,r)),(BT*r))[:,None]974 b_dg -= tl.sum(tl.reshape(b_w*b_dw,(BT,r*BK)),-1)975 976 b_dq = b_dq*scale*tl.exp(b_g)[:,None] 977 b_dg += tl.sum(b_dq*tl.trans(b_q),1)#BT*BK978 979 b_dk = b_dk * safe_exp(b_glast-b_g)[:,None]980 b_dg -= tl.sum(b_dk*b_k,1)#BT*BK981 982 b_dg_last += tl.sum(b_dk*b_k)983 984 b_ds = tl.where(o_i[:, None] >= o_i[None, :], b_ds* safe_exp(b_g[:, None] - b_g[None, :]) * scale, 0)985 b_ds2 = b_ds*(tl.dot(tl.trans(b_q),tl.trans(b_k)))986 987 b_dg += tl.sum(b_ds2,axis=1)988 b_dg -= tl.sum(b_ds2,axis=0)989 b_ds = b_ds.to(b_k.dtype)990 991 b_dq += tl.dot(b_ds, b_k, allow_tf32=False)992 b_dk += tl.trans(tl.dot(b_q, b_ds, allow_tf32=False)) #这些应该没啥问题993 994 995 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))996 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))997 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))998 p_dg = tl.make_block_ptr(dg + i_k * s_g_k + i_bh * T,(T,),(1,),(i_t*BT,),(BT,),(0,))999 b_dg = tl.where(o_i<min(BT, T-i_t*BT) - 1, b_dg, b_dg + b_dg_last)1000 1001 tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))1002 tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))1003 tl.store(p_dw, ((-b_dw.to(p_dw.dtype.element_ty))), boundary_check=(0, 1))1004 tl.store(p_dg,b_dg.to(p_dg.dtype.element_ty),boundary_check=(0,))1005 1006 1007 1008def gated_chunk_bwd_dqkw_fn(q, k, v_new, w, g, h, du, do, dh, BT):1009 B, H, T, K, V = *q.shape, v_new.shape[-1]1010 _,_,RT,_ = w.shape1011 r = RT // T1012 #最后一个函数，计算dw,dq,dk1013 BK = triton.next_power_of_2(K//r)#需要更细粒度的划分，确保不会使得 不同位置的划到一起1014 BK = min(triton.next_power_of_2(K//r), 64)1015 BV = min(triton.next_power_of_2(V), 64)1016 NK = triton.cdiv(K//r, BK)1017 NT = triton.cdiv(T, BT)1018 grid = (NK, NT, B * H * r)#通过NK控制位置1019 dq = torch.empty_like(q)1020 dk = torch.empty_like(k)#k_org1021 dw = torch.empty_like(w)#bh nt k1022 dg = torch.empty(NK,*g.shape,dtype=torch.float32,device=g.device)1023 1024 gated_chunk_delta_rule_bwd_kernel_dqkw[grid](1025 q, k, v_new, w, g, h, do, dh, dq, dk, du, dw,dg,1026 T*K,K,1,1027 T*V, V, 1,1028 NT*K*V,V,1029 B*H*T,1030 scale=K ** -0.5,1031 H=H, T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT,r = r1032 )1033 dg = dg.sum(0)1034 return dq.to(q.dtype), dk.to(k.dtype), dw.to(w.dtype),dg1035 1036#compute this 1037@triton.autotune(1038 configs=[1039 triton.Config({}, num_warps=1),1040 triton.Config({}, num_warps=2),1041 triton.Config({}, num_warps=4),1042 triton.Config({}, num_warps=8),1043 triton.Config({}, num_warps=16)1044 ],1045 key=["BT", "BK", "BV"],1046)1047@triton.jit1048def gated_bwd_prepare_wy_repr_kernel( 1049 k, v, beta,mask_ij,g_cumsum,Aw,Au,1050 dw, du,1051 dk, dv, dbeta,dmask,dg,1052 s_qk_h,1053 s_qk_t,1054 s_qk_d,1055 s_vo_h,1056 s_vo_t,1057 s_vo_d,1058 T,1059 K,1060 V,1061 r: tl.constexpr,1062 BT: tl.constexpr,1063 BK: tl.constexpr,1064 BV: tl.constexpr1065):1066 i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)1067 p_A = 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))1068 b_A = tl.load(p_A, boundary_check=(0, 1)).to(k.dtype.element_ty)1069 b_dbeta = tl.zeros([BT], dtype=tl.float32)1070 b_dA = tl.zeros([BT*r,BT*r], dtype=tl.float32)1071 p_beta = tl.make_block_ptr(beta + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))1072 b_beta = tl.load(p_beta, boundary_check=(0,))1073 b_dmask = tl.zeros([r,r],dtype=tl.float32)1074 block_k = K//r1075 for i_r in range(r):1076 p_mask = mask_ij + tl.arange(0,r)*r + i_r#读取第ir列1077 b_mask = tl.load(p_mask)#第r列1078 rmask = tl.arange(0, r) == i_r #第r列1079 for i_k in range(tl.cdiv(block_k, BK)):1080 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))1081 b_k = tl.load(p_k, boundary_check=(0, 1))1082 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))1083 b_k_beta = ((b_k * b_beta[:, None])[:,None,:]*b_mask[None,:,None]).to(b_k.dtype)#BT*r*d1084 b_k_beta = tl.reshape(b_k_beta,(BT*r,BK))1085 b_dw = tl.load(p_dw, boundary_check=(0, 1))1086 b_dA += tl.dot(b_dw, tl.trans(b_k_beta), allow_tf32=False)1087 b_dk_beta = tl.dot(tl.trans(b_A), b_dw, allow_tf32=False)1088 b_dk_beta = tl.reshape(b_dk_beta,(BT,r,BK))1089 sum_dk = tl.sum(b_dk_beta * b_mask[None,:,None],1)1090 b_dk = sum_dk* b_beta[:, None]1091 b_dbeta += tl.sum(sum_dk * b_k, 1)1092 b_ss = (tl.sum(tl.sum(b_dk_beta * b_beta[:,None,None] * b_k[:,None,:],0),-1))1093 b_dmask += (b_ss[:,None]*rmask[None,:]).to(tl.float32)1094 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))1095 tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))1096 1097 i = tl.arange(0, BT * r)[:, None]1098 j = tl.arange(0, BT * r)[None, :]1099 iB = i // r1100 jB = j // r1101 da_mask = iB > jB1102 b_dA = tl.where(da_mask, b_dA, 0)1103 b_dA = tl.dot(b_dA.to(b_A.dtype), tl.trans(b_A), allow_tf32=False)1104 b_dA = tl.dot(tl.trans(b_A), b_dA.to(b_A.dtype), allow_tf32=False)1105 b_dA = tl.where(da_mask, -b_dA, 0) #等价于 kkt的 dA 很多0，对角处1106 b_dA = tl.reshape(b_dA,(BT,r,BT,r))1107 1108 1109 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))1110 b_A = tl.load(p_A, boundary_check=(0, 1)).to(k.dtype.element_ty)1111 b_dA2 = tl.zeros([BT*r,BT*r], dtype=tl.float32)1112 1113 for i_v in range(tl.cdiv(V, BV)):#分块r 1114 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))1115 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 BV1116 b_v = tl.load(p_v, boundary_check=(0, 1))1117 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*BV1118 b_v_beta = tl.reshape(b_v_beta,(BT*r,BV))1119 b_du = tl.load(p_du, boundary_check=(0, 1))1120 b_dA2 += tl.dot(b_du, tl.trans(b_v_beta), allow_tf32=False)#BT*r,BT*r1121 b_dv_beta = tl.dot(tl.trans(b_A), b_du, allow_tf32=False)#BT*r,BV1122 b_dv_beta = tl.reshape(b_dv_beta,(BT,r,BV))#1123 sum_dv = tl.sum(b_dv_beta,-2)#这里不一样，结果1124 b_dv = (sum_dv * b_beta[:, None])#？哪一步结果不一样呢1125 b_dbeta += tl.sum(sum_dv * b_v, 1)1126 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))1127 tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))1128 1129 1130 b_dA2 = tl.where(da_mask, b_dA2, 0)1131 b_dA2 = tl.dot(b_dA2.to(b_A.dtype), tl.trans(b_A), allow_tf32=False)1132 b_dA2 = tl.dot(tl.trans(b_A), b_dA2.to(b_A.dtype), allow_tf32=False)1133 b_dA2 = tl.where(da_mask, -b_dA2, 0) #等价于 kkt的 dA 很多0，对角处1134 b_dA2 = tl.reshape(b_dA2,(BT,r,BT,r))1135 1136 1137 p_g = tl.make_block_ptr(g_cumsum + i_bh*T,(T,),(1,),(i_t*BT,),(BT,),(0,))1138 b_g = tl.load(p_g,boundary_check=(0,))1139 b_dA2 *= safe_exp(b_g[:,None]-b_g[None,:])[:,None,:,None]1140 b_dA += b_dA21141 1142 b_dA2 = tl.permute(b_dA2,(0,2,1,3))#Bt bt r r1143 b_A = tl.zeros([BT,BT,r,r], dtype=tl.float32)1144 1145 for i_r in range(r):#只取ir项 1146 p_mask = mask_ij + tl.arange(0,r)*r+i_r#读取第ir列1147 b_mask = tl.load(p_mask)#第ir列1148 rmask = tl.arange(0, r) == i_r #第ir列1149 g = tl.sum(tl.where(rmask[None,None,None,:], b_dA, 0), -1)#BT r BT #取出第ir列1150 ir_A = tl.sum(g * b_mask[None,:,None],1).to(k.dtype.element_ty)#BT BT1151 1152 for i_k in range(tl.cdiv(block_k, BK)):#ik = 11153 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))1154 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))1155 b_k = tl.load(p_k, boundary_check=(0, 1))1156 b_dk = tl.load(p_dk, boundary_check=(0, 1))1157 b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype)#BT*BK1158 1159 b_dk_beta = tl.dot(ir_A, b_k, allow_tf32=False)1160 b_dbeta += tl.sum(b_dk_beta * b_k, 1)1161 b_dk += tl.dot(tl.trans(ir_A), b_k_beta, allow_tf32=False)1162 b_dk += b_dk_beta * b_beta[:, None]1163 tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))1164 1165 beta_kkt = (tl.dot(b_k_beta,tl.trans(b_k), allow_tf32=False))#BT BT1166 b_A += beta_kkt[:,:,None,None] * (rmask[:,None] * b_mask[None,:])[None,None,:,:]1167 1168 betas = tl.sum(tl.sum(beta_kkt[:,None,:]*g,-1),0)1169 b_dmask += (betas[:,None]*rmask[None,:]).to(tl.float32)1170 1171 p_dbeta = tl.make_block_ptr(dbeta + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))1172 tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,))1173 1174 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))1175 tl.store(p_dmask, b_dmask.to(p_dmask.dtype.element_ty), boundary_check=(0,1))1176 1177 b_dA2 *= b_A #BT BT r r1178 b_dA2 = tl.sum(tl.reshape(b_dA2,(BT,BT,r*r)),-1)1179 1180 b_dg = tl.sum(b_dA2,1)-tl.sum(b_dA2,0)1181 p_dg = tl.make_block_ptr(dg+i_bh*T,(T,),(1,),(i_t*BT,),(BT,),(0,))1182 tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,))1183 1184 1185 1186def gated_bwd_prepare_wy_repr(k, v, beta, mask,g, Aw,Au, dw, du, BT):1187 B, H, T, K, V = *k.shape, v.shape[-1]1188 r = mask.shape[-1]1189 NT = triton.cdiv(T, BT)1190 BK = min(triton.next_power_of_2(K//r), 64)1191 BV = min(triton.next_power_of_2(V), 64)1192 NT = triton.cdiv(T, BT)1193 dk = torch.empty_like(k)1194 dv = torch.empty_like(v).contiguous()1195 dbeta = torch.zeros_like(beta)1196 dg = torch.empty_like(g)1197 dmask = torch.zeros([B*H*NT,r,r],device=k.device,dtype=k.dtype).contiguous()1198 assert BK <= K//r1199 gated_bwd_prepare_wy_repr_kernel[(NT, B*H)](1200 k, v, beta, mask, g, Aw,Au,