Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/platform/torch/custom_ops/gdn/triton/solve

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

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

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 -*-

20# pylint: disable=missing-public-type-hints,missing-public-docstring,invalid-name

21# pylint: disable=import-outside-toplevel,unused-argument,unused-import

22# pylint: disable=missing-module-docstring,missing-function-docstring

24import os

25from typing import Optional

27import torch

28import triton

29import triton.language as tl

31from .utils import prepare_chunk_indices, make_tensor_descriptor, input_guard

34def _ensure_slice_ops() -> bool:

35 """Probe and attach tl.extract_slice / insert_slice if missing; return success."""

36 if hasattr(tl, "extract_slice") and hasattr(tl, "insert_slice"):

37 return True

38 try:

39 from triton.language.extra.cann.extension import extract_slice, insert_slice

40 tl.extract_slice = extract_slice

41 tl.insert_slice = insert_slice

42 return True

43 except ImportError:

44 return False

46_TRITON_SLICE_AVAILABLE: bool = _ensure_slice_ops()

47FLA_TRIL_PRECISION = os.environ.get('FLA_TRIL_PRECISION', 'ieee')

50@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})

51@triton.jit(do_not_specialize=["T"])

52def solve_tril_16x16_loop_kernel_paral_v3(

53 A_ptr,

54 Ad_ptr,

55 cu_seqlens,

56 chunk_indices,

57 T,

58 H: tl.constexpr,

59 BT: tl.constexpr,

60 IS_VARLEN: tl.constexpr,

61 LARGE_BLOCK_T: tl.constexpr,

62 NT: tl.constexpr,

63 BH: tl.constexpr,

64):

65 worker_id = tl.program_id(0)

66 total_tasks = NT * BH

67 num_tasks = total_tasks // 48

68 remainder = total_tasks - num_tasks * 48

69 upper_bound = min(total_tasks, num_tasks * (worker_id + 1) + min(worker_id + 1, remainder))

70 lower_bound = num_tasks * worker_id + min(worker_id, remainder)

71 for task_id in range(lower_bound, upper_bound):

72 i_t = task_id // BH

73 i_bh = task_id % BH

74 i_b, i_h = i_bh // H, i_bh % H

75 if IS_VARLEN:

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

77 chunk_indices + i_t * 2 + 1

78 ).to(tl.int32)

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

80 cu_seqlens + i_n + 1

81 ).to(tl.int32)

82 T = eos - bos

83 else:

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

86 A = A_ptr + (bos * H + i_h) * BT

87 Ad = Ad_ptr + (bos * H + i_h) * 16

89 base_t = i_t * LARGE_BLOCK_T

91 NTASKS: tl.constexpr = 2

92 N_BLOCKS: tl.constexpr = LARGE_BLOCK_T // 16 // NTASKS

94 for taskid in range(0, NTASKS):

95 base_t += taskid * (LARGE_BLOCK_T // NTASKS)

97 b_A = tl.zeros((N_BLOCKS, 16, 16), dtype=tl.float32) # (N_BLOCKS, 16, 16)

98 for blkid in range(0, N_BLOCKS):

99 row_start_o = base_t + blkid * 16

100 col_start_o = row_start_o % BT

101 # using ptr with mask instead of tl.load(block_ptr)

102 offs_rows_in_block = tl.arange(0, 16)

103 offs_cols_in_block = tl.arange(0, 16)

104 ptr_A_subrec16 = (

105 A

106 + row_start_o * H * BT

107 + col_start_o

108 + offs_rows_in_block[:, None] * H * BT

109 + offs_cols_in_block[None, :]

110 )

111 global_rows = row_start_o + offs_rows_in_block[:, None]

112 global_cols = col_start_o + offs_cols_in_block[None, :]

113 load_mask = (global_rows < T) & (global_cols < BT)

114 b_A_subrec16 = tl.load(ptr_A_subrec16, mask=load_mask, other=0.0).to(

115 tl.float32

116 )

117 b_A = tl.insert_slice(

118 ful=b_A,

119 sub=b_A_subrec16[None, :, :], # (1, 16, 16)

120 offsets=[blkid, 0, 0],

121 sizes=[1, 16, 16],

122 strides=[1, 1, 1],

123 )

124

125 # load multi 16x16

126 local_ori_A = tl.trans(b_A, (1, 0, 2))

127 local_ori_A = tl.reshape(local_ori_A, (16, 16 * N_BLOCKS)) # (16, N_BLOCKS*16)

128

129 # change mask into matrix elementwise action

130 tmp = tl.arange(0, 16).to(tl.float32)

131 rows = tmp[:, None]

132 cols = tmp[None, :]

133 is_lower = (rows > cols).to(b_A.dtype)

134 b_A = -b_A * is_lower

135

136 for i in range(1, 16):

137 nblks_vec16 = -tl.extract_slice(

138 local_ori_A, (i, 0), (1, 16 * N_BLOCKS), (16 * N_BLOCKS, 1)

139 )

140 b_a = tl.reshape(nblks_vec16, (N_BLOCKS, 16))

141

142 dot_tmp = tl.trans(b_a[:, :, None] * b_A, (1, 0, 2))

143 dot_product = tl.sum(dot_tmp, 0)

144 b_a = b_a + dot_product # (N_BLOCKS, 16)

145

146 b_a_new_expanded = b_a[:, None, :] # (N_BLOCKS, 1, 16)

147 b_A = tl.insert_slice(

148 ful=b_A,

149 sub=b_a_new_expanded,

150 offsets=[0, i, 0],

151 sizes=[N_BLOCKS, 1, 16],

152 strides=[1, 1, 1],

153 )

154

155 on_diagonal = rows == cols

156 b_A = tl.where(on_diagonal, b_A + 1.0, b_A)

157

158 b_A = tl.reshape(b_A, (N_BLOCKS * 16, 16))

159 # using ptr with mask instead of tl.load(block_ptr)

160 offs_rows_to_store = tl.arange(0, N_BLOCKS * 16)

161 offs_cols_to_store = tl.arange(0, 16)

162 p_Ai = (

163 Ad

164 + base_t * H * 16

165 + 0

166 + offs_rows_to_store[:, None] * H * 16

167 + offs_cols_to_store[None, :]

168 )

169 global_store_rows = base_t + offs_rows_to_store[:, None]

170 store_mask = global_store_rows < T

171 tl.store(

172 p_Ai,

173 b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"),

174 mask=store_mask,

175 )

176

177

178@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})

179@triton.jit(do_not_specialize=["T", "NT"])

180def merge_16x16_to_32x32_loop_inverse_kernel(

181 A,

182 Ad,

183 Ai,

184 cu_seqlens,

185 chunk_indices,

186 T,

187 NT,

188 H: tl.constexpr,

189 BT: tl.constexpr,

190 IS_VARLEN: tl.constexpr,

191 BH: tl.constexpr,

192):

193 worker_id = tl.program_id(0)

194 total_tasks = NT * BH

195 num_tasks = total_tasks // 24

196 remainder = total_tasks - num_tasks * 24

197 upper_bound = min(total_tasks, num_tasks * (worker_id + 1) + min(worker_id + 1, remainder))

198 lower_bound = num_tasks * worker_id + min(worker_id, remainder)

199 for task_id in range(lower_bound, upper_bound):

200 i_tt = task_id // BH

201 i_bh = task_id % BH

202 i_b, i_h = i_bh // H, i_bh % H

203 if IS_VARLEN:

204 i_n, i_t = tl.load(chunk_indices + i_tt * 2).to(tl.int32), tl.load(

205 chunk_indices + i_tt * 2 + 1

206 ).to(tl.int32)

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

208 cu_seqlens + i_n + 1

209 ).to(tl.int32)

210 T = eos - bos

211 else:

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

213 i_t = i_tt

214

215 A_ptr = A + (bos * H + i_h) * BT

216 Ad_ptr = Ad + (bos * H + i_h) * 16

217 Ai_ptr = Ai + (bos * H + i_h) * 32

218

219 p_A_21 = tl.make_block_ptr(

220 A_ptr, (T, BT), (H * BT, 1), (i_t * 32 + 16, 0 + i_t % (BT // 32) * 32), (16, 16), (1, 0)

221 )

222 p_Ad_11 = tl.make_block_ptr(

223 Ad_ptr, (T, 16), (H * 16, 1), (i_t * 32, 0), (16, 16), (1, 0)

224 )

225 p_Ad_22 = tl.make_block_ptr(

226 Ad_ptr, (T, 16), (H * 16, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0)

227 )

228 p_Ai_11 = tl.make_block_ptr(

229 Ai_ptr, (T, 32), (H * 32, 1), (i_t * 32, 0), (16, 16), (1, 0)

230 )

231 p_Ai_22 = tl.make_block_ptr(

232 Ai_ptr, (T, 32), (H * 32, 1), (i_t * 32 + 16, 16), (16, 16), (1, 0)

233 )

234 p_Ai_21 = tl.make_block_ptr(

235 Ai_ptr, (T, 32), (H * 32, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0)

236 )

237

238 A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)

239 Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32)

240 Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32)

241 Ai_21 = -tl.dot(

242 tl.dot(Ai_22, A_21, input_precision="ieee"), Ai_11, input_precision="ieee"

243 )

244 tl.store(

245 p_Ai_11,

246 Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"),

247 boundary_check=(0, 1),

248 )

249 tl.store(

250 p_Ai_22,

251 Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"),

252 boundary_check=(0, 1),

253 )

254 tl.store(

255 p_Ai_21,

256 Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"),

257 boundary_check=(0, 1),

258 )

259

260

261@triton.heuristics(

262 {

263 "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,

264 }

265)

266@triton.jit(do_not_specialize=["T", "NT"])

267def merge_32x32_to_64x64_loop_inverse_kernel(

268 A,

269 Ad,

270 Ai,

271 cu_seqlens,

272 chunk_indices,

273 T,

274 NT,

275 H: tl.constexpr,

276 BT: tl.constexpr,

277 IS_VARLEN: tl.constexpr,

278 BH: tl.constexpr,

279):

280 worker_id = tl.program_id(0)

281 total_tasks = NT * BH

282 num_tasks = total_tasks // 24

283 remainder = total_tasks - num_tasks * 24

284 upper_bound = min(total_tasks, num_tasks * (worker_id + 1) + min(worker_id + 1, remainder))

285 lower_bound = num_tasks * worker_id + min(worker_id, remainder)

286 for task_id in range(lower_bound, upper_bound):

287 i_tt = task_id // BH

288 i_bh = task_id % BH

289 i_b, i_h = i_bh // H, i_bh % H

290 if IS_VARLEN:

291 i_n, i_t = tl.load(chunk_indices + i_tt * 2).to(tl.int32), tl.load(

292 chunk_indices + i_tt * 2 + 1

293 ).to(tl.int32)

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

295 cu_seqlens + i_n + 1

296 ).to(tl.int32)

297 T = eos - bos

298 else:

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

300 i_t = i_tt

301

302 A_ptr = A + (bos * H + i_h) * BT

303 Ad_ptr = Ad + (bos * H + i_h) * 32

304 Ai_ptr = Ai + (bos * H + i_h) * 64

305

306 p_A_21 = tl.make_block_ptr(

307 A_ptr, (T, BT), (H * BT, 1), (i_t * 64 + 32, 0 + i_t % (BT // 64) * 64), (32, 32), (1, 0)

308 )

309

310 p_Ad_11 = tl.make_block_ptr(

311 Ad_ptr, (T, 32), (H * 32, 1), (i_t * 64, 0), (32, 32), (1, 0)

312 )

313 p_Ad_22 = tl.make_block_ptr(

314 Ad_ptr, (T, 32), (H * 32, 1), (i_t * 64 + 32, 0), (32, 32), (1, 0)

315 )

316

317 p_Ai_11 = tl.make_block_ptr(

318 Ai_ptr, (T, 64), (H * 64, 1), (i_t * 64, 0), (32, 32), (1, 0)

319 )

320 p_Ai_22 = tl.make_block_ptr(

321 Ai_ptr, (T, 64), (H * 64, 1), (i_t * 64 + 32, 32), (32, 32), (1, 0)

322 )

323 p_Ai_21 = tl.make_block_ptr(

324 Ai_ptr, (T, 64), (H * 64, 1), (i_t * 64 + 32, 0), (32, 32), (1, 0)

325 )

326

327 A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)

328 Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32)

329 Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32)

330 Ai_21 = -tl.dot(

331 tl.dot(Ai_22, A_21, input_precision="ieee"), Ai_11, input_precision="ieee"

332 )

333 tl.store(

334 p_Ai_11,

335 Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"),

336 boundary_check=(0, 1),

337 )

338 tl.store(

339 p_Ai_22,

340 Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"),

341 boundary_check=(0, 1),

342 )

343 tl.store(

344 p_Ai_21,

345 Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"),

346 boundary_check=(0, 1),

347 )

348

349

350@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})

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

352def solve_tril_64x64_kernel(

353 A,

354 Ai,

355 cu_seqlens,

356 chunk_indices,

357 T,

358 H: tl.constexpr,

359 BT: tl.constexpr,

360 USE_TMA: tl.constexpr,

361 IS_VARLEN: tl.constexpr,

362 DOT_PRECISION: tl.constexpr

363):

364 i_t, i_bh = tl.program_id(0), tl.program_id(1)

365 i_b, i_h = i_bh // H, i_bh % H

366 if IS_VARLEN:

367 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)

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

369 T = eos - bos

370 else:

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

372 o_i = tl.arange(0, 64)

373 m_I = o_i[:, None] == o_i[None, :]

374

375 A = A + (bos * H + i_h) * BT

376 Ai = Ai + (bos * H + i_h) * 64

377

378 offset = (i_t * 64) % BT

379 if not USE_TMA:

380 p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * 64, offset), (64, 64), (1, 0))

381 b_A = -tl.load(p_A, boundary_check=(0, 1)).to(tl.float32)

382 else:

383 desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [64, 64])

384 desc_o = make_tensor_descriptor(Ai, [T, 64], [H * 64, 1], [64, 64])

385 b_A = -desc.load([i_t * 64, offset]).to(tl.float32)

386

387 for i in range(2, min(64, T - i_t * 64)):

388 b_a = -tl.load(A + (i_t * 64 + i) * H * BT + o_i + offset)

389 b_a = b_a + tl.sum(b_a[:, None] * b_A, 0)

390 b_A = tl.where((o_i == i)[:, None], b_a, b_A)

391 b_A += m_I

392 if not USE_TMA:

393 p_Ai = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64, 0), (64, 64), (1, 0))

394 tl.store(p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))

395 else:

396 desc_o.store([i_t * 64, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne"))

397

398

399def solve_tril_64(

400 A: torch.Tensor,

401 cu_seqlens: Optional[torch.Tensor] = None,

402 output_dtype: torch.dtype = torch.float,

403 ):

404 B, T, H, BT = A.shape

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

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

407

408 Ai = torch.zeros_like(A, dtype=output_dtype)

409 solve_tril_64x64_kernel[NT, B * H](

410 A=A,

411 Ai=Ai,

412 cu_seqlens=cu_seqlens,

413 chunk_indices=chunk_indices,

414 T=T,

415 H=H,

416 BT=BT,

417 USE_TMA=False,

418 DOT_PRECISION=FLA_TRIL_PRECISION,

419 )

420 return Ai

421

422

423@input_guard

424def solve_tril(

425 A: torch.Tensor,

426 cu_seqlens: Optional[torch.Tensor] = None,

427 output_dtype: torch.dtype = torch.float

428) -> torch.Tensor:

429 """

430 Compute the inverse of the matrix I + A

431 A should be strictly lower triangular, i.e., A.triu() == 0.

432

433 Args:

434 A (torch.Tensor):

435 [B, T, H, BT], where BT should only be 16, 32, or 64.

436 cu_seqlens (torch.Tensor):

437 The cumulative sequence lengths of the input tensor. Default: `None`.

438 output_dtype (torch.dtype):

439 The dtype of the output tensor. Default: `torch.float`.

440 If `None`, the output dtype will be the same as the input dtype.

441

442 Returns:

443 (I + A)^-1 with the same shape as A

444 """

445 output_dtype = A.dtype if output_dtype is None else output_dtype

446 if not _TRITON_SLICE_AVAILABLE:

447 if A.shape[-1] not in [64]:

448 raise ValueError(

449 f"A shape BT should in [64], but current is {A.shape[-1]}"

450 )

451 return solve_tril_64(A, cu_seqlens, output_dtype)

452 if A.shape[-1] not in [16, 32, 64]:

453 raise ValueError(

454 f"A shape BT should in [16, 32, 64], but current is {A.shape[-1]}"

455 )

456

457 B, T, H, BT = A.shape

458 # If BT matches the current processing level (final step), use output_dtype

459 # (e.g. BF16) so the kernel can downcast internally, avoiding an extra

460 # external cast that hurts performance. Otherwise, keep FP32 to preserve

461 # precision for subsequent computation stages.

462 Ad = torch.empty(

463 B, T, H, 16, device=A.device, dtype=torch.float if BT != 16 else output_dtype

464 )

465

466 LARGE_BLOCK_T = 608 * 2

467

468 chunk_indices = (

469 prepare_chunk_indices(cu_seqlens, LARGE_BLOCK_T)

470 if cu_seqlens is not None

471 else None

472 )

473 NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, LARGE_BLOCK_T)

474 solve_tril_16x16_loop_kernel_paral_v3[(48,)](

475 A,

476 Ad,

477 cu_seqlens=cu_seqlens,

478 chunk_indices=chunk_indices,

479 T=T,

480 H=H,

481 BT=BT,

482 LARGE_BLOCK_T=LARGE_BLOCK_T,

483 NT=NT,

484 BH=B * H,

485 )

486

487 if BT == 16:

488 return Ad

489

490 # Same dtype logic as above: output_dtype for the final step, FP32 otherwise.

491 Ai = torch.zeros(

492 B, T, H, 32, device=A.device, dtype=torch.float if BT != 32 else output_dtype

493 )

494

495 chunk_indices = (

496 prepare_chunk_indices(cu_seqlens, 32) if cu_seqlens is not None else None

497 )

498 NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, 32)

499 merge_16x16_to_32x32_loop_inverse_kernel[(24,)](

500 A=A,

501 Ad=Ad,

502 Ai=Ai,

503 cu_seqlens=cu_seqlens,

504 chunk_indices=chunk_indices,

505 T=T,

506 H=H,

507 BT=BT,

508 NT=NT,

509 BH=B * H,

510 )

511 if BT == 32:

512 return Ai

513

514 Ad = Ai

515 # Same dtype logic as above: output_dtype for the final step, FP32 otherwise.

516 Ai = torch.zeros(

517 B, T, H, 64, device=A.device, dtype=torch.float if BT != 64 else output_dtype

518 )

519 chunk_indices = (

520 prepare_chunk_indices(cu_seqlens, 64) if cu_seqlens is not None else None

521 )

522 NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, 64)

523 merge_32x32_to_64x64_loop_inverse_kernel[(24,)](

524 A=A,

525 Ad=Ad,

526 Ai=Ai,

527 cu_seqlens=cu_seqlens,

528 chunk_indices=chunk_indices,

529 T=T,

530 H=H,

531 BT=BT,

532 NT=NT,

533 BH=B * H,

534 )

535 if BT == 64:

536 return Ai

537 return Ai

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / custom_ops / gdn / triton / solve_tril.py: 0%

217 statements