Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / custom_ops / gdn / triton / 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# -*- coding: utf-8 -*- 

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

17 

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

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

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

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 assert K <= 256, "current kernel does not support head dimension larger than 256." 

278 

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

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

281 

282 BV = 128 

283 

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

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

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

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

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

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

290 k=k, 

291 v=u, 

292 w=w, 

293 v_new=v_new, 

294 g=g, 

295 gk=gk, 

296 h=h, 

297 h0=initial_state, 

298 ht=final_state, 

299 cu_seqlens=cu_seqlens, 

300 chunk_offsets=chunk_offsets, 

301 T=T, 

302 H=H, 

303 K=K, 

304 V=V, 

305 BT=BT, 

306 BV=BV, 

307 NT=NT, 

308 ) 

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

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

311 return h, v_new, final_state 

312 

313 

314@triton.heuristics({ 

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

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

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

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

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

320}) 

321@triton.autotune( 

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

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

324) 

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

326def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( 

327 q, 

328 k, 

329 w, 

330 g, 

331 gk, 

332 dht, 

333 dh0, 

334 do, 

335 dh, 

336 dv, 

337 dv2, 

338 cu_seqlens, 

339 chunk_offsets, 

340 scale, 

341 T, 

342 H: tl.constexpr, 

343 K: tl.constexpr, 

344 V: tl.constexpr, 

345 BT: tl.constexpr, 

346 BV: tl.constexpr, 

347 USE_G: tl.constexpr, 

348 USE_GK: tl.constexpr, 

349 USE_INITIAL_STATE: tl.constexpr, 

350 USE_FINAL_STATE_GRADIENT: tl.constexpr, 

351 IS_VARLEN: tl.constexpr, 

352): 

353 T_all = T 

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

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

356 if IS_VARLEN: 

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

358 T = eos - bos 

359 NT = tl.cdiv(T, BT) 

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

361 else: 

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

363 NT = tl.cdiv(T, BT) 

364 boh = i_n * NT 

365 

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

367 if K > 64: 

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

369 if K > 128: 

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

371 if K > 192: 

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

373 

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

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

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

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

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

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

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

381 if USE_GK: 

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

383 

384 if USE_INITIAL_STATE: 

385 dh0 += i_nh * K * V 

386 if USE_FINAL_STATE_GRADIENT: 

387 dht += i_nh * K * V 

388 

389 stride_v = H * V 

390 stride_h = H * K * V 

391 stride_k = H * K 

392 

393 if USE_FINAL_STATE_GRADIENT: 

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

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

396 if K > 64: 

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

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

399 if K > 128: 

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

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

402 if K > 192: 

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

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

405 

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

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

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

409 if K > 64: 

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

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

412 if K > 128: 

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

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

415 if K > 192: 

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

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

418 

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

420 if USE_G: 

421 if IS_VARLEN: 

422 bos_g = i_h * T_all + bos 

423 else: 

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

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

426 bg_last_exp = tl.exp(bg_last) 

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

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

429 b_g_exp = tl.exp(b_g) 

430 

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

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

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

434 

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

436 

437 # Update dv 

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

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

440 if USE_GK: 

441 o_k1 = tl.arange(0, 64) 

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

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

444 

445 if K > 64: 

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

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

448 if USE_GK: 

449 o_k2 = 64 + o_k1 

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

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

452 

453 if K > 128: 

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

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

456 if USE_GK: 

457 o_k3 = 128 + o_k1 

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

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

460 

461 if K > 192: 

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

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

464 if USE_GK: 

465 o_k4 = 192 + o_k1 

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

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

468 

469 if USE_G: 

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

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

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

473 

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

475 # Update dh 

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

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

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

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

480 if USE_G: 

481 b_dh1 *= bg_last_exp 

482 b_q = b_q * b_g_exp[None, :] 

483 if USE_GK: 

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

485 b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))