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 if USE_G: 

100 g = g + i_h * T_all + bos 

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

102 if SAVE_NEW_VALUE: 

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

104 else: 

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

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

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

108 if USE_G: 

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

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

111 if SAVE_NEW_VALUE: 

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

113 

114 if USE_INITIAL_STATE: 

115 h0_ptr = h0 + i_nh * K * V 

116 if STORE_FINAL_STATE: 

117 ht_ptr = ht + i_nh * K * V 

118 

119 # Load initial state 

120 if USE_INITIAL_STATE: 

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

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

123 if K > 64: 

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

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

126 if K > 128: 

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

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

129 if K > 192: 

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

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

132 

133 # Main recurrence over chunks 

134 for i_t in range(NT): 

135 # Store current hidden state h_t 

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

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

138 if K > 64: 

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

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

141 if K > 128: 

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

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

144 if K > 192: 

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

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

147 

148 # Compute v_residual = v - w @ h 

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

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

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

152 if K > 64: 

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

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

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

156 if K > 128: 

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

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

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

160 if K > 192: 

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

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

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

164 

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

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

167 

168 if SAVE_NEW_VALUE: 

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

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

171 

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

173 

174 # Apply output gate g 

175 if USE_G: 

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

177 b_g_last = tl.load(g + last_idx) 

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

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

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

181 b_g_last_exp = tl.exp(b_g_last) 

182 b_h1 *= b_g_last_exp 

183 if K > 64: 

184 b_h2 *= b_g_last_exp 

185 if K > 128: 

186 b_h3 *= b_g_last_exp 

187 if K > 192: 

188 b_h4 *= b_g_last_exp 

189 

190 # Apply key gate gk 

191 if USE_GK: 

192 o_k1 = tl.arange(0, 64) 

193 if IS_VARLEN: 

194 gk_base_ptr = gk + (i_h * T_all + bos) * K 

195 else: 

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

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

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

199 if K > 64: 

200 o_k2 = 64 + o_k1 

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

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

203 if K > 128: 

204 o_k3 = 128 + o_k1 

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

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

207 if K > 192: 

208 o_k4 = 192 + o_k1 

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

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

211 

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

213 

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

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

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

217 if USE_GK: 

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

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

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

221 

222 if K > 64: 

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

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

225 if USE_GK: 

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

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

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

229 

230 if K > 128: 

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

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

233 if USE_GK: 

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

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

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

237 

238 if K > 192: 

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

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

241 if USE_GK: 

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

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

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

245 

246 # Store final state 

247 if STORE_FINAL_STATE: 

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

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

250 if K > 64: 

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

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

253 if K > 128: 

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

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

256 if K > 192: 

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

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

259 

260 

261def chunk_gated_delta_rule_fwd_h( 

262 k: torch.Tensor, 

263 w: torch.Tensor, 

264 u: torch.Tensor, 

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

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

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

268 output_final_state: bool = False, 

269 chunk_size: int = 64, # default:64 

270 save_new_value: bool = True, 

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

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

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

274 BT = chunk_size 

275 

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

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

278 if cu_seqlens is None: 

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

280 else: 

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

282 if K > 256: 

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

284 

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

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

287 

288 BV = 128 

289 

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

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

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

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

294 g = g.permute(0, 2, 1).contiguous() if g is not None else None 

295 # Match the head-major layout used by the forward kernel's token strides. 

296 gk = gk.permute(0, 2, 1, 3).contiguous() if gk is not None else None 

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

298 k=k, 

299 v=u, 

300 w=w, 

301 v_new=v_new, 

302 g=g, 

303 gk=gk, 

304 h=h, 

305 h0=initial_state, 

306 ht=final_state, 

307 cu_seqlens=cu_seqlens, 

308 chunk_offsets=chunk_offsets, 

309 T=T, 

310 H=H, 

311 K=K, 

312 V=V, 

313 BT=BT, 

314 BV=BV, 

315 NT=NT, 

316 ) 

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

318 v_new = v_new.permute(0, 2, 1, 3).contiguous() if v_new is not None else None 

319 return h, v_new, final_state 

320 

321 

322@triton.heuristics({ 

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

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

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

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

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

328}) 

329# Preserve all cache keys without interpreting them as automatic tiling axes. 

330@triton.runtime.autotuner.autotune( 

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

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

333) 

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

335def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( 

336 q, 

337 k, 

338 w, 

339 g, 

340 gk, 

341 dht, 

342 dh0, 

343 do, 

344 dh, 

345 dv, 

346 dv2, 

347 cu_seqlens, 

348 chunk_offsets, 

349 scale, 

350 T, 

351 H: tl.constexpr, 

352 K: tl.constexpr, 

353 V: tl.constexpr, 

354 BT: tl.constexpr, 

355 BV: tl.constexpr, 

356 USE_G: tl.constexpr, 

357 USE_GK: tl.constexpr, 

358 USE_INITIAL_STATE: tl.constexpr, 

359 USE_FINAL_STATE_GRADIENT: tl.constexpr, 

360 IS_VARLEN: tl.constexpr, 

361): 

362 T_all = T 

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

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

365 if IS_VARLEN: 

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

367 T = eos - bos 

368 NT = tl.cdiv(T, BT) 

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

370 else: 

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

372 NT = tl.cdiv(T, BT) 

373 boh = i_n * NT 

374 

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

376 if K > 64: 

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

378 if K > 128: 

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

380 if K > 192: 

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

382 

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

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

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

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

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

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

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

390 if USE_GK: 

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

392 

393 if USE_INITIAL_STATE: 

394 dh0 += i_nh * K * V 

395 if USE_FINAL_STATE_GRADIENT: 

396 dht += i_nh * K * V 

397 

398 stride_v = H * V 

399 stride_h = H * K * V 

400 stride_k = H * K 

401 

402 if USE_FINAL_STATE_GRADIENT: 

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

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

405 if K > 64: 

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

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

408 if K > 128: 

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

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

411 if K > 192: 

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

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

414 

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

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

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

418 if K > 64: 

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

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

421 if K > 128: 

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

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

424 if K > 192: 

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

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

427 

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

429 if USE_G: 

430 if IS_VARLEN: 

431 bos_g = i_h * T_all + bos 

432 else: 

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

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

435 bg_last_exp = tl.exp(bg_last) 

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

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

438 b_g_exp = tl.exp(b_g) 

439 

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

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

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

443 

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

445 

446 # Update dv 

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

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

449 if USE_GK: 

450 o_k1 = tl.arange(0, 64) 

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

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

453 

454 if K > 64: 

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

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

457 if USE_GK: 

458 o_k2 = 64 + o_k1 

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

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

461 

462 if K > 128: 

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

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

465 if USE_GK: 

466 o_k3 = 128 + o_k1 

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

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

469 

470 if K > 192: 

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

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

473 if USE_GK: 

474 o_k4 = 192 + o_k1 

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

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

477 

478 if USE_G: 

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

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

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

482 

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

484 # Update dh 

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

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