Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / custom_ops / gdn / triton / chunk_o.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=missing-public-type-hints,missing-public-docstring,disallowed-name 

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

20# pylint: disable=unused-variable,too-many-nested-blocks 

21 

22from typing import Optional, Tuple 

23 

24import torch 

25import triton 

26import triton.language as tl 

27 

28from .utils import prepare_chunk_indices, exp, prepare_chunk_offsets 

29 

30 

31@triton.heuristics({ 

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

33 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, 

34 'USE_DW': lambda args: args['dw'] is not None, 

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

36}) 

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

38def chunk_bwd_kernel_dqkwg( 

39 q, 

40 k, 

41 v, 

42 h, 

43 g, 

44 g_gamma, 

45 do, 

46 dh, 

47 dq, 

48 dk, 

49 dg, 

50 w, 

51 dv, 

52 dw, 

53 cu_seqlens, 

54 chunk_indices, 

55 scale, 

56 B: tl.constexpr, 

57 T, 

58 H: tl.constexpr, 

59 K: tl.constexpr, 

60 V: tl.constexpr, 

61 BT: tl.constexpr, 

62 BK: tl.constexpr, 

63 BV: tl.constexpr, 

64 USE_G: tl.constexpr, 

65 USE_G_GAMMA: tl.constexpr, 

66 USE_DW: tl.constexpr, 

67 IS_VARLEN: tl.constexpr, 

68 gdiff, 

69): 

70 i_t, i_b = tl.program_id(0), tl.program_id(1) 

71 T_max = T 

72 if IS_VARLEN: 

73 i_tg = i_t 

74 i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) 

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

76 total = B * T_max 

77 T = eos - bos 

78 else: 

79 NT = tl.cdiv(T, BT) 

80 i_tg = i_b * NT + i_t 

81 bos, eos = i_b * T, i_b * T + T 

82 total = B * T_max 

83 

84 NK = tl.cdiv(K, BK) 

85 for i_k in range(NK): 

86 if USE_G: 

87 dg_k = dg + i_k * total * H 

88 

89 for i_h in range(H): 

90 v_h = v + (bos * H + i_h) * V 

91 do_h = do + (bos * H + i_h) * V 

92 h_h = h + (i_tg * H + i_h).to(tl.int64) * K * V 

93 dh_h = dh + (i_tg * H + i_h).to(tl.int64) * K * V 

94 q_h = q + (bos * H + i_h) * K 

95 k_h = k + (bos * H + i_h) * K 

96 dq_h = dq + (bos * H + i_h) * K 

97 dk_h = dk + (bos * H + i_h) * K 

98 

99 if USE_DW: 

100 w_h = w + (bos * H + i_h) * K 

101 dw_h = dw + (bos * H + i_h) * K 

102 dv_h = dv + (bos * H + i_h) * V 

103 

104 if USE_G: 

105 if IS_VARLEN: 

106 dg_h = dg_k + i_h * T_max + bos 

107 g_h = g + i_h * T_max + bos 

108 else: 

109 dg_h = dg_k + (i_b * H + i_h) * T_max 

110 g_h = g + (i_b * H + i_h) * T_max 

111 b_dg_last = tl.zeros([1, ], dtype=tl.float32) 

112 

113 if USE_G_GAMMA: 

114 b_gamma = tl.load(g_gamma + i_h) 

115 b_g = b_gamma * (tl.arange(0, BT) + 1) 

116 b_g_last = b_gamma * min(BT, T - i_t * BT) 

117 

118 b_dq = tl.zeros([BT, BK], dtype=tl.float32) 

119 b_dk = tl.zeros([BT, BK], dtype=tl.float32) 

120 b_ds = tl.zeros([BT, BT], dtype=tl.float32) 

121 b_dw = tl.zeros([BT, BK], dtype=tl.float32) if USE_DW else None 

122 

123 for i_v in range(tl.cdiv(V, BV)): 

124 p_v = tl.make_block_ptr(v_h, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) 

125 p_do = tl.make_block_ptr(do_h, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) 

126 p_h = tl.make_block_ptr(h_h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) 

127 p_dh = tl.make_block_ptr(dh_h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) 

128 

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

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

131 b_h = tl.load(p_h, boundary_check=(0, 1)) 

132 b_dh = tl.load(p_dh, boundary_check=(0, 1)) 

133 

134 if USE_G: 

135 b_dg_last += (tl.sum(b_h * b_dh)) 

136 

137 b_ds += tl.dot(b_do, tl.trans(b_v)) 

138 b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) 

139 b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) 

140 

141 if USE_DW: 

142 p_dv = tl.make_block_ptr(dv_h, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) 

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

144 b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype)) 

145 

146 if USE_DW: 

147 p_dw = tl.make_block_ptr(dw_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) 

148 tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) 

149 

150 tl.debug_barrier() 

151 

152 p_q = tl.make_block_ptr(q_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) 

153 p_k = tl.make_block_ptr(k_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) 

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

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

156 

157 p_dq = tl.make_block_ptr(dq_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) 

158 p_dk = tl.make_block_ptr(dk_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) 

159 

160 o_t = i_t * BT + tl.arange(0, BT) 

161 m_t = o_t < T 

162 m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) 

163 

164 if USE_G: 

165 b_dg = tl.zeros([BT, ], dtype=tl.float32) 

166 p_g = tl.make_block_ptr(g_h, (T,), (1,), (i_t * BT,), (BT,), (0,)) 

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

168 b_g_last = tl.load(g_h + (min(i_t * BT + BT, T) - 1) * 1) 

169 b_dg_last *= tl.exp(b_g_last) 

170 

171 b_dq = b_dq * tl.exp(b_g)[:, None] * scale 

172 b_dg += tl.sum(b_dq * b_q, axis=1) 

173 

174 b_dk = b_dk * tl.where(m_t, tl.exp(-b_g + b_g_last), 0)[:, None] 

175 b_dg -= tl.sum(b_k * b_dk, axis=1) 

176 b_dg_last += tl.sum(b_dk * b_k) 

177 

178 if IS_VARLEN: 

179 b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale 

180 else: 

181 p_gdiff = tl.make_block_ptr(gdiff + i_b * H * NT * BT * BT + i_h * NT * BT * BT + i_t * BT * BT, 

182 (BT, BT), (BT, 1), (0, 0), (BT, BT), (1, 0)) 

183 gdiff_ = tl.load(p_gdiff) 

184 b_ds = b_ds * gdiff_ * scale 

185 

186 b_ds2 = b_ds * tl.dot(b_q, tl.trans(b_k)) 

187 b_dg += tl.sum(b_ds2, axis=1) 

188 b_dg -= tl.sum(b_ds2, axis=0) 

189 

190 b_ds = b_ds.to(b_k.dtype) 

191 b_dq += tl.dot(b_ds, b_k) 

192 b_dk += tl.dot(tl.trans(b_ds), b_q) 

193 p_dg = tl.make_block_ptr(dg_h, (T,), (1,), (i_t * BT,), (BT,), (0,)) 

194 

195 last_index_local = min(BT, T - i_t * BT) - 1 

196 if last_index_local >= 0: 

197 is_last_mask = tl.arange(0, BT) == last_index_local 

198 b_dg = tl.where(is_last_mask, b_dg + b_dg_last, b_dg) 

199 else: 

200 pass 

201 

202 tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) 

203 tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) 

204 tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) 

205 

206 elif USE_G_GAMMA: 

207 b_dq = b_dq * exp(b_g)[:, None] * scale 

208 b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] 

209 b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale 

210 b_ds = b_ds.to(b_k.dtype) 

211 b_dq += tl.dot(b_ds, b_k) 

212 b_dk += tl.dot(tl.trans(b_ds), b_q) 

213 tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) 

214 tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) 

215 

216 else: 

217 b_ds = tl.where(m_A, b_ds, 0) 

218 b_ds = b_ds.to(b_k.dtype) 

219 b_dq += tl.dot(b_ds, b_k) 

220 b_dk += tl.dot(tl.trans(b_ds), b_q) * scale 

221 b_dq *= scale 

222 tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) 

223 tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) 

224 

225 

226@triton.heuristics({ 

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

228 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, 

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

230}) 

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

232def chunk_bwd_kernel_dv_local( 

233 q, 

234 k, 

235 g, 

236 g_gamma, 

237 do, 

238 dv, 

239 cu_seqlens, 

240 chunk_indices, 

241 scale, 

242 T, 

243 H: tl.constexpr, 

244 K: tl.constexpr, 

245 V: tl.constexpr, 

246 BT: tl.constexpr, 

247 BK: tl.constexpr, 

248 BV: tl.constexpr, 

249 USE_G: tl.constexpr, 

250 USE_G_GAMMA: tl.constexpr, 

251 IS_VARLEN: tl.constexpr, 

252): 

253 i_t, i_b = tl.program_id(0), tl.program_id(1) 

254 T_max = T 

255 

256 if IS_VARLEN: 

257 i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) 

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

259 T = eos - bos 

260 else: 

261 bos, eos = i_b * T, i_b * T + T 

262 

263 for i_h in range(H): 

264 offset_kh = (bos * H + i_h) * K 

265 offset_vh = (bos * H + i_h) * V 

266 

267 b_A = tl.zeros([BT, BT], dtype=tl.float32) 

268 for i_k in range(tl.cdiv(K, BK)): 

269 p_k = tl.make_block_ptr(k + offset_kh, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) 

270 p_q = tl.make_block_ptr(q + offset_kh, (K, T), (1, H * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) 

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

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

273 b_A += tl.dot(b_k, b_q) 

274 

275 if USE_G: 

276 if IS_VARLEN: 

277 offset_g = i_h * T_max + bos 

278 else: 

279 offset_g = i_b * H * T_max + i_h * T_max 

280 

281 p_g = tl.make_block_ptr(g + offset_g, (T,), (1,), (i_t * BT,), (BT,), (0,)) 

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

283 

284 if USE_G_GAMMA: 

285 b_gamma = tl.load(g_gamma + i_h) 

286 b_g = b_gamma * (tl.arange(0, BT) + 1) 

287 

288 o_t = i_t * BT + tl.arange(0, BT) 

289 m_t = o_t < T 

290 m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) 

291 

292 if USE_G: 

293 b_A = tl.where(m_A, b_A * tl.exp(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty) 

294 else: 

295 b_A = tl.where(m_A, b_A * scale, 0).to(do.dtype.element_ty) 

296 

297 for i_v in range(tl.cdiv(V, BV)): 

298 p_do = tl.make_block_ptr(do + offset_vh, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) 

299 p_dv = tl.make_block_ptr(dv + offset_vh, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) 

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

301 b_dv = tl.dot(b_A.to(b_do.dtype), b_do) 

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

303 

304 

305@triton.heuristics({ 

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

307 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, 

308 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None 

309}) 

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

311def chunk_fwd_kernel_o( 

312 q, 

313 k, 

314 v, 

315 h, 

316 g, 

317 g_gamma, 

318 o, 

319 cu_seqlens, 

320 chunk_offsets, 

321 scale, 

322 T, 

323 H: tl.constexpr, 

324 N: tl.constexpr, 

325 Hg: tl.constexpr, 

326 K: tl.constexpr, 

327 V: tl.constexpr, 

328 BT: tl.constexpr, 

329 BK: tl.constexpr, 

330 BV: tl.constexpr, 

331 USE_G: tl.constexpr, 

332 USE_G_GAMMA: tl.constexpr, 

333 IS_VARLEN: tl.constexpr, 

334): 

335 T_max = T 

336 for i_v in range(tl.cdiv(V, BV)): 

337 for i_n in range(N): 

338 if IS_VARLEN: 

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

340 cu_seqlens + i_n + 1 

341 ).to(tl.int32) 

342 T = eos - bos 

343 NT = tl.cdiv(T, BT) 

344 boh = tl.load(chunk_offsets + i_n).to(tl.int64) 

345 else: 

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

347 NT = tl.cdiv(T, BT) 

348 boh = i_n * NT 

349 

350 core_id = tl.program_id(0) 

351 total_cores = tl.num_programs(0) 

352 base_chunks_per_pid = NT // total_cores 

353 remainder = NT % total_cores 

354 

355 if core_id < remainder: 

356 chunks_this_pid = base_chunks_per_pid + 1 

357 start_idx = core_id * chunks_this_pid 

358 else: 

359 chunks_this_pid = base_chunks_per_pid 

360 start_idx = core_id * base_chunks_per_pid + remainder 

361 

362 # offset calculation 

363 for i_h in range(0, H): 

364 q_offset = (bos * Hg + i_h // (H // Hg)) * K 

365 k_offset = (bos * Hg + i_h // (H // Hg)) * K 

366 v_offset = (bos * H + i_h) * V 

367 o_offset = (bos * H + i_h) * V 

368 

369 for i_t in range(start_idx, start_idx + chunks_this_pid): 

370 i_tg = boh + i_t 

371 h_base = h + (i_tg * H + i_h).to(tl.int64) * K * V 

372 b_o = tl.zeros([BT, BV], dtype=tl.float32) 

373 b_A = tl.zeros([BT, BT], dtype=tl.float32) 

374 for i_k in range(tl.cdiv(K, BK)): 

375 p_q = tl.make_block_ptr( 

376 q + q_offset, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) 

377 ) 

378 p_k = tl.make_block_ptr( 

379 k + k_offset, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) 

380 ) 

381 p_h = tl.make_block_ptr( 

382 h_base, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0) 

383 ) 

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

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

386 b_h = tl.load(p_h, boundary_check=(0, 1)) 

387 

388 # [BT, BK] @ [BK, BV] -> [BT, BV] 

389 b_o += tl.dot(b_q, b_h) 

390 # [BT, BK] @ [BK, BT] -> [BT, BT] 

391 b_A += tl.dot(b_q, b_k) 

392 

393 if USE_G: 

394 if IS_VARLEN: 

395 p_g = tl.make_block_ptr(g + bos + i_h * T_max, (T,), (1,), (i_t * BT,), (BT,), (0,)) 

396 else: 

397 p_g = tl.make_block_ptr(g + bos * H + i_h * T_max, (T,), (1,), (i_t * BT,), (BT,), (0,)) 

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

399 b_o = b_o * exp(b_g)[:, None] 

400 b_A = b_A * exp(b_g[:, None] - b_g[None, :]) 

401 if USE_G_GAMMA: 

402 b_gamma = tl.load(g_gamma + i_h) 

403 b_g = b_gamma * (tl.arange(0, BT) + 1) 

404 

405 o_i = tl.arange(0, BT) 

406 m_A = o_i[:, None] >= o_i[None, :] 

407 b_A = tl.where(m_A, b_A, 0) 

408 

409 p_v = tl.make_block_ptr( 

410 v + v_offset, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) 

411 ) 

412 p_o = tl.make_block_ptr( 

413 o + o_offset, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) 

414 ) 

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

416 

417 # to fix mma -> mma layout conversion 

418 # already solved by triton v3.2 or higher 

419 b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale 

420 tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) 

421 

422 

423def chunk_bwd_dqkwg( 

424 q: torch.Tensor, 

425 k: torch.Tensor, 

426 v: torch.Tensor, 

427 do: torch.Tensor, 

428 h: torch.Tensor, 

429 dh: torch.Tensor, 

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

431 g_gamma: Optional[torch.Tensor] = None, 

432 dv: Optional[torch.Tensor] = None, 

433 w: Optional[torch.Tensor] = None, 

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

435 chunk_size: int = 64, 

436 scale: float = 1.0, 

437) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 

438 B, T, H, K, V = *k.shape, v.shape[-1] 

439 BT = min(chunk_size, max(16, triton.next_power_of_2(T))) 

440 chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None 

441 NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) 

442 

443 BK = 128 if cu_seqlens is None else 64 

444 BV = 64 

445 NK = triton.cdiv(K, BK) 

446 dq = torch.empty_like(q) 

447 dk = torch.empty_like(k) 

448 g = g.transpose(1, 2).contiguous() 

449 dg = torch.empty(NK, *g.shape, dtype=torch.float32, device=g.device) if g is not None else None 

450 dw = torch.empty_like(w) if w is not None else None 

451 grid = (NT, B) 

452 

453 if cu_seqlens is None: 

454 if NT * BT == T: 

455 g_ = g.reshape(B, H, NT, BT) 

456 g_diff = g_[:, :, :, :, None] - g_[:, :, :, None, :] 

457 g_diff = g_diff.clamp(-60, 60).exp() 

458 g_diff[:, :, :] *= torch.tril(torch.ones(BT, BT), diagonal=0).to(g.device) 

459 else: 

460 diff = NT * BT - T 

461 g_ = torch.cat((g, torch.zeros(B, H, diff).to(g.device)), dim=-1).reshape(B, H, NT, BT) 

462 g_diff = g_[:, :, :, :, None] - g_[:, :, :, None, :] 

463 g_diff = g_diff.clamp(-60, 60).exp() 

464 g_diff[:, :, :] *= torch.tril(torch.ones(BT, BT), diagonal=0).to(g.device) 

465 bias = torch.arange(0, BT).to(g.device) 

466 o_t = (NT - 1) * BT + bias 

467 m_t = o_t < T 

468 m_A = (m_t[:, None] & m_t) 

469 g_diff[:, :, -1] *= m_A 

470 else: 

471 g_diff = None 

472 

473 chunk_bwd_kernel_dqkwg[grid]( 

474 q=q, 

475 k=k, 

476 v=v, 

477 h=h, 

478 g=g, 

479 g_gamma=g_gamma, 

480 do=do, 

481 dh=dh, 

482 dv=dv, 

483 w=w, 

484 dw=dw, 

485 dq=dq, 

486 dk=dk, 

487 dg=dg, 

488 cu_seqlens=cu_seqlens, 

489 chunk_indices=chunk_indices, 

490 scale=scale, 

491 B=B, 

492 T=T, 

493 H=H, 

494 K=K, 

495 V=V, 

496 BT=BT, 

497 BK=BK, 

498 BV=BV, 

499 gdiff=g_diff, 

500 ) 

501 

502 if dg is not None: 

503 dg = dg.sum(0) 

504 dg = dg.transpose(1, 2).contiguous() 

505 return dq, dk, dw, dg 

506 

507 

508def chunk_bwd_dv_local( 

509 q: torch.Tensor, 

510 k: torch.Tensor, 

511 do: torch.Tensor, 

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

513 g_gamma: Optional[torch.Tensor] = None, 

514 scale: float = None, 

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

516 chunk_size: int = 64 

517) -> torch.Tensor: 

518 B, T, H, K, V = *k.shape, do.shape[-1] 

519 BT = min(chunk_size, max(16, triton.next_power_of_2(T))) 

520 chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None 

521 

522 BK = 128 

523 BV = 128 

524 NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) 

525 

526 g = g.transpose(1, 2).contiguous() 

527 dv = torch.empty_like(do) 

528 grid = (NT, B) 

529 chunk_bwd_kernel_dv_local[grid]( 

530 q=q, 

531 k=k, 

532 g=g, 

533 g_gamma=g_gamma, 

534 do=do, 

535 dv=dv, 

536 cu_seqlens=cu_seqlens, 

537 chunk_indices=chunk_indices, 

538 scale=scale, 

539 T=T, 

540 H=H, 

541 K=K, 

542 V=V, 

543 BT=BT, 

544 BK=BK, 

545 BV=BV, 

546 ) 

547 return dv 

548 

549 

550def chunk_fwd_o( 

551 q: torch.Tensor, 

552 k: torch.Tensor, 

553 v: torch.Tensor, 

554 h: torch.Tensor, 

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

556 g_gamma: Optional[torch.Tensor] = None,