()
| 424 | num_warps = metadata.num_warps |
| 425 | num_stages = metadata.num_stages |
| 426 | cluster_x, cluster_y, cluster_z = unpack_grid((metadata.num_ctas,)) |
| 427 | shared_memory = metadata.shared |
| 428 | return { |
| 429 | "name": f"fused_mm_sample_triton_<grid:{grid_x}x{grid_y}x{grid_z}>_<cluster:{cluster_x}x{cluster_y}x{cluster_z}>_<warps:{num_warps}>_<shared:{shared_memory}>_<stages:{num_stages}>", |
| 430 | } |
| 431 | |
| 432 | |
| 433 | def unpack_grid(grid): |
| 434 | if len(grid) == 1: |
| 435 | return grid[0], 1, 1 |
| 436 | if len(grid) == 2: |
| 437 | return grid[0], grid[1], 1 |
| 438 | if len(grid) == 3: |
| 439 | return grid[0], grid[1], grid[2] |
| 440 | |
| 441 | |
| 442 | def get_autotuning_configs() -> list[triton.Config]: |
| 443 | cc = torch.cuda.get_device_capability() |
| 444 | is_dev_machine: bool = cc == (8, 6) # RTX 3090 config |
| 445 | if is_dev_machine: |
| 446 | return [ |
| 447 | triton.Config( |
| 448 | {"BLOCK_SIZE_V": MIN_BLOCK_SIZE_V, "BLOCK_SIZE_D": 32, "GROUP_SIZE_V": 4}, |
| 449 | num_warps=4, |
| 450 | num_stages=2, |
| 451 | # Persistent kernel: grid = NUM_SMS, so only 1 block per SM. |
| 452 | # No occupancy benefit from limiting registers, so let ptxas |
| 453 | # use the full register file instead of spilling to local memory. |
| 454 | maxnreg=255, |
| 455 | ) |
| 456 | ] |
| 457 | return [ |
nothing calls this directly
no outgoing calls
no test coverage detected