Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / _triton / gated_delta_net / chunk_delta_h.py: 0%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang 

16 

17# pylint: disable=line-too-long,missing-public-type-hints,missing-public-docstring 

18# pylint: disable=used-before-assignment,unsupported-binary-operation,unused-argument 

19# pylint: disable=invalid-name,missing-module-docstring,missing-function-docstring 

20# pylint: disable=forbidden-backend-import 

21 

22from typing import Optional, Tuple 

23 

24import torch 

25import triton 

26import triton.language as tl 

27 

28from .utils import prepare_chunk_indices, prepare_chunk_offsets, get_autotune_config, get_npu_properties 

29 

30CUBE_CORE_NUM = get_npu_properties()['num_aicore'] 

31 

32 

33@triton.heuristics({ 

34 'USE_G': lambda args: args['g'] is not None, 

35 'USE_GK': lambda args: args['gk'] is not None, 

36 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 

37 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, 

38 'SAVE_NEW_VALUE': lambda args: args['v_new'] is not None, 

39 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, 

40}) 

41@triton.autotune( 

42 configs=get_autotune_config(multibuffer_list=(False,)), 

43 key=['H', 'K', 'V', 'BT'], 

44) 

45@triton.jit(do_not_specialize=['T']) 

46def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( 

47 k, 

48 v, 

49 w, 

50 v_new, 

51 g, 

52 gk, 

53 h, 

54 h0, 

55 ht, 

56 cu_seqlens, 

57 chunk_offsets, 

58 T, 

59 H: tl.constexpr, 

60 K: tl.constexpr, 

61 V: tl.constexpr, 

62 BT: tl.constexpr, 

63 BV: tl.constexpr, 

64 NT: tl.constexpr, 

65 USE_G: tl.constexpr, 

66 USE_GK: tl.constexpr, 

67 USE_INITIAL_STATE: tl.constexpr, 

68 STORE_FINAL_STATE: tl.constexpr, 

69 SAVE_NEW_VALUE: tl.constexpr, 

70 IS_VARLEN: tl.constexpr, 

71): 

72 T_all = T 

73 NT_all = NT 

74 i_v, i_nh = tl.program_id(0), tl.program_id(1) 

75 i_n, i_h = i_nh // H, i_nh % H 

76 if IS_VARLEN: 

77 bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) 

78 T = eos - bos 

79 NT = tl.cdiv(T, BT) 

80 boh = tl.load(chunk_offsets + i_n).to(tl.int32) 

81 else: 

82 bos, eos = i_n * T, i_n * T + T 

83 NT = tl.cdiv(T, BT) 

84 boh = i_n * NT 

85 

86 # Initialize hidden states 

87 b_h1 = tl.zeros([64, BV], dtype=tl.float32) 

88 if K > 64: 

89 b_h2 = tl.zeros([64, BV], dtype=tl.float32) 

90 if K > 128: 

91 b_h3 = tl.zeros([64, BV], dtype=tl.float32) 

92 if K > 192: 

93 b_h4 = tl.zeros([64, BV], dtype=tl.float32) 

94 

95 if IS_VARLEN: 

96 v = v + (i_h * T_all + bos) * V 

97 k = k + (i_h * T_all + bos) * K 

98 w = w + (i_h * T_all + bos) * K 

99 g = g + i_h * T_all + bos 

100 h = h + (i_h * NT_all + boh) * K * V 

101 if SAVE_NEW_VALUE: 

102 v_new_base = v_new + (i_h * T_all + bos) * V 

103 else: 

104 v = v + (i_n * H + i_h) * T * V 

105 k = k + (i_n * H + i_h) * T * K 

106 w = w + (i_n * H + i_h) * T * K 

107 g = g + (i_n * H + i_h) * T 

108 h = h + (i_n * H + i_h) * NT * K * V 

109 if SAVE_NEW_VALUE: 

110 v_new_base = v_new + (i_n * H + i_h) * T * V 

111 

112 if USE_INITIAL_STATE: 

113 h0_ptr = h0 + i_nh * K * V 

114 if STORE_FINAL_STATE: 

115 ht_ptr = ht + i_nh * K * V 

116 

117 # Load initial state 

118 if USE_INITIAL_STATE: 

119 p_h0_1 = tl.make_block_ptr(h0_ptr, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) 

120 b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) 

121 if K > 64: 

122 p_h0_2 = tl.make_block_ptr(h0_ptr, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) 

123 b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) 

124 if K > 128: 

125 p_h0_3 = tl.make_block_ptr(h0_ptr, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) 

126 b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) 

127 if K > 192: 

128 p_h0_4 = tl.make_block_ptr(h0_ptr, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) 

129 b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) 

130 

131 # Main recurrence over chunks 

132 for i_t in range(NT): 

133 # Store current hidden state h_t 

134 p_h1 = tl.make_block_ptr(h + i_t * K * V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) 

135 tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) 

136 if K > 64: 

137 p_h2 = tl.make_block_ptr(h + i_t * K * V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) 

138 tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) 

139 if K > 128: 

140 p_h3 = tl.make_block_ptr(h + i_t * K * V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) 

141 tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) 

142 if K > 192: 

143 p_h4 = tl.make_block_ptr(h + i_t * K * V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) 

144 tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) 

145 

146 # Compute v_residual = v - w @ h 

147 p_w = tl.make_block_ptr(w, (T, K), (K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) 

148 b_w = tl.load(p_w, boundary_check=(0, 1)) 

149 b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) 

150 if K > 64: 

151 p_w = tl.make_block_ptr(w, (T, K), (K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) 

152 b_w = tl.load(p_w, boundary_check=(0, 1)) 

153 b_v += tl.dot(b_w, b_h2.to(b_w.dtype)) 

154 if K > 128: 

155 p_w = tl.make_block_ptr(w, (T, K), (K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) 

156 b_w = tl.load(p_w, boundary_check=(0, 1)) 

157 b_v += tl.dot(b_w, b_h3.to(b_w.dtype)) 

158 if K > 192: 

159 p_w = tl.make_block_ptr(w, (T, K), (K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) 

160 b_w = tl.load(p_w, boundary_check=(0, 1)) 

161 b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) 

162 

163 p_v = tl.make_block_ptr(v, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) 

164 b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v 

165 

166 if SAVE_NEW_VALUE: 

167 p_v_new = tl.make_block_ptr(v_new_base, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) 

168 tl.store(p_v_new, b_v.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) 

169 

170 last_idx = min((i_t + 1) * BT, T) - 1 

171 

172 # Apply output gate g 

173 if USE_G: 

174 m_t = (i_t * BT + tl.arange(0, BT)).to(tl.float32) < T 

175 b_g_last = tl.load(g + last_idx) 

176 p_g = tl.make_block_ptr(g, (T,), (1,), (i_t * BT,), (BT,), (0,)) 

177 b_g = tl.load(p_g, boundary_check=(0,)) 

178 b_v *= (m_t * tl.exp(b_g_last - b_g))[:, None] 

179 b_g_last_exp = tl.exp(b_g_last) 

180 b_h1 *= b_g_last_exp 

181 if K > 64: 

182 b_h2 *= b_g_last_exp 

183 if K > 128: 

184 b_h3 *= b_g_last_exp 

185 if K > 192: 

186 b_h4 *= b_g_last_exp 

187 

188 # Apply key gate gk 

189 if USE_GK: 

190 o_k1 = tl.arange(0, 64).to(tl.float32) 

191 gk_base_ptr = gk + (i_n * H + i_h) * T * K 

192 b_gk_last1 = tl.load(gk_base_ptr + last_idx * K + o_k1, mask=(o_k1 < K), other=0.) 

193 b_h1 *= tl.exp(b_gk_last1)[:, None] 

194 if K > 64: 

195 o_k2 = 64 + o_k1 

196 b_gk_last2 = tl.load(gk_base_ptr + last_idx * K + o_k2, mask=(o_k2 < K), other=0.) 

197 b_h2 *= tl.exp(b_gk_last2)[:, None] 

198 if K > 128: 

199 o_k3 = 128 + o_k1 

200 b_gk_last3 = tl.load(gk_base_ptr + last_idx * K + o_k3, mask=(o_k3 < K), other=0.) 

201 b_h3 *= tl.exp(b_gk_last3)[:, None] 

202 if K > 192: 

203 o_k4 = 192 + o_k1 

204 b_gk_last4 = tl.load(gk_base_ptr + last_idx * K + o_k4, mask=(o_k4 < K), other=0.) 

205 b_h4 *= tl.exp(b_gk_last4)[:, None] 

206 

207 b_v = b_v.to(k.dtype.element_ty) 

208 

209 # Update hidden state: h += k @ v 

210 p_k = tl.make_block_ptr(k, (K, T), (1, K), (0, i_t * BT), (64, BT), (0, 1)) 

211 b_k = tl.load(p_k, boundary_check=(0, 1)) 

212 if USE_GK: 

213 p_gk = tl.make_block_ptr(gk_base_ptr, (K, T), (1, K), (0, i_t * BT), (64, BT), (0, 1)) 

214 b_k = (b_k * tl.exp(b_gk_last1[:, None] - tl.load(p_gk, boundary_check=(0, 1)))).to(b_k.dtype) 

215 b_h1 += tl.dot(b_k, b_v) 

216 

217 if K > 64: 

218 p_k = tl.make_block_ptr(k, (K, T), (1, K), (64, i_t * BT), (64, BT), (0, 1)) 

219 b_k = tl.load(p_k, boundary_check=(0, 1)) 

220 if USE_GK: 

221 p_gk = tl.make_block_ptr(gk_base_ptr, (K, T), (1, K), (64, i_t * BT), (64, BT), (0, 1)) 

222 b_k = (b_k * tl.exp(b_gk_last2[:, None] - tl.load(p_gk, boundary_check=(0, 1)))).to(b_k.dtype) 

223 b_h2 += tl.dot(b_k, b_v) 

224 

225 if K > 128: 

226 p_k = tl.make_block_ptr(k, (K, T), (1, K), (128, i_t * BT), (64, BT), (0, 1)) 

227 b_k = tl.load(p_k, boundary_check=(0, 1)) 

228 if USE_GK: 

229 p_gk = tl.make_block_ptr(gk_base_ptr, (K, T), (1, K), (128, i_t * BT), (64, BT), (0, 1)) 

230 b_k = (b_k * tl.exp(b_gk_last3[:, None] - tl.load(p_gk, boundary_check=(0, 1)))).to(b_k.dtype) 

231 b_h3 += tl.dot(b_k, b_v) 

232 

233 if K > 192: 

234 p_k = tl.make_block_ptr(k, (K, T), (1, K), (192, i_t * BT), (64, BT), (0, 1)) 

235 b_k = tl.load(p_k, boundary_check=(0, 1)) 

236 if USE_GK: 

237 p_gk = tl.make_block_ptr(gk_base_ptr, (K, T), (1, K), (192, i_t * BT), (64, BT), (0, 1)) 

238 b_k = (b_k * tl.exp(b_gk_last4[:, None] - tl.load(p_gk, boundary_check=(0, 1)))).to(b_k.dtype) 

239 b_h4 += tl.dot(b_k, b_v) 

240 

241 # Store final state 

242 if STORE_FINAL_STATE: 

243 p_ht = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) 

244 tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) 

245 if K > 64: 

246 p_ht = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) 

247 tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) 

248 if K > 128: 

249 p_ht = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) 

250 tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) 

251 if K > 192: 

252 p_ht = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) 

253 tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) 

254 

255 

256def chunk_gated_delta_rule_fwd_h( 

257 k: torch.Tensor, 

258 w: torch.Tensor, 

259 u: torch.Tensor, 

260 g: Optional[torch.Tensor] = None, 

261 gk: Optional[torch.Tensor] = None, 

262 initial_state: Optional[torch.Tensor] = None, 

263 output_final_state: bool = False, 

264 chunk_size: int = 64, # default:64 

265 save_new_value: bool = True, 

266 cu_seqlens: Optional[torch.LongTensor] = None, 

267) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: 

268 B, T, H, K, V = *k.shape, u.shape[-1] 

269 BT = chunk_size 

270 

271 chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None 

272 # N: the actual number of sequences in the batch with either equal or variable lengths 

273 if cu_seqlens is None: 

274 N, NT, chunk_offsets = B, triton.cdiv(T, BT), None 

275 else: 

276 N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) 

277 if K > 256: 

278 raise ValueError("current kernel does not support head dimension larger than 256.") 

279 

280 h = k.new_empty(B, NT, H, K, V).permute(0, 2, 1, 3, 4).contiguous() 

281 final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None 

282 

283 BV = 128 

284 

285 v_new = torch.empty_like(u).permute(0, 2, 1, 3).contiguous() if save_new_value else None 

286 k = k.permute(0, 2, 1, 3).contiguous() 

287 w = w.permute(0, 2, 1, 3).contiguous() 

288 u = u.permute(0, 2, 1, 3).contiguous() 

289 g = g.permute(0, 2, 1).contiguous() 

290 chunk_gated_delta_rule_fwd_kernel_h_blockdim64[(triton.cdiv(V, BV), N * H)]( 

291 k=k, 

292 v=u, 

293 w=w, 

294 v_new=v_new, 

295 g=g, 

296 gk=gk, 

297 h=h, 

298 h0=initial_state, 

299 ht=final_state, 

300 cu_seqlens=cu_seqlens, 

301 chunk_offsets=chunk_offsets, 

302 T=T, 

303 H=H, 

304 K=K, 

305 V=V, 

306 BT=BT, 

307 BV=BV, 

308 NT=NT, 

309 ) 

310 h = h.permute(0, 2, 1, 3, 4).contiguous() 

311 v_new = v_new.permute(0, 2, 1, 3).contiguous() 

312 return h, v_new, final_state 

313 

314 

315@triton.heuristics({ 

316 'USE_G': lambda args: args['g'] is not None, 

317 'USE_GK': lambda args: args['gk'] is not None, 

318 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, 

319 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, 

320 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, 

321}) 

322@triton.autotune( 

323 configs=get_autotune_config(multibuffer_list=(True, False)), 

324 key=['H', 'K', 'V', 'BT', 'BV', 'USE_G', 'IS_VARLEN'], 

325) 

326@triton.jit(do_not_specialize=['T']) 

327def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( 

328 q, 

329 k, 

330 w, 

331 g, 

332 gk, 

333 dht, 

334 dh0, 

335 do, 

336 dh, 

337 dv, 

338 dv2, 

339 cu_seqlens, 

340 chunk_offsets, 

341 scale, 

342 T, 

343 H: tl.constexpr, 

344 K: tl.constexpr, 

345 V: tl.constexpr, 

346 BT: tl.constexpr, 

347 BV: tl.constexpr, 

348 USE_G: tl.constexpr, 

349 USE_GK: tl.constexpr, 

350 USE_INITIAL_STATE: tl.constexpr, 

351 USE_FINAL_STATE_GRADIENT: tl.constexpr, 

352 IS_VARLEN: tl.constexpr, 

353): 

354 T_all = T 

355 i_v, i_nh = tl.program_id(0), tl.program_id(1) 

356 i_n, i_h = i_nh // H, i_nh % H 

357 if IS_VARLEN: 

358 bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) 

359 T = eos - bos 

360 NT = tl.cdiv(T, BT) 

361 boh = tl.load(chunk_offsets + i_n).to(tl.int32) 

362 else: 

363 bos, eos = i_n * T, i_n * T + T 

364 NT = tl.cdiv(T, BT) 

365 boh = i_n * NT 

366 

367 b_dh1 = tl.zeros([64, BV], dtype=tl.float32) 

368 if K > 64: 

369 b_dh2 = tl.zeros([64, BV], dtype=tl.float32) 

370 if K > 128: 

371 b_dh3 = tl.zeros([64, BV], dtype=tl.float32) 

372 if K > 192: 

373 b_dh4 = tl.zeros([64, BV], dtype=tl.float32) 

374 

375 q += (bos * H + i_h) * K 

376 k += (bos * H + i_h) * K 

377 w += (bos * H + i_h) * K 

378 do += (bos * H + i_h) * V 

379 dv += (bos * H + i_h) * V 

380 dv2 += (bos * H + i_h) * V 

381 dh += (boh * H + i_h) * K * V 

382 if USE_GK: 

383 gk += (bos * H + i_h) * K 

384 

385 if USE_INITIAL_STATE: 

386 dh0 += i_nh * K * V 

387 if USE_FINAL_STATE_GRADIENT: 

388 dht += i_nh * K * V 

389 

390 stride_v = H * V 

391 stride_h = H * K * V 

392 stride_k = H * K 

393 

394 if USE_FINAL_STATE_GRADIENT: 

395 p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) 

396 b_dh1 += tl.load(p_dht1, boundary_check=(0, 1)) 

397 if K > 64: 

398 p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) 

399 b_dh2 += tl.load(p_dht2, boundary_check=(0, 1)) 

400 if K > 128: 

401 p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) 

402 b_dh3 += tl.load(p_dht3, boundary_check=(0, 1)) 

403 if K > 192: 

404 p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) 

405 b_dh4 += tl.load(p_dht4, boundary_check=(0, 1)) 

406 

407 for i_t in range(NT - 1, -1, -1): 

408 p_dh1 = tl.make_block_ptr(dh + i_t * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) 

409 tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) 

410 if K > 64: 

411 p_dh2 = tl.make_block_ptr(dh + i_t * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) 

412 tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) 

413 if K > 128: 

414 p_dh3 = tl.make_block_ptr(dh + i_t * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) 

415 tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) 

416 if K > 192: 

417 p_dh4 = tl.make_block_ptr(dh + i_t * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) 

418 tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) 

419 

420 last_idx = min((i_t + 1) * BT, T) - 1 

421 if USE_G: 

422 if IS_VARLEN: 

423 bos_g = i_h * T_all + bos 

424 else: 

425 bos_g = (i_n * H + i_h) * T_all 

426 bg_last = tl.load(g + bos_g + last_idx) 

427 bg_last_exp = tl.exp(bg_last) 

428 p_g = tl.make_block_ptr(base=g + bos_g, shape=(T,), strides=(1,), offsets=(i_t * BT,), block_shape=(BT,), order=(0,)) 

429 b_g = tl.load(p_g, boundary_check=(0,)) 

430 b_g_exp = tl.exp(b_g) 

431 

432 p_dv = tl.make_block_ptr(dv, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) 

433 p_dv2 = tl.make_block_ptr(dv2, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) 

434 p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) 

435 

436 b_do = tl.load(p_do, boundary_check=(0, 1)) 

437 

438 # Update dv 

439 p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) 

440 b_k = tl.load(p_k, boundary_check=(0, 1)) 

441 if USE_GK: 

442 o_k1 = tl.arange(0, 64) 

443 b_gk_last1 = tl.load(gk + last_idx * H * K + o_k1, mask=(o_k1 < K), other=0.) 

444 b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) 

445 

446 if K > 64: 

447 p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) 

448 b_k = tl.load(p_k, boundary_check=(0, 1)) 

449 if USE_GK: 

450 o_k2 = 64 + o_k1 

451 b_gk_last2 = tl.load(gk + last_idx * H * K + o_k2, mask=(o_k2 < K), other=0.) 

452 b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) 

453 

454 if K > 128: 

455 p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) 

456 b_k = tl.load(p_k, boundary_check=(0, 1)) 

457 if USE_GK: 

458 o_k3 = 128 + o_k1 

459 b_gk_last3 = tl.load(gk + last_idx * H * K + o_k3, mask=(o_k3 < K), other=0.) 

460 b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) 

461 

462 if K > 192: 

463 p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) 

464 b_k = tl.load(p_k, boundary_check=(0, 1)) 

465 if USE_GK: 

466 o_k4 = 192 + o_k1 

467 b_gk_last4 = tl.load(gk + last_idx * H * K + o_k4, mask=(o_k4 < K), other=0.) 

468 b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) 

469 

470 if USE_G: 

471 m_t = (i_t * BT + tl.arange(0, BT)).to(tl.float32) < T 

472 b_dv *= (m_t * tl.exp(bg_last - b_g))[:, None] 

473 b_dv += tl.load(p_dv, boundary_check=(0, 1)) 

474 

475 tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) 

476 # Update dh 

477 p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) 

478 p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) 

479 b_w = tl.load(p_w, boundary_check=(0, 1)) 

480 b_q = tl.load(p_q, boundary_check=(0, 1)) 

481 if USE_G: 

482 b_dh1 *= bg_last_exp 

483 b_q = b_q * b_g_exp[None, :] 

484 if USE_GK: 

485 b_dh1 *= tl.exp(b_gk_last1[:, None])