Coverage for src/flag_gems/runtime/backend/_cambricon/ops/softmax.py: 0%

1import copy 

2import logging 

3import math 

4 

5import torch 

6import triton 

7import triton.language as tl 

8 

9from flag_gems import runtime 

10from flag_gems.runtime import torch_device_fn 

11from flag_gems.utils import libentry, libtuner 

12 

13from ..utils import MAX_NRAM_SIZE, TOTAL_CORE_NUM 

14from .zeros import zero_ 

15 

16logger = logging.getLogger("flag_gems").getChild(__name__.lstrip(".")) 

17MAX_N = 16384 

18 

19 

20def align(max_block): 

21 a = triton.next_power_of_2(max_block) 

22 return max_block if max_block == a else a // 2 

23 

24 

25def config_prune1(configs, named_args, **kwargs): 

26 M = named_args["M"] 

27 N = named_args["N"] 

28 K = named_args["K"] 

29 input = named_args["input_ptr"] 

30 configs_map = {} 

31 for config in configs: 

32 kw = config.kwargs 

33 TILE_K, TILE_N, num_warps, num_stages = ( 

34 kw["TILE_K"], 

35 kw["TILE_N"], 

36 config.num_warps, 

37 config.num_stages, 

38 ) 

39 if N < MAX_N: 

40 config = copy.deepcopy(config) 

41 TILE_N = config.kwargs["TILE_N"] = N 

42 k_per_core = math.ceil(K / max(TOTAL_CORE_NUM // M, 1)) 

43 TILE_K = config.kwargs["TILE_K"] = k_per_core 

44 num_stages = config.num_stages = 1 

45 key = (TILE_K, TILE_N, num_warps, num_stages) 

46 configs_map.setdefault(key, config) 

47 

48 config = copy.deepcopy(config) 

49 max_tile_k_without_pipe = MAX_NRAM_SIZE // 4 // (2 * TILE_N + 1) 

50 TILE_K = config.kwargs["TILE_K"] = align(max_tile_k_without_pipe) 

51 num_stages = config.num_stages = 1 

52 key = (TILE_K, TILE_N, num_warps, num_stages) 

53 configs_map.setdefault(key, config) 

54 

55 config = copy.deepcopy(config) 

56 max_tile_k_without_pipe = MAX_NRAM_SIZE // 4 // (3 * TILE_N + 1) 

57 if input.dtype == torch.float32: 

58 max_tile_k_without_pipe = MAX_NRAM_SIZE // 4 // (4 * TILE_N + 1) 

59 TILE_K = config.kwargs["TILE_K"] = align(max_tile_k_without_pipe) 

60 num_stages = config.num_stages = 3 

61 key = (TILE_K, TILE_N, num_warps, num_stages) 

62 configs_map.setdefault(key, config) 

63 else: 

64 key = (TILE_K, TILE_N, num_warps, num_stages) 

65 configs_map.setdefault(key, config) 

66 pruned_configs = [] 

67 for k, v in configs_map.items(): 

68 pruned_configs.append(v) 

69 extra_config = copy.deepcopy(pruned_configs[0]) 

70 extra_config.kwargs["TILE_K"] = 1 

71 extra_config.kwargs["TILE_N"] = N 

72 extra_config.num_warps = 1 

73 extra_config.num_stages = 3 

74 pruned_configs.append(extra_config) 

75 extra_config2 = copy.deepcopy(extra_config) 

76 extra_config2.num_stages = 1 

77 pruned_configs.append(extra_config2) 

78 return pruned_configs 

79 

80 

81def softmax_tile_mode_for_non_inner(M, N, K, TILE_N, TILE_K): 

82 one_tile_k = TILE_K * max(TOTAL_CORE_NUM // M, 1) >= K 

83 one_tile_n = TILE_N >= N 

84 if one_tile_n and one_tile_k: 

85 return 0 

86 elif one_tile_n and not one_tile_k: 

87 return 1 

88 else: 

89 return 2 

90 

91 

92@libentry() 

93@libtuner( 

94 configs=runtime.get_tuned_config("softmax_non_inner"), 

95 key=[ 

96 "N", 

97 "K", 

98 ], 

99 prune_configs_by={"early_config_prune": config_prune1}, 

100) 

101@triton.heuristics(runtime.get_heuristic_config("softmax_non_inner")) 

102@triton.jit 

103def softmax_kernel_non_inner( 

104 output_ptr, 

105 input_ptr, 

106 M, 

107 N, 

108 K, 

109 TILE_N: tl.constexpr, 

110 TILE_K: tl.constexpr, 

111 TILE_MODE: tl.constexpr, 

112): 

113 pid_m = tl.program_id(0) 

114 pid_k = tl.program_id(1) 

115 

116 p_k_num = tl.num_programs(axis=1) 

117 split_k = tl.cdiv(K, p_k_num) 

118 k_start = pid_k * split_k 

119 

120 if TILE_MODE == 0: 

121 n_offset = tl.arange(0, TILE_N) 

122 k_offset = pid_k * TILE_K + tl.arange(0, TILE_K) 

123 offset = pid_m * N * K + n_offset[:, None] * K + k_offset[None, :] 

124 mask = k_offset[None, :] < K 

125 input_ptrs = input_ptr + offset 

126 inp = tl.load(input_ptrs, mask=mask, other=-float("inf")).to(tl.float32) 

127 row_minus_max = inp - tl.max(inp, axis=0)[None, :] 

128 numerator = tl.exp(row_minus_max) 

129 denominator = tl.sum(numerator, axis=0)[None, :] 

130 recip = 1.0 / denominator 

131 softmax_output = numerator * recip 

132 output_ptrs = output_ptr + offset 

133 tl.store(output_ptrs, softmax_output, mask=mask) 

134 elif TILE_MODE == 1: 

135 for k_idx in range(0, split_k, TILE_K): 

136 k_offset = k_start + k_idx + tl.arange(0, TILE_K) 

137 n_offset = tl.arange(0, TILE_N) 

138 offset = pid_m * N * K + n_offset[:, None] * K + k_offset[None, :] 

139 mask = k_offset[None, :] < K 

140 input_ptrs = input_ptr + offset 

141 inp = tl.load(input_ptrs, mask=mask, other=-float("inf")).to(tl.float32) 

142 row_minus_max = inp - tl.max(inp, axis=0)[None, :] 

143 numerator = tl.exp(row_minus_max) 

144 denominator = tl.sum(numerator, axis=0)[None, :] 

145 recip = 1.0 / denominator 

146 softmax_output = numerator * recip 

147 output_ptrs = output_ptr + offset 

148 tl.store(output_ptrs, softmax_output, mask=mask) 

149 else: 

150 for k_idx in range(0, split_k, TILE_K): 

151 k_offset = k_start + k_idx + tl.arange(0, TILE_K) 

152 m = tl.full([TILE_N, TILE_K], value=float("-inf"), dtype=tl.float32) 

153 z = tl.full([TILE_N, TILE_K], value=0.0, dtype=tl.float32) 

154 # specialization does not improve performance inn this example, as tested 

155 for start_n in range(0, N, TILE_N): 

156 n_offset = start_n + tl.arange(0, TILE_N) 

157 offset = pid_m * N * K + n_offset[:, None] * K + k_offset[None, :] 

158 mask = (n_offset[:, None] < N) & (k_offset[None, :] < K) 

159 inp = tl.load(input_ptr + offset, mask=mask, other=-float("inf")).to( 

160 tl.float32 

161 ) 

162 m_new = tl.maximum(m, inp) 

163 all_neg_inf = m_new == float("-inf") 

164 z = tl.where( 

165 all_neg_inf, z, z * tl.exp(m - m_new) + tl.exp(inp - m_new) 

166 ) 

167 m = m_new 

168 m_reduced = tl.max(m, 0) # (TILE_K,) 

169 z = tl.sum(z * tl.exp(m - m_reduced[None, :]), 0) # (TILE_K, ) 

170 recip_z = 1.0 / z 

171 m = m_reduced 

172 # specialization does not improve performance inn this example, as tested 

173 for start_n in range(0, N, TILE_N): 

174 n_offset = start_n + tl.arange(0, TILE_N) 

175 offset = pid_m * N * K + n_offset[:, None] * K + k_offset[None, :] 

176 mask = (n_offset[:, None] < N) & (k_offset[None, :] < K) 

177 inp = tl.load(input_ptr + offset, mask=mask, other=-float("inf")).to( 

178 tl.float32 

179 ) 

180 o = tl.exp(inp - m[None, :]) * recip_z[None, :] 

181 tl.store(output_ptr + offset, o, mask=mask) 

182 

183 

184def config_prune2(configs, named_args, **kwargs): 

185 M = named_args["M"] 

186 N = named_args["N"] 

187 input = named_args["input_ptr"] 

188 configs_map = {} 

189 # When N is less than MAX_C_MLU_SOFTMAX_FORWARD, no reduction loops 

190 for config in configs: 

191 kw = config.kwargs 

192 BLOCK_M, BLOCK_N, num_warps, num_stages = ( 

193 kw["BLOCK_M"], 

194 kw["BLOCK_N"], 

195 config.num_warps, 

196 config.num_stages, 

197 ) 

198 if N < MAX_N: 

199 config = copy.deepcopy(config) 

200 BLOCK_N = config.kwargs["BLOCK_N"] = N 

201 m_per_core = math.ceil(M / TOTAL_CORE_NUM) 

202 BLOCK_M = config.kwargs["BLOCK_M"] = m_per_core 

203 num_stages = config.num_stages = 1 

204 key = (BLOCK_M, BLOCK_N, num_warps, num_stages) 

205 configs_map.setdefault(key, config) 

206 

207 config = copy.deepcopy(config) 

208 max_block_m_without_pipe = MAX_NRAM_SIZE // 4 // (2 * BLOCK_N + 1) 

209 BLOCK_M = config.kwargs["BLOCK_M"] = align(max_block_m_without_pipe) 

210 num_stages = config.num_stages = 1 

211 key = (BLOCK_M, BLOCK_N, num_warps, num_stages) 

212 configs_map.setdefault(key, config) 

213 

214 config = copy.deepcopy(config) 

215 max_block_m_without_pipe = MAX_NRAM_SIZE // 4 // (4 * BLOCK_N + 1) 

216 if input.dtype == torch.float32: 

217 max_block_m_without_pipe = MAX_NRAM_SIZE // 4 // (6 * BLOCK_N + 1) 

218 BLOCK_M = config.kwargs["BLOCK_M"] = align(max_block_m_without_pipe) 

219 num_stages = config.num_stages = 3 

220 key = (BLOCK_M, BLOCK_N, num_warps, num_stages) 

221 configs_map.setdefault(key, config) 

222 key = (BLOCK_M, BLOCK_N, num_warps, num_stages) 

223 # Only keep one config for the same key 

224 configs_map.setdefault(key, config) 

225 pruned_configs = [] 

226 for k, v in configs_map.items(): 

227 pruned_configs.append(v) 

228 # Add a heuristic config. 

229 extra_config = copy.deepcopy(pruned_configs[0]) 

230 extra_config.kwargs["BLOCK_M"] = 1 

231 extra_config.kwargs["BLOCK_N"] = N 

232 extra_config.num_warps = 1 

233 extra_config.num_stages = 3 

234 pruned_configs.append(extra_config) 

235 extra_config2 = copy.deepcopy(extra_config) 

236 extra_config2.num_stages = 1 

237 pruned_configs.append(extra_config2) 

238 return pruned_configs 

239 

240 

241def softmax_tile_mode_for_inner(args): 

242 one_tile_m = args["BLOCK_M"] * TOTAL_CORE_NUM >= args["M"] 

243 one_tile_n = args["BLOCK_N"] >= args["N"] 

244 if one_tile_n and one_tile_m: 

245 return 0 

246 elif one_tile_n and not one_tile_m: 

247 return 1 

248 else: 

249 return 2 

250 

251 

252@libentry() 

253@libtuner( 

254 configs=runtime.get_tuned_config("softmax_inner"), 

255 key=[ 

256 "M", 

257 "N", 

258 ], 

259 prune_configs_by={"early_config_prune": config_prune2}, 

260) 

261@triton.heuristics(runtime.get_heuristic_config("softmax_inner")) 

262@triton.jit 

263def softmax_kernel_inner( 

264 output_ptr, 

265 input_ptr, 

266 M, 

267 N, 

268 BLOCK_M: tl.constexpr, 

269 BLOCK_N: tl.constexpr, 

270 TILE_MODE: tl.constexpr, 

271): 

272 pid_m = tl.program_id(0) 

273 pnum = tl.num_programs(axis=0) 

274 split_m = tl.cdiv(M, pnum) 

275 m_start = pid_m * split_m 

276 

277 if TILE_MODE == 0: 

278 m_offset = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) 

279 n_offset = tl.arange(0, BLOCK_N) 

280 offset = m_offset[:, None] * N + n_offset[None, :] 

281 mask = m_offset[:, None] < M 

282 input_ptrs = input_ptr + offset 

283 inp = tl.load(input_ptrs, mask=mask, other=-float("inf")).to(tl.float32) 

284 row_minus_max = inp - tl.max(inp, axis=1)[:, None] 

285 numerator = tl.exp(row_minus_max) 

286 denominator = tl.sum(numerator, axis=1)[:, None] 

287 recip = 1.0 / denominator 

288 softmax_output = numerator * recip 

289 output_ptrs = output_ptr + offset 

290 tl.store(output_ptrs, softmax_output, mask=mask) 

291 elif TILE_MODE == 1: 

292 for m_idx in range(0, split_m, BLOCK_M): 

293 m_offset = m_start + m_idx + tl.arange(0, BLOCK_M) 

294 n_offset = tl.arange(0, BLOCK_N) 

295 offset = m_offset[:, None] * N + n_offset[None, :] 

296 mask = m_offset[:, None] < M 

297 input_ptrs = input_ptr + offset 

298 inp = tl.load(input_ptrs, mask=mask, other=-float("inf")).to(tl.float32) 

299 trans_inp = tl.trans(inp) 

300 row_minus_max = trans_inp - tl.max(trans_inp, axis=0)[None, :] 

301 numerator = tl.exp(row_minus_max) 

302 denominator = tl.sum(numerator, axis=0)[None, :] 

303 recip = 1.0 / denominator 

304 softmax_output = tl.trans(numerator * recip) 

305 output_ptrs = output_ptr + offset 

306 tl.store(output_ptrs, softmax_output, mask=mask) 

307 else: 

308 for m_idx in range(0, split_m, BLOCK_M): 

309 m_offset = m_start + m_idx + tl.arange(0, BLOCK_M) 

310 block_max = tl.full( 

311 [BLOCK_M, BLOCK_N], value=float("-inf"), dtype=tl.float32 

312 ) 

313 block_sum = tl.full([BLOCK_M, BLOCK_N], value=0.0, dtype=tl.float32) 

314 # specialization does not improve performance inn this example, as tested 

315 for start_n in range(0, N, BLOCK_N): 

316 n_offset = start_n + tl.arange(0, BLOCK_N) 

317 offset = m_offset[:, None] * N + n_offset[None, :] 

318 mask = m_offset[:, None] < M and n_offset[None, :] < N 

319 inp = tl.load(input_ptr + offset, mask=mask, other=-float("inf")).to( 

320 tl.float32 

321 ) 

322 cur_max = tl.maximum(block_max, inp) 

323 all_neg_inf = cur_max == float("-inf") 

324 block_sum = tl.where( 

325 all_neg_inf, 

326 block_sum, 

327 block_sum * tl.exp(block_max - cur_max) + tl.exp(inp - cur_max), 

328 ) 

329 block_max = cur_max 

330 

331 trans_block_max = tl.trans(block_max) 

332 trans_block_sum = tl.trans(block_sum) 

333 max_reduced = tl.max(trans_block_max, 0) 

334 total_sum = tl.sum( 

335 trans_block_sum * tl.exp(trans_block_max - max_reduced[None, :]), 0 

336 ) 

337 recip_total_sum = 1.0 / total_sum 

338 total_max = max_reduced 

339 

340 for start_n in range(0, N, BLOCK_N): 

341 n_offset = start_n + tl.arange(0, BLOCK_N) 

342 offset = m_offset[:, None] * N + n_offset[None, :] 

343 mask = m_offset[:, None] < M and n_offset[None, :] < N 

344 inp = tl.load(input_ptr + offset, mask=mask, other=-float("inf")).to( 

345 tl.float32 

346 ) 

347 o = tl.exp(inp - total_max[:, None]) * recip_total_sum[:, None] 

348 tl.store(output_ptr + offset, o, mask=mask) 

349 

350 

351@triton.jit 

352def softmax_kernel_inner_k_partial_stats( 

353 x_ptr, 

354 max_buf_ptr, 

355 sum_buf_ptr, 

356 M, 

357 N, 

358 T, 

359 BLOCK_M: tl.constexpr, 

360 BLOCK_N: tl.constexpr, 

361): 

362 pnum = tl.num_programs(axis=0) 

363 pid = tl.program_id(0) 

364 total_blocks = (M // BLOCK_M) * T 

365 work_per_core = (total_blocks + pnum - 1) // pnum 

366 start = pid * work_per_core 

367 end = tl.minimum(start + work_per_core, total_blocks) 

368 

369 for task in range(start, end): 

370 row_id = task // T 

371 tile_id = task % T 

372 

373 offs_m = row_id * BLOCK_M + tl.arange(0, BLOCK_M) 

374 offs_n = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) 

375 mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) 

376 

377 tile = tl.load( 

378 x_ptr + offs_m[:, None] * N + offs_n[None, :], 

379 mask=mask, 

380 other=-float("inf"), 

381 ).to(tl.float32) 

382 

383 tile_max = tl.max(tile, axis=1) 

384 all_neg_inf = tile_max == -float("inf") 

385 

386 tile_sum = tl.where( 

387 all_neg_inf, 

388 0.0, 

389 tl.sum(tl.exp(tile - tile_max[:, None]), axis=1), 

390 ) 

391 

392 tl.store(max_buf_ptr + offs_m * T + tile_id, tile_max, mask=(offs_m < M)) 

393 tl.store(sum_buf_ptr + offs_m * T + tile_id, tile_sum, mask=(offs_m < M)) 

394 

395 

396@triton.jit 

397def softmax_kernel_inner_k_merge_stats( 

398 max_buf_ptr, 

399 sum_buf_ptr, 

400 gmax_ptr, 

401 gsum_ptr, 

402 M: tl.constexpr, 

403 T: tl.constexpr, 

404 BLOCK_M: tl.constexpr, 

405): 

406 pid_m = tl.program_id(axis=0) 

407 offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) # [BM] 

408 mask_m = offs_m < M 

409 tile_max = tl.load( 

410 max_buf_ptr + offs_m[:, None] * T + tl.arange(0, T)[None, :], 

411 mask=(offs_m[:, None] < M), 

412 other=-float("inf"), 

413 ) 

414 tile_sum = tl.load( 

415 sum_buf_ptr + offs_m[:, None] * T + tl.arange(0, T)[None, :], 

416 mask=(offs_m[:, None] < M), 

417 other=0.0, 

418 ).to(tl.float32) 

419 

420 gmax = tl.max(tile_max, axis=1) 

421 scale = tl.exp(tile_max - gmax[:, None]) 

422 scale = tl.where(gmax[:, None] == -float("inf"), 0.0, scale) 

423 gsum = tl.sum(tile_sum * scale, axis=1) 

424 

425 tl.store(gmax_ptr + offs_m, gmax, mask=mask_m) 

426 tl.store(gsum_ptr + offs_m, gsum, mask=mask_m) 

427 

428 

429@triton.jit 

430def softmax_kernel_inner_k_write_softmax( 

431 x_ptr, 

432 y_ptr, 

433 gmax_ptr, 

434 gsum_ptr, 

435 M, 

436 N, 

437 T, 

438 BLOCK_M: tl.constexpr, 

439 BLOCK_N: tl.constexpr, 

440): 

441 pnum = tl.num_programs(axis=0) 

442 pid = tl.program_id(0) 

443 total_blocks = (M // BLOCK_M) * T 

444 work_per_core = (total_blocks + pnum - 1) // pnum 

445 start = pid * work_per_core 

446 end = tl.minimum(start + work_per_core, total_blocks) 

447 

448 for task in range(start, end): 

449 row_id = task // T 

450 tile_id = task % T 

451 

452 offs_m = row_id * BLOCK_M + tl.arange(0, BLOCK_M) 

453 offs_n = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) 

454 mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) 

455 

456 # load global stats 

457 gmax = tl.load(gmax_ptr + offs_m, mask=(offs_m < M), other=-float("inf")).to( 

458 tl.float32 

459 ) 

460 gsum = tl.load(gsum_ptr + offs_m, mask=(offs_m < M), other=0.0).to(tl.float32) 

461 

462 # load tile 

463 tile = tl.load( 

464 x_ptr + offs_m[:, None] * N + offs_n[None, :], 

465 mask=mask, 

466 other=-float("inf"), 

467 ).to(tl.float32) 

468 

469 valid = gsum[:, None] > 0 

470 

471 out = tl.where( 

472 valid, 

473 tl.exp(tile - gmax[:, None]) / gsum[:, None], 

474 0.0, 

475 ) 

476 

477 tl.store(y_ptr + offs_m[:, None] * N + offs_n[None, :], out, mask=mask) 

478 

479 

480# ------------------------ backward ------------------------------- 

481 

482 

483def nram_usage_for_backward_non_inner(bn, bk, tile_mode, num_stages, dtype): 

484 coef = 1 

485 if tile_mode == 0: 

486 coef = 3 

487 elif tile_mode == 1: 

488 if num_stages == 1: 

489 coef = 3 

490 else: 

491 if dtype == torch.float32: 

492 coef = 7 

493 else: 

494 coef = 6 

495 else: 

496 if num_stages == 1: 

497 coef = 5 

498 else: 

499 if dtype == torch.float32: 

500 coef = 13 

501 else: 

502 coef = 10 

503 return (coef * bn + 1) * bk * 4 

504 

505 

506def config_prune3(configs, named_args, **kwargs): 

507 M = named_args["M"] 

508 N = named_args["N"] 

509 K = named_args["K"] 

510 output = named_args["output_ptr"] 

511 dtype = output.dtype 

512 k_per_core = math.ceil(K / max(TOTAL_CORE_NUM // M, 1)) 

513 # No need for any loop. 

514 if nram_usage_for_backward_non_inner(N, k_per_core, 0, 1, dtype) < MAX_NRAM_SIZE: 

515 config = copy.deepcopy(configs[0]) 

516 config.kwargs["TILE_K"] = k_per_core 

517 config.kwargs["TILE_N"] = N 

518 config.num_stages = 1 

519 return [config] 

520 align_num = 256 // 4 if dtype == torch.float32 else 256 // 2 

521 pruned_configs = [] 

522 for config in configs: 

523 kw = config.kwargs 

524 TILE_K, TILE_N, num_stages = ( 

525 kw["TILE_K"], 

526 kw["TILE_N"], 

527 config.num_stages, 

528 ) 

529 # Align the lowest dimension to 256B while loading/storing data. 

530 if TILE_K % align_num != 0: 

531 continue 

532 # nram usage shoule be smaller than MAX_NRAM_SIZE 

533 mode = softmax_tile_mode_for_non_inner(M, N, K, TILE_N, TILE_K) 

534 nram = nram_usage_for_backward_non_inner( 

535 TILE_N, TILE_K, mode, num_stages, dtype 

536 ) 

537 if nram > MAX_NRAM_SIZE or nram < MAX_NRAM_SIZE // 2: 

538 continue 

539 pruned_configs.append(config) 

540 return pruned_configs 

541 

542 

543@libentry() 

544@libtuner( 

545 configs=runtime.get_tuned_config("softmax_non_inner_bw"), 

546 key=[ 

547 "N", 

548 "K", 

549 ], 

550 prune_configs_by={"early_config_prune": config_prune3}, 

551) 

552@triton.heuristics(runtime.get_heuristic_config("softmax_backward_non_inner")) 

553@triton.jit 

554def softmax_backward_kernel_non_inner( 

555 output_ptr, 

556 out_grad_ptr, 

557 in_grad_ptr, 

558 M, 

559 N, 

560 K, 

561 TILE_N: tl.constexpr, 

562 TILE_K: tl.constexpr, 

563 TILE_MODE: tl.constexpr, 

564): 

565 pid_m = tl.program_id(0) 

566 pid_k = tl.program_id(1) 

567 

568 p_k_num = tl.num_programs(axis=1) 

569 split_k = tl.cdiv(K, p_k_num) 

570 k_start = pid_k * split_k 

571 

572 if TILE_MODE == 0: 

573 n_offset = tl.arange(0, TILE_N) 

574 k_offset = pid_k * TILE_K + tl.arange(0, TILE_K) 

575 offset = pid_m * N * K + n_offset[:, None] * K + k_offset[None, :] 

576 mask = k_offset[None, :] < K 

577 out_tile = tl.load(output_ptr + offset, mask=mask).to(tl.float32) 

578 out_grad_tile = tl.load(out_grad_ptr + offset, mask=mask).to(tl.float32) 

579 scale = tl.sum(out_tile * out_grad_tile, axis=0) 

580 in_grad_tile = out_tile * (out_grad_tile - scale[None, :]) 

581 tl.store(in_grad_ptr + offset, in_grad_tile, mask=mask) 

582 elif TILE_MODE == 1: 

583 for k_idx in range(0, split_k, TILE_K): 

584 k_offset = k_start + k_idx + tl.arange(0, TILE_K) 

585 n_offset = tl.arange(0, TILE_N) 

586 offset = pid_m * N * K + n_offset[:, None] * K + k_offset[None, :] 

587 mask = k_offset[None, :] < K and n_offset[:, None] < N 

588 out_tile = tl.load(output_ptr + offset, mask=mask).to(tl.float32) 

589 out_grad_tile = tl.load(out_grad_ptr + offset, mask=mask).to(tl.float32) 

590 scale = tl.sum(out_tile * out_grad_tile, axis=0) 

591 in_grad_tile = out_tile * (out_grad_tile - scale[None, :]) 

592 tl.store(in_grad_ptr + offset, in_grad_tile, mask=mask) 

593 else: 

594 for k_idx in range(0, split_k, TILE_K): 

595 k_offset = k_start + k_idx + tl.arange(0, TILE_K) 

596 scale = tl.zeros([TILE_N, TILE_K], dtype=tl.float32) 

597 # specialization does not improve performance inn this example, as tested 

598 for start_n in range(0, N, TILE_N): 

599 n_offset = start_n + tl.arange(0, TILE_N) 

600 offset = pid_m * N * K + n_offset[:, None] * K + k_offset[None, :] 

601 mask = (n_offset[:, None] < N) & (k_offset[None, :] < K) 

602 out_tile = tl.load(output_ptr + offset, mask=mask).to(tl.float32) 

603 out_grad_tile = tl.load(out_grad_ptr + offset, mask=mask).to(tl.float32) 

604 scale += out_tile * out_grad_tile 

605 scale = tl.sum(scale, axis=0) 

606 for start_n in range(0, N, TILE_N): 

607 n_offset = start_n + tl.arange(0, TILE_N) 

608 offset = pid_m * N * K + n_offset[:, None] * K + k_offset[None, :] 

609 mask = (n_offset[:, None] < N) & (k_offset[None, :] < K) 

610 out_tile = tl.load(output_ptr + offset, mask=mask).to(tl.float32) 

611 out_grad_tile = tl.load(out_grad_ptr + offset, mask=mask).to(tl.float32) 

612 in_grad_tile = out_tile * (out_grad_tile - scale[None, :]) 

613 tl.store(in_grad_ptr + offset, in_grad_tile, mask=mask) 

614 

615 

616def config_prune4(configs, named_args, **kwargs): 

617 M = named_args["M"] 

618 N = named_args["N"] 

619 output = named_args["output_ptr"] 

620 configs_map = {} 

621 # When N is less than MAX_C_MLU_SOFTMAX_FORWARD, no reduction loops 

622 for config in configs: 

623 kw = config.kwargs 

624 BLOCK_M, BLOCK_N, num_warps, num_stages = ( 

625 kw["BLOCK_M"], 

626 kw["BLOCK_N"], 

627 config.num_warps, 

628 config.num_stages, 

629 ) 

630 if N < MAX_N: 

631 config = copy.deepcopy(config) 

632 BLOCK_N = config.kwargs["BLOCK_N"] = N 

633 m_per_core = math.ceil(M / TOTAL_CORE_NUM) 

634 BLOCK_M = config.kwargs["BLOCK_M"] = m_per_core 

635 num_stages = config.num_stages = 1 

636 key = (BLOCK_M, BLOCK_N, num_warps, num_stages)