From 4d06fb8d21cff895d713e508fa7638997bbd6de9 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Mon, 5 Oct 2026 15:44:35 +0000 Subject: [PATCH 1/4] [FlyDSL] Add full-expert DeepSeek V4.1 Flash MegaMoE tuning for gfx950 Co-authored-by: Codex Signed-off-by: Fangzhou Ai --- aiter/configs/mega_moe_flash_gfx950.json | 3130 +++++++++++++++++ aiter/ops/flydsl/kernels/mega_moe/dispatch.py | 45 +- aiter/ops/flydsl/kernels/mega_moe/envs.py | 4 + aiter/ops/flydsl/kernels/mega_moe/gemm1.py | 4 +- aiter/ops/flydsl/kernels/mega_moe/gemm2.py | 12 +- .../ops/flydsl/kernels/mega_moe/gemm_util.py | 4 +- .../kernels/mega_moe/mega_moe_config.py | 208 +- .../kernels/mega_moe/mega_moe_stage1.py | 16 +- .../kernels/mega_moe/mega_moe_stage2.py | 2 +- .../flydsl/kernels/mega_moe/mega_moe_v2.py | 53 +- op_tests/multigpu_tests/bench_mega_moe_v2.py | 223 +- op_tests/multigpu_tests/test_mega_moe_v2.py | 136 +- op_tests/multigpu_tests/tune_mega_moe_v2.py | 516 +++ 13 files changed, 4223 insertions(+), 130 deletions(-) create mode 100644 aiter/configs/mega_moe_flash_gfx950.json create mode 100644 op_tests/multigpu_tests/tune_mega_moe_v2.py diff --git a/aiter/configs/mega_moe_flash_gfx950.json b/aiter/configs/mega_moe_flash_gfx950.json new file mode 100644 index 00000000000..1388f02f79a --- /dev/null +++ b/aiter/configs/mega_moe_flash_gfx950.json @@ -0,0 +1,3130 @@ +{ + "schema_version": 1, + "architecture": "gfx950", + "shape": [ + 5120, + 2304, + 384, + 6 + ], + "fixed": { + "2": { + "configs": [ + { + "stage1": { + "sort_block_m": 16, + "tile_n": 128, + "num_waves": 2, + "grid_mult": 1, + "num_dispatch_cu": 8, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false + }, + "p2p_quant": "none", + "combine_blocks": 32, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 32, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "waves_per_eu_hint": 1, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 32, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 32, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 16, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 2 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "b_hoist": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "b_hoist": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 32, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + } + ], + "tokens": { + "1": 0, + "2": 1, + "3": 2, + "4": 3, + "5": 4, + "6": 5, + "7": 6, + "8": 7, + "9": 8, + "10": 9, + "11": 10, + "12": 11, + "13": 12, + "14": 13, + "15": 14, + "16": 15, + "17": 16, + "18": 17, + "19": 17, + "20": 18, + "21": 19, + "22": 20, + "23": 21, + "24": 22, + "25": 23, + "26": 24, + "27": 25, + "28": 26, + "29": 27, + "30": 28, + "31": 29, + "32": 30, + "33": 24, + "34": 31, + "35": 32, + "36": 33, + "37": 34, + "38": 35, + "39": 36, + "40": 27, + "41": 37, + "42": 38, + "43": 28, + "44": 39, + "45": 40, + "46": 41, + "47": 42, + "48": 43, + "49": 25, + "50": 44, + "51": 45, + "52": 20, + "53": 38, + "54": 46, + "55": 47, + "56": 48, + "57": 31, + "58": 49, + "59": 24, + "60": 50, + "61": 48, + "62": 51, + "63": 44, + "64": 52, + "96": 48, + "128": 53, + "192": 54, + "256": 55, + "512": 56 + } + }, + "4": { + "configs": [ + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 4, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "waves_per_eu_hint": 1, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 16, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 8, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "waves_per_eu_hint": 1, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false + }, + "p2p_quant": "none", + "combine_blocks": 16, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "b_hoist": false + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 32, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 32, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "waves_per_eu_hint": 1, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "waves_per_eu_hint": 1, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 2 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 202, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "waves_per_eu_hint": 1, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 0, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "b_hoist": false, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "ascale_prefetch": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "ascale_prefetch": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "ascale_prefetch": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 64, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 8 + } + ], + "tokens": { + "1": 0, + "2": 1, + "3": 2, + "4": 3, + "5": 4, + "6": 5, + "7": 6, + "8": 7, + "9": 8, + "10": 9, + "11": 10, + "12": 11, + "13": 12, + "14": 13, + "15": 14, + "16": 15, + "17": 16, + "18": 17, + "19": 18, + "20": 19, + "21": 20, + "22": 21, + "23": 22, + "24": 23, + "25": 24, + "26": 25, + "27": 26, + "28": 27, + "29": 28, + "30": 29, + "31": 30, + "32": 31, + "33": 32, + "34": 26, + "35": 33, + "36": 34, + "37": 35, + "38": 36, + "39": 37, + "40": 38, + "41": 39, + "42": 28, + "43": 39, + "44": 40, + "45": 28, + "46": 41, + "47": 42, + "48": 43, + "49": 44, + "50": 45, + "51": 24, + "52": 39, + "53": 28, + "54": 25, + "55": 46, + "56": 47, + "57": 48, + "58": 49, + "59": 42, + "60": 50, + "61": 51, + "62": 48, + "63": 42, + "64": 52, + "96": 53, + "128": 54, + "192": 55, + "256": 56, + "512": 57, + "1024": 58 + } + } + } +} diff --git a/aiter/ops/flydsl/kernels/mega_moe/dispatch.py b/aiter/ops/flydsl/kernels/mega_moe/dispatch.py index 79ccff13962..cdb38436505 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/dispatch.py +++ b/aiter/ops/flydsl/kernels/mega_moe/dispatch.py @@ -575,7 +575,7 @@ def emit_direct_fixed_slot_finalize( *, fz_npes, fz_epr, fz_cap, fz_mtpr, fz_rank, fz_tile_m, n_tiles, addr_disp, parity, expected ): """Finalize local fixed slots as soon as every source publishes this destination.""" - assert 0 < fz_epr <= 64, "direct fixed-slot finalize requires 1..64 experts per rank" + assert 0 < fz_epr <= 256 dispatch_table = ptr_buf_tensor(addr_disp, fx.Int64) def dp(i): @@ -612,23 +612,27 @@ def dp(i): comm_ops.wait_i32_until_equals(a_source_done + fx.Int64(done_index) * fx.Int64(4), expected) comm_ops.fence_system_acquire() - valid_expert = lane < fx.Int32(fz_epr) - safe_expert = valid_expert.select(lane, fx.Int32(0)) - count = running[safe_expert] - count = valid_expert.select(count, fx.Int32(0)) - overflow_flag = (count > fx.Int32(fz_cap)).select(fx.Int32(1), fx.Int32(0)) - overflow_prefix = _wave_inclusive_scan_i32(overflow_flag, lane) - overflow_count = fx.Int32(fx.rocdl.readlane(T.i32, overflow_prefix, fz_epr - 1)) - no_overflow = overflow_count == fx.Int32(0) - safe_count = (count <= fx.Int32(fz_cap)).select(count, fx.Int32(0)) - num_expert_tiles = (safe_count + fx.Int32(fz_tile_m - 1)) // fx.Int32(fz_tile_m) - max_expert_tiles = _wave_reduce_max_i32(num_expert_tiles, lane) - inclusive_tiles = _wave_inclusive_scan_i32(num_expert_tiles, lane) - metadata_base = inclusive_tiles - num_expert_tiles - total_tiles = fx.Int32(fx.rocdl.readlane(T.i32, inclusive_tiles, fz_epr - 1)) - - if valid_expert: - if no_overflow: + total_tiles = fx.Int32(0) + overflow_count = fx.Int32(0) + max_expert_tiles = fx.Int32(0) + for chunk in range_constexpr((fz_epr + 63) // 64): + expert = lane + fx.Int32(chunk * 64) + valid_expert = expert < fx.Int32(fz_epr) + safe_expert = valid_expert.select(expert, fx.Int32(0)) + count = valid_expert.select(running[safe_expert], fx.Int32(0)) + overflow_flag = (count > fx.Int32(fz_cap)).select(fx.Int32(1), fx.Int32(0)) + overflow_prefix = _wave_inclusive_scan_i32(overflow_flag, lane) + last_lane = min(63, fz_epr - chunk * 64 - 1) + overflow_count += fx.Int32(fx.rocdl.readlane(T.i32, overflow_prefix, last_lane)) + safe_count = (count <= fx.Int32(fz_cap)).select(count, fx.Int32(0)) + num_expert_tiles = (safe_count + fx.Int32(fz_tile_m - 1)) // fx.Int32(fz_tile_m) + chunk_max = _wave_reduce_max_i32(num_expert_tiles, lane) + max_expert_tiles = (chunk_max > max_expert_tiles).select(chunk_max, max_expert_tiles) + inclusive_tiles = _wave_inclusive_scan_i32(num_expert_tiles, lane) + metadata_base = total_tiles + inclusive_tiles - num_expert_tiles + total_tiles += fx.Int32(fx.rocdl.readlane(T.i32, inclusive_tiles, last_lane)) + + if valid_expert: global_expert = fx.Int32(fz_rank * fz_epr) + safe_expert payload_base = safe_expert * fx.Int32(fz_cap) for tile in range(fx.Int32(0), num_expert_tiles, 1): @@ -646,10 +650,9 @@ def dp(i): fz_npes * fz_mtpr ) expert_tile_end[safe_expert] = metadata_base + num_expert_tiles - else: - expert_tile_end[safe_expert] = fx.Int32(0) - running[safe_expert] = fx.Int32(0) + running[safe_expert] = fx.Int32(0) + no_overflow = overflow_count == fx.Int32(0) if lane == fx.Int32(0): num_valid = no_overflow.select(total_tiles * fx.Int32(fz_tile_m), fx.Int32(0)) ready_work = no_overflow.select(total_tiles * fx.Int32(n_tiles), fx.Int32(0)) diff --git a/aiter/ops/flydsl/kernels/mega_moe/envs.py b/aiter/ops/flydsl/kernels/mega_moe/envs.py index 82e9b716614..a1f2df14438 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/envs.py +++ b/aiter/ops/flydsl/kernels/mega_moe/envs.py @@ -16,6 +16,10 @@ def _flag(value: str) -> bool: return value == "1" +def fixed_slot_max_mtpr(default: int) -> int: + return int(os.environ.get("AITER_MEGA_FIXED_SLOT_MAX_MTPR", str(default))) + + # name -> (parser, default) environment_variables: dict[str, tuple[Callable[[str], Any], str]] = { # Largest MTPR still dispatched fixed-slot (direct expert slots, no count diff --git a/aiter/ops/flydsl/kernels/mega_moe/gemm1.py b/aiter/ops/flydsl/kernels/mega_moe/gemm1.py index 3e05c3e2ccf..6bdf1e1e21e 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/gemm1.py +++ b/aiter/ops/flydsl/kernels/mega_moe/gemm1.py @@ -38,7 +38,7 @@ def do_tile(m_tile, n_tile_base, expert, sched, a_gather, a_s2r, b_loader, b_sca # fmt: on N_ACC = M_REPEAT * NUM_ACC_N NUM_B_SCALE = NUM_ACC_N // _PACK - NUM_A_SCALE = M_REPEAT // _PACK + NUM_A_SCALE = (M_REPEAT + _PACK - 1) // _PACK B_STATE_END = ( N_ACC + NUM_ACC_N * _PACK ) @@ -378,7 +378,7 @@ def compile_gemm1( n_tiles = (2 * inter_dim) // tile_n m_repeat = sort_block_m // 16 num_acc_n = n_per_wave // 16 - assert num_acc_n % 2 == 0 and m_repeat % 2 == 0 + assert num_acc_n % 2 == 0 and sort_block_m in (16, 32, 64, 128) a_k_step_bytes = tile_k k_iters = model_dim // tile_k diff --git a/aiter/ops/flydsl/kernels/mega_moe/gemm2.py b/aiter/ops/flydsl/kernels/mega_moe/gemm2.py index 6ca3a5eeaa1..ca609bfc980 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/gemm2.py +++ b/aiter/ops/flydsl/kernels/mega_moe/gemm2.py @@ -186,6 +186,7 @@ def gemm2_compute_v2( explicit_m_row=None, explicit_n_block=None, explicit_expert=None, + dense_bm16_scales=False, ): """Run GEMM2, optionally using an explicitly selected expert row/tile.""" # SBM is the sort padding unit; BM is the compute tile and must divide SBM. @@ -324,7 +325,9 @@ def issue_a_ds_read(slot): asc_per_mb = fx.Int32(kScaleSubBlocks) * kAS_per_chunk_dw * fx.Int32(4) asc_num = fx.Int64(i32_max_m_blocks) * fx.Int64(asc_per_mb) - scale_chunk0 = m_block_idx if const_expr(is_bm16) else m_row // 32 + scale_chunk0 = ( + m_block_idx if const_expr(is_bm16 and not dense_bm16_scales) else m_row // 32 + ) def make_ascale_view(sub): base_dw = (scale_chunk0 + fx.Int32(sub)) * kAS_per_chunk_dw @@ -355,7 +358,12 @@ def load_a_scale_tile(kt): ascale_views[sub][lane_div_16, lane_mod_16, chunk_kt, None], saf, ) - out.append(Vec(saf.load())[0]) + scale_word = Vec(saf.load())[0] + if const_expr(is_bm16 and dense_bm16_scales): + scale_word = scale_word.shrui( + ((m_row // fx.Int32(16)) & fx.Int32(1)) * fx.Int32(8) + ) + out.append(scale_word) return out # Stream B weights and scales through registers so use_nt reaches the ISA cache policy. diff --git a/aiter/ops/flydsl/kernels/mega_moe/gemm_util.py b/aiter/ops/flydsl/kernels/mega_moe/gemm_util.py index 2a550dec39e..9157b14490b 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/gemm_util.py +++ b/aiter/ops/flydsl/kernels/mega_moe/gemm_util.py @@ -355,7 +355,7 @@ def __init__( self._rsrc = scale_rsrc self._n_scale = model_dim // 32 self._lane = fx.thread_idx.x % 64 - self._n_groups = m_repeat // _PACK + self._n_groups = (m_repeat + _PACK - 1) // _PACK self._sort_block_m = sort_block_m self._total_threads = total_threads self._tx = fx.thread_idx.x @@ -440,7 +440,7 @@ def load_step(self, lds_ascale, kstep_i32): out = [] for g in range_constexpr(self._n_groups): r0 = fx.Int32(g * 32) + lane_row - r1 = r0 + fx.Int32(16) + r1 = r0 + fx.Int32(16 if self._sort_block_m >= 32 else 0) b = [] for ksub in range_constexpr(_PACK): for rr in (r0, r1): diff --git a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py index db77b4bac41..ba01a2a7ca2 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py +++ b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py @@ -2,9 +2,11 @@ # Copyright (c) 2025 FlyDSL Project Contributors """Static MegaMoEV2 configuration rules for MI355X.""" +import json from bisect import bisect_left from dataclasses import dataclass, replace from functools import cache +from pathlib import Path from . import envs @@ -25,6 +27,7 @@ 16384, 32768, ) +FLASH_TOKEN_BUCKETS = tuple(range(1, 65)) + (96, 128, 192) + TOKEN_BUCKETS[7:] FIXED_GRID_MULT_VALUES = (1, 2, 3, 4, 6, 8, 12, 16, 24, 32) P2P_FP8_MIN_MTPR = 1024 # Fixed-slot (direct expert slots, no count exchange) is the default up to MTPR @@ -76,6 +79,7 @@ class Stage1Config: work_shards: int = 8 payload_chunk_rows: int = 0 prepare_quant_cu: int = 64 + pair_k: bool = False def stage1_bundle_identity(config: Stage1Config) -> Stage1Config: @@ -108,6 +112,8 @@ class MegaMoEConfig: stage1: Stage1Config stage2: Stage2Config p2p_quant: str + combine_blocks: int = 128 + combine_waves: int = 8 def __post_init__(self): sbm = self.stage1.sort_block_m @@ -161,7 +167,8 @@ def entry_for_tokens(self, tokens: int) -> MegaMoEBundleEntry: # Empty DP ranks must still launch the same collective MegaMoE protocol # as non-empty ranks. Use the smallest bundle geometry for that rank; # returning early would strand peers that have already entered dispatch. - bucket = TOKEN_BUCKETS[0] if tokens == 0 else nearest_token_bucket(tokens) + buckets = tuple(entry.token_bucket for entry in self.entries) + bucket = buckets[0] if tokens == 0 else nearest_token_bucket(tokens, buckets) for entry in self.entries: if entry.token_bucket == bucket: return entry @@ -170,15 +177,15 @@ def entry_for_tokens(self, tokens: int) -> MegaMoEBundleEntry: ) -def nearest_token_bucket(tokens: int) -> int: +def nearest_token_bucket(tokens: int, buckets=TOKEN_BUCKETS) -> int: if tokens <= 0: raise ValueError(f"tokens must be positive, got {tokens}") - index = bisect_left(TOKEN_BUCKETS, tokens) + index = bisect_left(buckets, tokens) if index == 0: - return TOKEN_BUCKETS[0] - if index == len(TOKEN_BUCKETS): - return TOKEN_BUCKETS[-1] - lower, upper = TOKEN_BUCKETS[index - 1], TOKEN_BUCKETS[index] + return buckets[0] + if index == len(buckets): + return buckets[-1] + lower, upper = buckets[index - 1], buckets[index] return upper if upper - tokens <= tokens - lower else lower @@ -186,6 +193,15 @@ def mtpr_config_class(mtpr: int) -> int: return mtpr if mtpr <= P2P_FP8_MIN_MTPR else MAX_MTPR_CLASS +def _use_fixed_slots(mtpr: int, world_size: int, experts_per_rank: int) -> bool: + geometry = (world_size, experts_per_rank) + if geometry in ((2, 192), (4, 96)): + limit = envs.fixed_slot_max_mtpr(512 if world_size == 2 else 1024) + else: + limit = FIXED_SLOT_MAX_MTPR + return mtpr <= limit and geometry in ((8, 48), (2, 192), (4, 96)) + + def _fixed_dispatch_cu(bucket: int) -> int: if bucket <= 1: return 64 @@ -358,6 +374,140 @@ def _select_large_stage2( ) +@cache +def _flash_tuning_data(): + path = Path(__file__).resolve().parents[4] / "configs/mega_moe_flash_gfx950.json" + return json.loads(path.read_text()) + + +def _select_flash_config(bucket, mtpr_class, world_size, fixed): + if fixed: + policy = _flash_tuning_data()["fixed"][str(world_size)] + config_id = policy["tokens"].get(str(bucket)) + if config_id is not None: + config = policy["configs"][config_id] + return MegaMoEConfig( + stage1=Stage1Config(**config["stage1"]), + stage2=Stage2Config(**config["stage2"]), + p2p_quant=config["p2p_quant"], + combine_blocks=config["combine_blocks"], + combine_waves=config["combine_waves"], + ) + bucket = nearest_token_bucket(bucket) + # W13: rows, columns, waves, producer CUs, work shards, weight cache hint. + # W2: columns, persistent CUs, non-temporal weights, strided scheduling. + if fixed and world_size == 2: + rows = { + 1: ((16, 128, 2, 8, 4, 3), (128, 256, True, False), (16, 4)), + 4: ((16, 256, 4, 16, 1, 3), (256, 256, True, False), (64, 4)), + 8: ((16, 512, 4, 16, 4, 3), (256, 256, True, False), (128, 4)), + 16: ((16, 256, 4, 64, 4, 0), (256, 128, True, False), (128, 8)), + 32: ((32, 256, 8, 64, 4, 0), (256, 192, True, False), (128, 8)), + 64: ((32, 512, 8, 128, 4, 0), (256, 256, False, False), (128, 8)), + 128: ((32, 512, 8, 128, 4, 0), (256, 256, False, False), (32, 4)), + 256: ((32, 512, 8, 128, 4, 0), (256, 256, False, False), (128, 8)), + 512: ((32, 512, 8, 128, 4, 0), (128, 256, True, False), (128, 8)), + 1024: ((64, 512, 8, 128, 4, 0), (128, 256, True, False), (256, 8)), + } + elif fixed: + rows = { + 1: ((16, 256, 2, 32, 4, 3), (256, 256, True, False), (16, 4)), + 4: ((16, 256, 2, 32, 4, 3), (256, 128, True, False), (16, 4)), + 8: ((16, 512, 4, 32, 4, 3), (256, 256, True, False), (128, 8)), + 16: ((32, 256, 8, 32, 2, 3), (128, 256, True, False), (32, 4)), + 32: ((32, 256, 8, 64, 4, 0), (256, 256, True, True), (128, 8)), + 64: ((32, 256, 8, 32, 4, 3), (256, 128, True, False), (128, 8)), + 128: ((16, 512, 4, 64, 4, 0), (256, 256, True, False), (64, 4)), + 256: ((32, 512, 8, 32, 4, 3), (256, 128, True, False), (128, 8)), + 512: ((32, 512, 8, 128, 4, 0), (256, 256, False, False), (128, 8)), + 1024: ((64, 512, 8, 128, 4, 0), (128, 256, False, False), (128, 8)), + } + elif world_size == 2: + rows = { + 1: ((16, 128, 2, 128, 4, 3), (128, 256, True, False), (64, 4)), + 4: ((16, 256, 4, 128, 4, 3), (128, 256, True, False), (64, 4)), + 8: ((16, 512, 4, 128, 1, 3), (128, 256, True, False), (64, 4)), + 16: ((16, 512, 4, 16, 4, 0), (256, 192, True, False), (64, 4)), + 32: ((16, 512, 4, 16, 4, 0), (256, 192, True, False), (64, 4)), + 64: ((16, 512, 4, 128, 4, 0), (128, 256, True, False), (32, 4)), + 128: ((16, 512, 4, 128, 4, 0), (256, 192, True, False), (128, 8)), + 256: ((16, 512, 4, 128, 4, 0), (256, 256, True, True), (128, 8)), + 512: ((32, 512, 8, 32, 4, 0), (128, 256, True, False), (128, 8)), + 1024: ((64, 512, 8, 32, 1, 3), (128, 256, True, False), (128, 8)), + } + else: + rows = { + 1: ((16, 256, 4, 32, 1, 3), (256, 256, True, False), (32, 4)), + 4: ((16, 512, 4, 64, 2, 3), (512, 128, True, False), (128, 4)), + 8: ((16, 512, 4, 64, 2, 3), (512, 128, True, False), (128, 4)), + 16: ((16, 512, 4, 32, 4, 0), (256, 256, True, True), (128, 4)), + 32: ((16, 512, 4, 32, 4, 0), (256, 256, True, True), (128, 4)), + 64: ((16, 512, 4, 128, 4, 0), (512, 256, True, False), (128, 4)), + 128: ((16, 512, 4, 64, 4, 0), (256, 256, True, True), (128, 8)), + 256: ((32, 512, 8, 128, 1, 3), (256, 128, True, False), (64, 4)), + 512: ((64, 512, 8, 64, 1, 3), (128, 256, True, False), (128, 8)), + 1024: ((64, 512, 8, 32, 4, 0), (128, 256, True, False), (128, 8)), + } + if bucket <= 1024: + (bm, bn, waves, producers, shards, hint), (n2, cu2, nt2, strided), combine = ( + rows[bucket] + ) + stage1 = Stage1Config( + sort_block_m=bm, + tile_n=bn, + num_waves=waves, + grid_mult=1, + num_dispatch_cu=producers, + mfma_amajor=bm >= 32, + async_a_copy=True, + use_tile_resource=True, + b_nt=hint, + work_shards=shards, + payload_chunk_rows=0 if fixed else 256, + ) + stage2 = Stage2Config( + block_m=min(bm, 64), + block_n=n2, + persist=True, + persist_cu=cu2, + use_nt=nt2, + persist_strided=strided, + ) + if not fixed: + if world_size == 2 and bucket == 64: + stage2 = replace(stage2, b_hoist=False) + elif world_size == 2 and bucket == 128: + stage2 = replace(stage2, spatial_partition=0) + elif world_size == 2 and bucket == 512 and mtpr_class <= 1024: + stage2 = replace(stage2, bf16_lds=True) + elif world_size == 2 and bucket == 1024 and mtpr_class <= 1024: + stage1 = replace(stage1, payload_chunk_rows=1024, prepare_quant_cu=32) + stage2 = replace(stage2, bf16_lds=True) + elif world_size == 4 and bucket == 64: + stage2 = replace(stage2, spatial_partition=102) + elif world_size == 4 and bucket == 128: + stage1 = replace(stage1, mfma_amajor=True) + stage2 = replace(stage2, b_hoist=False) + elif world_size == 4 and bucket == 256: + stage2 = replace(stage2, spatial_partition=0) + elif world_size == 4 and bucket >= 512: + stage2 = replace(stage2, spatial_partition=202) + else: + stage1 = _select_large_stage1(bucket, 2304) + # N512 at M128 spills registers and consumes 151.6 KiB LDS. + if stage1.sort_block_m == 128: + stage1 = replace(stage1, tile_n=256) + stage2 = _select_large_stage2(bucket, stage1.sort_block_m, 5120) + combine = (128, 8) + return MegaMoEConfig( + stage1=stage1, + stage2=stage2, + p2p_quant="none" if mtpr_class <= 1024 else "fp8_blockwise_1x32", + combine_blocks=combine[0], + combine_waves=combine[1], + ) + + @cache def _select_bucket_config( bucket: int, @@ -365,7 +515,15 @@ def _select_bucket_config( model_dim: int, inter_dim: int, fixed_slot_dispatch: bool, + world_size: int, + experts_per_rank: int, ) -> MegaMoEConfig: + if (model_dim, inter_dim, experts_per_rank * world_size) == ( + 5120, + 2304, + 384, + ) and world_size in (2, 4): + return _select_flash_config(bucket, mtpr_class, world_size, fixed_slot_dispatch) if mtpr_class == MAX_MTPR_CLASS: stage1 = _select_large_stage1(bucket, inter_dim) stage2 = _select_large_stage2(bucket, stage1.sort_block_m, model_dim) @@ -398,6 +556,12 @@ def select_mega_moe_config( raise ValueError(f"tokens={tokens} exceeds mtpr={mtpr}") if experts_per_rank <= 0: raise ValueError(f"experts_per_rank must be positive, got {experts_per_rank}") + if (model_dim, inter_dim, experts_per_rank * world_size) == ( + 5120, + 2304, + 384, + ) and world_size not in (2, 4): + raise ValueError("DeepSeek-V4.1-Flash MegaMoE requires EP2 or EP4") if not 0 < world_size <= 8: raise ValueError(f"world_size must be in [1, 8], got {world_size}") if model_dim <= 0 or inter_dim <= 0: @@ -407,15 +571,14 @@ def select_mega_moe_config( "MegaMoE v2 fanout pair ids support at most " f"{MAX_FANOUT_EXPERTS_PER_RANK} experts per rank" ) - bucket = nearest_token_bucket(tokens) - mtpr_class = mtpr_config_class(mtpr) - fixed_slot_dispatch = ( - mtpr_class <= FIXED_SLOT_MAX_MTPR - and world_size == 8 - and experts_per_rank == REFERENCE_EXPERTS_PER_RANK + buckets = ( + FLASH_TOKEN_BUCKETS + if (model_dim, inter_dim, experts_per_rank * world_size) == (5120, 2304, 384) + else TOKEN_BUCKETS ) - if fixed_slot_dispatch and bucket > FIXED_SLOT_MAX_MTPR: - raise ValueError(f"fixed-slot does not support token bucket {bucket}") + bucket = nearest_token_bucket(tokens, buckets) + mtpr_class = mtpr_config_class(mtpr) + fixed_slot_dispatch = _use_fixed_slots(mtpr_class, world_size, experts_per_rank) total_segments = world_size * experts_per_rank + world_size if total_segments > MAX_FANOUT_SEGMENTS: raise ValueError( @@ -428,6 +591,8 @@ def select_mega_moe_config( model_dim, inter_dim, fixed_slot_dispatch, + world_size, + experts_per_rank, ) @@ -443,15 +608,16 @@ def build_mega_moe_bundle_plan( """Deduplicate variants while keeping Stage1/Stage2 selection atomic.""" if mtpr <= 0 or mtpr & (mtpr - 1): raise ValueError(f"mtpr={mtpr} must be a positive power of two") - buckets = tuple(bucket for bucket in TOKEN_BUCKETS if bucket <= mtpr) + token_buckets = ( + FLASH_TOKEN_BUCKETS + if (model_dim, inter_dim, experts_per_rank * world_size) == (5120, 2304, 384) + else TOKEN_BUCKETS + ) + buckets = tuple(bucket for bucket in token_buckets if bucket <= mtpr) if not buckets or buckets[-1] != mtpr: raise ValueError(f"mtpr={mtpr} has no exact token bucket") - fixed_slot_dispatch = ( - mtpr <= FIXED_SLOT_MAX_MTPR - and world_size == 8 - and experts_per_rank == REFERENCE_EXPERTS_PER_RANK - ) + fixed_slot_dispatch = _use_fixed_slots(mtpr, world_size, experts_per_rank) stage1_variants: list[Stage1Config] = [] stage2_variants: list[Stage2BundleKey] = [] stage1_ids: dict[Stage1Config, int] = {} diff --git a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage1.py b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage1.py index 1277daa5109..60700839d73 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage1.py +++ b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage1.py @@ -59,8 +59,9 @@ def _validate_fixed_slot_geometry( if tile_m <= 0 or max_tokens_per_rank <= 0: raise ValueError("fixed-slot requires positive tile and token capacities") required_cap = ((npes * max_tokens_per_rank + tile_m - 1) // tile_m) * tile_m - if npes != 8 or experts_per_rank != 48 or cap != required_cap: - raise ValueError("fixed-slot dispatch requires the direct EP8/EPR48 layout") + supported = (npes, experts_per_rank) in ((8, 48), (2, 192), (4, 96)) + if not supported or cap != required_cap: + raise ValueError("unsupported fixed-slot EP geometry or capacity") def _validate_dispatch_capacity( @@ -100,6 +101,7 @@ def compile_mega_moe_stage1( work_shards: int | None = None, payload_chunk_rows: int = 0, tile_state_stride: int = 0, swiglu_limit: float = 0.0, + pair_k: bool = False, _return_kernel_spec: bool = False, ): arch = str(get_rocm_arch() or "") @@ -148,7 +150,7 @@ def compile_mega_moe_stage1( assert launch_grid_x <= num_cu * 33 + 1 M_REPEAT = sort_block_m // 16 NUM_ACC_N = n_per_wave // 16 - assert NUM_ACC_N % 2 == 0 and M_REPEAT % 2 == 0 + assert NUM_ACC_N % 2 == 0 and sort_block_m in (16, 32, 64, 128) TILE_K_BYTES = tile_k // 2 assert TILE_K_BYTES % 128 == 0 @@ -222,7 +224,9 @@ class SharedStorage: FAST_COPY = compact_dispatch and envs.AITER_MEGA_DISPATCH_FAST_COPY # Paired GEMM1 K loop for fixed-slot tiles with LDS-DMA A copies; compact # prefill is slower with it (its 8-wave tiles already hide the latency). - PAIR_K = fixed_slot_dispatch and async_a_copy and envs.AITER_MEGA_S1_FIXED_KPAIR + PAIR_K = async_a_copy and ( + pair_k or (fixed_slot_dispatch and envs.AITER_MEGA_S1_FIXED_KPAIR) + ) EPI_EVEC = envs.AITER_MEGA_S1_EPI_EVEC dispatch_path = "fixedslot" if fixed_slot_dispatch else "compact" swiglu_suffix = "" if swiglu_limit <= 0 else f"_sl{str(float(swiglu_limit)).replace('.', 'p')}" @@ -728,6 +732,7 @@ def compile_mega_moe_stage1_bundle( payload_chunk_rows=config.payload_chunk_rows, tile_state_stride=tile_state_stride, swiglu_limit=swiglu_limit, + pair_k=config.pair_k, _return_kernel_spec=True, ) for config in variants @@ -813,7 +818,7 @@ def run_mega_moe_stage1(out, x, w, scale_x, scale_w, sorted_token_ids, expert_id use_tile_resource=True, waves_per_eu_hint=2, b_nt=-1, work_shards=None, payload_chunk_rows=0, tile_state_stride=0, - swiglu_limit=0.0): + swiglu_limit=0.0, pair_k=False): launch = compile_mega_moe_stage1( model_dim=model_dim, inter_dim=inter_dim, rank=rank, experts_per_rank=experts_per_rank, fuse_npes=fuse_npes, fuse_topk=fuse_topk, fuse_cap=fuse_cap, fuse_mtpr=fuse_mtpr, @@ -825,6 +830,7 @@ def run_mega_moe_stage1(out, x, w, scale_x, scale_w, sorted_token_ids, expert_id b_nt=b_nt, work_shards=work_shards, payload_chunk_rows=payload_chunk_rows, tile_state_stride=tile_state_stride, swiglu_limit=swiglu_limit, + pair_k=pair_k, ) _run_compiled( launch, out, x, w, scale_x, scale_w, sorted_token_ids, expert_ids, num_valid_ids, out_scale, diff --git a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage2.py b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage2.py index d189254f22e..a2cb681a499 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage2.py +++ b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage2.py @@ -499,7 +499,7 @@ def run_unit(unit_bx, m_block_idx): arg_bscale, arg_eids, arg_aq, i32_max_m_blocks, unit_bx, lane, wave, i32_inter, i32_hidden, i32_kpad, i32_npad, BM=BM, BN=BN, BK=BK, use_nt=use_nt, INTER_MAX=INTER_MAX, aStages=aStages, a_dtype=a_dtype, has_pad=has_pad, SBM=SBM, g2_bhoist=g2_bhoist, g2_ascale_pf=g2_ascale_pf, - expert_offset=_expert_offset) + expert_offset=_expert_offset, dense_bm16_scales=True) p2p_scatter_epilog(lds_base_i32, accm_vecs, n_block_idx, wave, lane, N_OUT=N_OUT, BM=BM, BN=BN, npes=npes, topk=topk, log2_max_tok=log2_max_tok, mask_max_tok=mask_max_tok, recv_cap=_recv_cap, diff --git a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_v2.py b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_v2.py index 7bca56d9ff6..af76af86c31 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_v2.py +++ b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_v2.py @@ -33,6 +33,7 @@ class MegaMoEV2: # forward(mask_invalid_slots=...) is accepted: combine can skip -1 top-k slots. supports_combine_mask = True + supports_layer_weights = True # fmt: off def __init__(self, *, rank: int, world_size: int, model_dim: int, inter_dim: int, experts: int, topk: int, @@ -115,7 +116,7 @@ def _build_fused_stage1(self, w1, w1_scale): run_mega_moe_stage1_bundle, ) - self.sort_block_m = 32 + self.sort_block_m = min(entry.config.stage1.sort_block_m for entry in self._bundle_plan.entries) self._s1_w1 = w1.contiguous().view(torch.uint8) self._s1_w1_scale = w1_scale.contiguous().view(torch.uint8) op = self.comb_op._gm @@ -619,6 +620,7 @@ def _run_fused_stage1( mfma_amajor=config.mfma_amajor, swizzle_a=config.swizzle_a, async_a_copy=config.async_a_copy, + pair_k=config.pair_k, num_dispatch_cu=config.num_dispatch_cu, use_tile_resource=config.use_tile_resource, waves_per_eu_hint=config.waves_per_eu_hint, @@ -674,6 +676,11 @@ def _run_stage2( self, run_tokens, stream, slice_output, config: MegaMoEConfig, combine_ids=None ): """``combine_ids``: the top-k ids, when combine must skip their -1 slots.""" + self.comb_op.cfg = replace( + self.comb_cfg, + combine_block_num=config.combine_blocks, + combine_warp_num_per_block=config.combine_waves, + ) if config.stage2.aligned_pair: return self._run_aligned_pair_stage2( run_tokens, config, stream, slice_output, combine_ids @@ -697,11 +704,36 @@ def _stage2_output(self, result, run_tokens, slice_output): return out_tok[:run_tokens] if slice_output else out_tok def forward( - self, x_bf16, wts, topk_ids, *, stream=None, slice_output=True, mask_invalid_slots=None + self, x_bf16, wts, topk_ids, *, stream=None, slice_output=True, + mask_invalid_slots=None, w1=None, w1_scale=None, w2=None, w2_scale=None, + config_tokens=None, ): - """``mask_invalid_slots``: ids may contain -1 (skipped, contributing zero). - None defers to AITER_MEGA_COMBINE_MASK (default off).""" + """Run one collective MoE invocation with optional prepared layer weights. + + ``config_tokens`` must agree across ranks and cover their source counts. + Replacement weights use the constructor's shuffled MXFP4 layout. + ``mask_invalid_slots`` skips ids of -1; None uses AITER_MEGA_COMBINE_MASK. + """ run_tokens = int(x_bf16.shape[0]) + weights = (w1, w1_scale, w2, w2_scale) + if any(weight is not None for weight in weights): + if any(weight is None for weight in weights): + raise ValueError("per-layer weights require W13, W2 and both scales") + expected = ( + self.epr * 2 * self.inter_dim * self.model_dim // 2, + self.epr * 2 * self.inter_dim * self.model_dim // 32, + self.epr * self.model_dim * self.inter_dim // 2, + self.epr * self.model_dim * self.inter_dim // 32, + ) + for weight, size in zip(weights, expected): + if ( + weight.element_size() != 1 or weight.numel() != size + or weight.device != self.dev or not weight.is_contiguous() + ): + raise ValueError("per-layer weights must match the prepared geometry") + self._s1_w1, self._s1_w1_scale, self.w2, self.w2_scale = ( + weight.view(torch.uint8) for weight in weights + ) if run_tokens > self.mtpr: raise ValueError(f"run_tokens={run_tokens} > max_tok_per_rank={self.mtpr}") if x_bf16.dtype != torch.bfloat16 or not x_bf16.is_contiguous(): @@ -710,14 +742,16 @@ def forward( raise ValueError("wts must be contiguous float32") if topk_ids.dtype != torch.int32 or not topk_ids.is_contiguous(): raise ValueError("topk_ids must be contiguous int32") + policy_tokens = run_tokens if config_tokens is None else int(config_tokens) + if policy_tokens < run_tokens or policy_tokens > self.mtpr: + raise ValueError("config_tokens must cover every source rank within MTPR") + config = self._select_config(policy_tokens) if self._s1_fixed_slot: x_q, scales = self.quantize(x_bf16) return self._run_joint( x_q, scales, wts, topk_ids, run_tokens, stream, slice_output, - mask_invalid_slots=mask_invalid_slots, + config=config, mask_invalid_slots=mask_invalid_slots, ) - - config = self._select_config(run_tokens) prepare_stream = stream if prepare_stream is None: prepare_stream = fx.Stream(torch.cuda.current_stream().cuda_stream) @@ -934,6 +968,11 @@ def preload_stage2_bundle(self): ) self._preload_aligned_pair_stage2(config, stream) for entry in self._bundle_plan.entries: + self.comb_op.cfg = replace( + self.comb_cfg, + combine_block_num=entry.config.combine_blocks, + combine_warp_num_per_block=entry.config.combine_waves, + ) for mask_topk_ids in (False, True): self.comb_op.preload_combine_no_stage1( self._g2_combine_placeholder, diff --git a/op_tests/multigpu_tests/bench_mega_moe_v2.py b/op_tests/multigpu_tests/bench_mega_moe_v2.py index 25f72dd3811..41e4da79147 100644 --- a/op_tests/multigpu_tests/bench_mega_moe_v2.py +++ b/op_tests/multigpu_tests/bench_mega_moe_v2.py @@ -6,6 +6,7 @@ import argparse import os +import statistics from dataclasses import replace from pathlib import Path @@ -101,7 +102,10 @@ def make_inputs(tokens, rank, world, model_dim, experts, topk, route, hot_bias, destination_scores = torch.rand( (tokens, world), device=device, generator=generator ) - destination = torch.topk(destination_scores, topk, dim=-1).indices + destination_order = destination_scores.argsort(dim=-1, descending=True) + slots = torch.arange(topk, device=device) + destination = destination_order[:, slots % world] + local_slot = slots // world if route == "rank-balanced-last": hot = torch.ones_like(destination, dtype=torch.bool) elif route == "rank-mixed-skew": @@ -114,7 +118,12 @@ def make_inputs(tokens, rank, world, model_dim, experts, topk, route, hot_bias, cold_expert = torch.randint( 1, local_experts, (tokens, topk), device=device, generator=generator ) - hot_expert = local_experts - 1 if route == "rank-balanced-last" else 0 + hot_expert = ( + local_experts - 1 - local_slot + if route == "rank-balanced-last" + else local_slot + ) + cold_expert = cold_expert // topk * topk + local_slot ids = destination * local_experts + torch.where(hot, hot_expert, cold_expert) values = torch.randn( (tokens, topk), dtype=torch.float32, device=device, generator=generator @@ -169,31 +178,93 @@ def capture(body): for _ in range(5): graph.replay() barrier() + graph._benchmark_body = body return graph -def time_graph(graph, iters, device): - barrier() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - for _ in range(iters): - graph.replay() - end.record() - torch.cuda.synchronize() - local_ms = start.elapsed_time(end) / iters - mean = torch.tensor(local_ms, dtype=torch.float64, device=device) +def run_online_mori( + op, + x, + weights, + ids, + expert_mask, + layer_weights, + block_num, + warp_per_block, + trim_rows, +): + w1, w1_scale, w2, w2_scale = layer_weights + dispatched, recv_weights, recv_scales, recv_ids, recv_tokens = op.dispatch( + x, + weights, + None, + ids, + block_num=block_num, + warp_per_block=warp_per_block, + ) + dispatched, recv_weights, recv_scales, recv_ids = trim_mori_dispatch_output( + dispatched, + recv_weights, + recv_scales, + recv_ids, + trim_rows, + ) + local_out = fused_moe( + dispatched, + w1, + w2, + recv_weights, + recv_ids, + expert_mask, + quant_type=aiter.QuantType.per_1x32, + num_local_tokens=recv_tokens, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=None, + dtype=torch.bfloat16, + swiglu_limit=SWIGLU_LIMIT, + gate_mode=GateMode.INTERLEAVE.value, + ) + return op.combine( + local_out, None, ids, block_num=block_num, warp_per_block=warp_per_block + )[0][: x.shape[0]] + + +def time_graph(graph, iters, device, cache_graph=None): + start = torch.cuda.Event(enable_timing=True, external=True) + end = torch.cuda.Event(enable_timing=True, external=True) + measured = torch.cuda.CUDAGraph() + with torch.cuda.graph(measured, stream=torch.cuda.Stream()): + ms.shmem_barrier_on_stream(torch.cuda.current_stream()) + start.record() + graph._benchmark_body() + end.record() + samples = [] + for iteration in range(iters + 5): + if cache_graph is not None: + cache_graph.replay() + measured.replay() + torch.cuda.synchronize() + if iteration >= 5: + samples.append(start.elapsed_time(end)) + mean = torch.tensor(samples, dtype=torch.float64, device=device) maximum = mean.clone() dist.all_reduce(mean, op=dist.ReduceOp.SUM) dist.all_reduce(maximum, op=dist.ReduceOp.MAX) - return float(mean.item() / dist.get_world_size()), float(maximum.item()) + return ( + statistics.median(mean.tolist()) / dist.get_world_size(), + statistics.median(maximum.tolist()), + ) -def profile_graph(graph, name, rank, out_dir, replays=3): +def profile_graph(graph, name, rank, out_dir, replays=3, cache_graph=None): barrier() with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: dist.barrier() for _ in range(replays): + if cache_graph is not None: + cache_graph.replay() + ms.shmem_barrier_on_stream(torch.cuda.current_stream()) graph.replay() torch.cuda.synchronize() path = Path(out_dir) @@ -204,6 +275,7 @@ def profile_graph(graph, name, rank, out_dir, replays=3): def main(): parser = argparse.ArgumentParser() + parser.add_argument("--network", choices=("v4_pro", "v4_1_flash"), default="v4_pro") parser.add_argument("--tokens", type=int, default=8192) parser.add_argument("--rank-tokens", default="") parser.add_argument("--config-tokens", type=int, default=0) @@ -231,6 +303,7 @@ def main(): parser.add_argument("--experts", type=int, default=EXPERTS) parser.add_argument("--topk", type=int, default=TOPK) parser.add_argument("--iters", type=int, default=10) + parser.add_argument("--cold-cache-mb", type=int, default=512) parser.add_argument( "--route", choices=( @@ -252,12 +325,19 @@ def main(): parser.add_argument("--stage1-dispatch-cu", type=int, default=0) parser.add_argument("--stage1-grid-mult", type=int, default=0) parser.add_argument("--stage1-b-nt", type=int, default=-1) + parser.add_argument("--stage1-tile-n", type=int, default=0) + parser.add_argument("--stage1-sort-block-m", type=int, default=0) + parser.add_argument("--stage1-num-waves", type=int, default=0) + parser.add_argument("--stage2-block-n", type=int, default=0) + parser.add_argument("--stage2-block-m", type=int, default=0) parser.add_argument("--stage1-tile-resource", action="store_true") parser.add_argument("--check-variant", action="store_true") parser.add_argument("--profile-dir", default="") parser.add_argument("--mega-only", action="store_true") parser.add_argument("--perf-guard", action="store_true") args = parser.parse_args() + if args.network == "v4_1_flash": + args.model_dim, args.inter_dim, args.experts, args.topk = 5120, 2304, 384, 6 if not args.mega_only and args.phase is None: parser.error("--phase {decode,prefill} is required for a Mori comparison") @@ -274,8 +354,8 @@ def main(): os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0" rank, world, device = setup_dist() - if world != 8: - raise ValueError("This comparison requires eight ranks") + if world not in ((2, 4) if args.network == "v4_1_flash" else (8,)): + raise ValueError("Flash requires EP2/EP4; V4-Pro comparison requires EP8") if args.experts % world: raise ValueError(f"experts={args.experts} must be divisible by world={world}") rank_tokens = [int(value) for value in args.rank_tokens.split(",") if value] @@ -323,6 +403,12 @@ def main(): torch.ones_like(ids.flatten(), dtype=torch.int64), ) dist.all_reduce(route_counts, op=dist.ReduceOp.SUM) + received_tokens = ( + (ids[:, :, None] // local_experts == torch.arange(world, device=device)) + .any(dim=1) + .sum(dim=0) + ) + dist.all_reduce(received_tokens, op=dist.ReduceOp.SUM) expert_counts = torch.bincount( ids.flatten().to(torch.int64), minlength=args.experts ) @@ -358,6 +444,11 @@ def main(): or args.stage1_dispatch_cu or args.stage1_grid_mult or args.stage1_b_nt >= 0 + or args.stage1_tile_n + or args.stage1_sort_block_m + or args.stage1_num_waves + or args.stage2_block_n + or args.stage2_block_m or args.stage1_tile_resource or args.config_tokens ): @@ -377,6 +468,12 @@ def select_strided_config(tokens): stage1_updates["grid_mult"] = args.stage1_grid_mult if args.stage1_b_nt >= 0: stage1_updates["b_nt"] = args.stage1_b_nt + if args.stage1_tile_n: + stage1_updates["tile_n"] = args.stage1_tile_n + if args.stage1_sort_block_m: + stage1_updates["sort_block_m"] = args.stage1_sort_block_m + if args.stage1_num_waves: + stage1_updates["num_waves"] = args.stage1_num_waves if args.stage1_tile_resource: stage1_updates["use_tile_resource"] = True if stage1_updates: @@ -386,9 +483,13 @@ def select_strided_config(tokens): or args.stage2_persist_cu or args.stage2_skew_cu or args.disable_stage2_skew + or args.stage2_block_n + or args.stage2_block_m ): stage2 = replace( stage2, + block_n=args.stage2_block_n or stage2.block_n, + block_m=args.stage2_block_m or stage2.block_m, persist_strided=args.stage2_strided, persist_cu=args.stage2_persist_cu or stage2.persist_cu, skew_cu=( @@ -433,54 +534,51 @@ def select_strided_config(tokens): holders = {} def mori_body(): - dispatched, recv_weights, recv_scales, recv_ids, recv_tokens = mori_op.dispatch( + holders["mori"] = run_online_mori( + mori_op, x, route_weights, - None, ids, - block_num=mori_block_num, - warp_per_block=mori_warp_per_block, - ) - dispatched, recv_weights, recv_scales, recv_ids = trim_mori_dispatch_output( - dispatched, - recv_weights, - recv_scales, - recv_ids, + expert_mask, + (w1, w1_scale, w2, w2_scale), + mori_block_num, + mori_warp_per_block, mori_trim_rows, ) - local_out = fused_moe( - dispatched, - w1, - w2, - recv_weights, - recv_ids, - expert_mask, - quant_type=aiter.QuantType.per_1x32, - num_local_tokens=recv_tokens, + + def mega_body(): + holders["mega"] = mega( + x, + route_weights, + ids, + config_tokens=max_rank_tokens, + w1=w1, w1_scale=w1_scale, + w2=w2, w2_scale=w2_scale, - a1_scale=None, - dtype=torch.bfloat16, - swiglu_limit=SWIGLU_LIMIT, - gate_mode=GateMode.INTERLEAVE.value, ) - holders["mori"] = mori_op.combine( - local_out, - None, - ids, - block_num=mori_block_num, - warp_per_block=mori_warp_per_block, - )[0][:tokens] - def mega_body(): - holders["mega"] = mega(x, route_weights, ids) + if not args.mega_only: + mori_body() + baseline_output = holders["mori"].clone() + barrier() + mega_body() + error_sq = (holders["mega"].float() - baseline_output.float()).square().sum() + norm_sq = baseline_output.float().square().sum() + dist.all_reduce(error_sq) + dist.all_reduce(norm_sq) + baseline_rel_l2 = float((error_sq / norm_sq).sqrt()) + if not baseline_rel_l2 < 0.1: + raise AssertionError(f"MegaMoE vs baseline relL2={baseline_rel_l2}") + if rank == 0: + print(f"[ACCURACY] baseline_rel_l2={baseline_rel_l2:.6f}", flush=True) mori_graph = None if args.mega_only else capture(mori_body) print(f"[STEP] rank={rank} mori-capture-done", flush=True) if rank == 0 and mori_graph is not None: full_recv_rows = mori_op.max_num_tokens_to_recv() fused_moe_rows = min(full_recv_rows, mori_trim_rows or full_recv_rows) - tune_topk = args.topk - 1 # fused_moe's expert-parallel lookup convention + tune_topk = args.topk print( "[MORI_CONFIG] " f"mode=MANUAL phase={phase} graph_bs={graph_bs} dp_size={world} " @@ -497,16 +595,22 @@ def mega_body(): print(f"[STEP] rank={rank} mega-capture-done", flush=True) if rank == 0: print(f"[MEGA_CONFIG] {mega._active_config}", flush=True) + cache_storage = torch.zeros( + args.cold_cache_mb * 1024 * 1024, dtype=torch.uint8, device=device + ) + cache_graph = capture(lambda: cache_storage.add_(1)) if args.cold_cache_mb else None mori_ms = ( (float("nan"), float("nan")) if mori_graph is None - else time_graph(mori_graph, args.iters, device) + else time_graph(mori_graph, args.iters, device, cache_graph) ) - mega_ms = time_graph(mega_graph, args.iters, device) + mega_ms = time_graph(mega_graph, args.iters, device, cache_graph) x_q, x_scale = mega.quantize(x) def mega_stage1(): - mega._run_fused_stage1(x_q, route_weights, x_scale, ids) + mega._run_fused_stage1( + x_q, route_weights, x_scale, ids, config=mega._active_config.stage1 + ) stage1_graph = capture(mega_stage1) print(f"[STEP] rank={rank} stage1-capture-done", flush=True) @@ -518,10 +622,10 @@ def mega_stage2(): stage2_graph = capture(mega_stage2) print(f"[STEP] rank={rank} stage2-capture-done", flush=True) - stage1_ms = time_graph(stage1_graph, args.iters, device) + stage1_ms = time_graph(stage1_graph, args.iters, device, cache_graph) mega_stage1() barrier() - stage2_ms = time_graph(stage2_graph, args.iters, device) + stage2_ms = time_graph(stage2_graph, args.iters, device, cache_graph) rel_l2 = None if args.check_variant: @@ -537,6 +641,8 @@ def mega_stage2(): candidate.float() - reference.float() ).norm() / reference.float().norm() dist.all_reduce(rel_l2, op=dist.ReduceOp.MAX) + if not float(rel_l2) < 0.001: + raise AssertionError(f"variant vs default relL2={float(rel_l2)}") if args.profile_dir: if mori_graph is not None: @@ -571,6 +677,10 @@ def mega_stage2(): guard_pass = guard_floor is None or speedup >= guard_floor if rank == 0: print(f"[ROUTES] per-destination-rank={route_counts.tolist()}", flush=True) + print( + f"[RECEIVED] unique_tokens={received_tokens.tolist()} max_tokens={int(received_tokens.max())}", + flush=True, + ) print( f"[EXPERTS] active={(expert_counts > 0).sum().item()} max_routes={expert_counts.max().item()} " f"mean_routes={expert_counts.float().mean().item():.1f}", @@ -588,7 +698,8 @@ def mega_stage2(): f"mori_e2e={mori_ms[0]:.4f}/{mori_ms[1]:.4f}ms " f"mega_e2e={mega_ms[0]:.4f}/{mega_ms[1]:.4f}ms speedup={speedup:.2f}% " f"stage1={stage1_ms[0]:.4f}/{stage1_ms[1]:.4f}ms " - f"stage2_combine={stage2_ms[0]:.4f}/{stage2_ms[1]:.4f}ms rank-mean/max", + f"stage2_combine={stage2_ms[0]:.4f}/{stage2_ms[1]:.4f}ms " + f"cold_cache_mb={args.cold_cache_mb} median-rank-mean/max", flush=True, ) if guard_floor is not None: diff --git a/op_tests/multigpu_tests/test_mega_moe_v2.py b/op_tests/multigpu_tests/test_mega_moe_v2.py index 22321c256d1..f685465f846 100644 --- a/op_tests/multigpu_tests/test_mega_moe_v2.py +++ b/op_tests/multigpu_tests/test_mega_moe_v2.py @@ -22,6 +22,13 @@ from aiter.utility import fp4_utils NETWORKS = { + "v4_1_flash": { + "model_dim": 5120, + "inter_dim": 2304, + "experts": 384, + "topk": 6, + "swiglu_limit": 10.0, + }, # swiglu_limit 0.0 exercises the operator's default, unclamped SwiGLU. "glm52": { "model_dim": 6144, @@ -111,6 +118,7 @@ def _make_inputs( force_fanout_boundary, inject_invalid_route, force_padding_boundary, + route="uniform", ): generator = torch.Generator(device=device).manual_seed(seed + rank) x = torch.randn( @@ -120,6 +128,19 @@ def _make_inputs( (tokens, experts), dtype=torch.float32, device=device, generator=generator ) values, ids = torch.topk(scores, topk, dim=-1) + if route != "uniform": + world = dist.get_world_size() + epr = experts // world + slot = torch.arange(topk, device=device)[None, :] + token = torch.arange(tokens, device=device)[:, None] + if route == "expert-boundary": + destination = (token + rank + slot) % world + local = torch.where(slot % 2 == 0, slot // world, epr - 1 - slot // world) + ids = destination * epr + local + elif route == "hot-rank0": + ids = (slot + epr - topk).expand(tokens, -1).clone() + else: + raise ValueError(f"unsupported route {route}") if force_fanout_boundary or inject_invalid_route or force_padding_boundary: if topk != 16 or experts % dist.get_world_size(): raise ValueError("fanout adversarial cases require topk=16 and EP") @@ -285,7 +306,7 @@ def _reference( return partial[start : start + x.shape[0]] -def _time_graph(fn, device, iters): +def _time_graph(fn, device, iters, prepare=None): _barrier() fn() _barrier() @@ -294,18 +315,32 @@ def _time_graph(fn, device, iters): with torch.cuda.graph(graph, stream=capture_stream): fn() for _ in range(10): + if prepare is not None: + prepare() graph.replay() torch.cuda.synchronize() start, end = ( torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True), ) - start.record() - for _ in range(iters): - graph.replay() - end.record() + if prepare is None: + start.record() + for _ in range(iters): + graph.replay() + end.record() + events = [(start, end)] + else: + events = [ + (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) + for _ in range(iters) + ] + for start, end in events: + prepare() + start.record() + graph.replay() + end.record() torch.cuda.synchronize() - local_ms = start.elapsed_time(end) / iters + local_ms = sum(start.elapsed_time(end) for start, end in events) / iters mean_ms = _reduce_float(local_ms, device, dist.ReduceOp.SUM) / dist.get_world_size() max_ms = _reduce_float(local_ms, device, dist.ReduceOp.MAX) return mean_ms, max_ms @@ -365,7 +400,9 @@ def end_to_end(): stage1_ms = _time_graph(stage1, device, args.iters) stage1() _barrier() - stage2_ms = _time_graph(stage2, device, args.iters) + # Stage1's peer handshake protects the previous combine's consumers before + # the next W2 overwrites their P2P inputs. Keep it outside the timed region. + stage2_ms = _time_graph(stage2, device, args.iters, prepare=stage1) e2e_ms = _time_graph(end_to_end, device, args.iters) graph_output = state["output"][:tokens] graph_replay_exact = torch.tensor( @@ -447,14 +484,26 @@ def capture(): ) -def _run_burst(moe, x, weights, ids, depth, rank): - moe(x, weights, ids) +def _run_burst(moe, x, weights, ids, depth, rank, graph_mode=False): + expected = moe(x, weights, ids).clone() _barrier() + graph = None + if graph_mode: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=torch.cuda.Stream()): + output = moe(x, weights, ids) + _barrier() for _ in range(depth): - moe(x, weights, ids) + if graph is None: + output = moe(x, weights, ids) + else: + graph.replay() torch.cuda.synchronize() + if not _reduce_float(torch.equal(output, expected), x.device, dist.ReduceOp.MIN): + raise AssertionError("burst execution changed the output") if rank == 0: - print(f"[BURST] completed={depth}/{depth}", flush=True) + mode = "graph" if graph_mode else "eager" + print(f"[BURST] mode={mode} completed={depth}/{depth} output=PASS", flush=True) def _install_config_policy(moe, config_tokens, unify_fields): @@ -503,7 +552,12 @@ def select_config(tokens): def main(): parser = argparse.ArgumentParser() parser.add_argument("--network", choices=NETWORKS, default="v4_pro") - parser.add_argument("--bs-list", default="128") + parser.add_argument("--bs-list") + parser.add_argument( + "--route", + choices=("uniform", "expert-boundary", "hot-rank0"), + default="uniform", + ) parser.add_argument("--iters", type=int, default=30) parser.add_argument("--seed", type=int, default=123) parser.add_argument("--accuracy-max-bs", type=int, default=128) @@ -513,6 +567,7 @@ def main(): parser.add_argument("--config-tokens", type=int, default=0) parser.add_argument("--unify-fields", default="") parser.add_argument("--burst-depth", type=int, default=0) + parser.add_argument("--burst-graph", action="store_true") parser.add_argument("--force-fanout-boundary", action="store_true") parser.add_argument("--inject-invalid-route", action="store_true") parser.add_argument("--force-padding-boundary", action="store_true") @@ -524,6 +579,14 @@ def main(): "forward(mask_invalid_slots=True) returns them as zeros", ) args = parser.parse_args() + if args.bs_list is None: + sizes = list(range(1, 65)) + [96, 128, 192, 256, 512, 1024] + args.bs_list = ( + ",".join(map(str, sizes)) if args.network == "v4_1_flash" else "128" + ) + if args.network == "v4_1_flash": + args.accuracy_max_bs = max(args.accuracy_max_bs, 1024) + args.max_tok_per_rank = args.max_tok_per_rank or 1024 if ( sum( ( @@ -543,6 +606,8 @@ def main(): rank, world, device = _setup_dist() try: network = NETWORKS[args.network] + if args.network == "v4_1_flash" and world not in (2, 4): + raise ValueError("DeepSeek-V4.1-Flash MegaMoE requires EP2 or EP4") if network["experts"] % world: raise ValueError( f"experts={network['experts']} must be divisible by world={world}" @@ -584,6 +649,7 @@ def main(): force_fanout_boundary=args.force_fanout_boundary, inject_invalid_route=args.inject_invalid_route, force_padding_boundary=args.force_padding_boundary, + route=args.route, ) ref_weights = w1_q, w1_ref_scale, w2_q, w2_ref_scale # MegaMoE's producer/consumer wire geometry is rank-invariant. A @@ -624,6 +690,44 @@ def main(): **network, ) moe = shared_moe + if args.network == "v4_1_flash" and batch_size == batch_sizes[0]: + local_x = x[:local_batch_size].contiguous() + local_weights = weights[:local_batch_size].contiguous() + local_ids = ids[:local_batch_size].contiguous() + expected = moe( + local_x, + local_weights, + local_ids, + config_tokens=config_tokens or None, + ).clone() + zero_w2 = torch.zeros_like(w2.view(torch.uint8)).view(w2.dtype) + zero_output = moe( + local_x, + local_weights, + local_ids, + w1=w1, + w1_scale=w1_scale, + w2=zero_w2, + w2_scale=w2_scale, + config_tokens=config_tokens or None, + ) + if not torch.equal(zero_output, torch.zeros_like(zero_output)): + raise AssertionError( + "per-layer W2 replacement did not produce zeros" + ) + restored = moe( + local_x, + local_weights, + local_ids, + w1=w1, + w1_scale=w1_scale, + w2=w2, + w2_scale=w2_scale, + config_tokens=config_tokens or None, + ) + if not torch.equal(restored, expected): + raise AssertionError("per-layer weight restore changed the output") + del zero_w2 _install_config_policy(moe, config_tokens, args.unify_fields) if rank_tokens: selected = moe._select_config(local_batch_size) @@ -646,7 +750,13 @@ def main(): local_ids = ids[:local_batch_size].contiguous() if args.burst_depth: _run_burst( - moe, local_x, local_weights, local_ids, args.burst_depth, rank + moe, + local_x, + local_weights, + local_ids, + args.burst_depth, + rank, + args.burst_graph, ) else: _run_size( diff --git a/op_tests/multigpu_tests/tune_mega_moe_v2.py b/op_tests/multigpu_tests/tune_mega_moe_v2.py new file mode 100644 index 00000000000..f52e6bfb3b6 --- /dev/null +++ b/op_tests/multigpu_tests/tune_mega_moe_v2.py @@ -0,0 +1,516 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""Layout and scheduling ablations for Flash EP2/EP4, with cold-cache timing.""" + +import argparse +import json +import os +import random +import re +from dataclasses import asdict, replace +from pathlib import Path + +from bench_mega_moe_v2 import ( + MegaMoEV2, + barrier, + capture, + dist, + make_inputs, + make_weights, + mori, + ms, + profile_graph, + run_online_mori, + setup_dist, + time_graph, + torch, +) + + +def layouts(base): + result = [base] + for bm, bn, waves in ( + (16, 128, 2), + (16, 256, 2), + (16, 256, 4), + (16, 512, 4), + (32, 128, 4), + (32, 64, 2), + (32, 128, 2), + (32, 256, 4), + (32, 256, 8), + (32, 512, 8), + (64, 256, 4), + (64, 128, 2), + (64, 128, 4), + (64, 256, 8), + (64, 512, 8), + (128, 256, 8), + (128, 128, 4), + ): + result.append( + replace( + base, + stage1=replace( + base.stage1, + sort_block_m=bm, + tile_n=bn, + num_waves=waves, + mfma_amajor=bm >= 32, + ), + stage2=replace(base.stage2, block_m=min(bm, 64)), + ) + ) + return list(dict.fromkeys(result)) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", default="1,8,32,64,128,256,512,1024") + parser.add_argument("--mtpr", type=int, default=1024) + parser.add_argument("--iters", type=int, default=15) + parser.add_argument("--route", choices=("uniform", "hot-rank0"), default="uniform") + parser.add_argument("--output", required=True) + parser.add_argument("--schedule", action="store_true") + parser.add_argument("--paired-k", action="store_true") + parser.add_argument("--baseline", action="store_true") + parser.add_argument("--seed-config-json") + parser.add_argument("--kernel-knobs", action="store_true") + parser.add_argument("--transport", action="store_true") + parser.add_argument("--selected-only", action="store_true") + parser.add_argument("--profile-dir", default="") + args = parser.parse_args() + seed_records = ( + json.loads(Path(args.seed_config_json).read_text()) + if args.seed_config_json + else [] + ) + os.environ["MORI_EP_LAUNCH_CONFIG_MODE"] = "MANUAL" + os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0" + rank, world, device = setup_dist() + if world not in (2, 4): + raise ValueError("Flash requires EP2 or EP4") + w1, s1, w2, s2 = make_weights(384 // world, 5120, 2304, rank, device) + mega = MegaMoEV2( + rank=rank, + world_size=world, + model_dim=5120, + inter_dim=2304, + experts=384, + topk=6, + quant="a8w4", + w1=w1, + w1_scale=s1, + w2=w2, + w2_scale=s2, + max_tok_per_rank=args.mtpr, + swiglu_limit=10.0, + ) + default_select = mega._select_config + mori_op = None + expert_mask = torch.zeros(385, dtype=torch.int32, device=device) + expert_mask[rank * mega.epr : (rank + 1) * mega.epr] = 1 + if args.baseline: + mori_op = mori.ops.EpDispatchCombineOp( + mori.ops.EpDispatchCombineConfig( + data_type=torch.bfloat16, + rank=rank, + world_size=world, + hidden_dim=5120, + scale_dim=0, + scale_type_size=4, + max_token_type_size=2, + max_num_inp_token_per_rank=args.mtpr, + num_experts_per_rank=mega.epr, + num_experts_per_token=6, + warp_num_per_block=16, + block_num=80, + rdma_block_num=0, + kernel_type=mori.ops.EpDispatchCombineKernelType.IntraNode, + gpu_per_node=world, + ) + ) + cache = torch.zeros(512 * 1024 * 1024, dtype=torch.uint8, device=device) + cache_graph = capture(lambda: cache.add_(1)) + records = [] + + def tune_case(tokens): + x, weights, ids = make_inputs( + tokens, rank, world, 5120, 384, 6, args.route, 8.0, device + ) + counts = torch.bincount(ids.flatten().long(), minlength=384) + dist.all_reduce(counts) + local_counts = counts[rank * mega.epr : (rank + 1) * mega.epr] + received = ( + (ids[:, :, None] // mega.epr == torch.arange(world, device=device)) + .any(1) + .sum(0) + ) + dist.all_reduce(received) + mega._select_config = default_select + reference = mega(x, weights, ids, config_tokens=tokens).clone() + barrier() + base = default_select(tokens) + if seed_records: + from aiter.ops.flydsl.kernels.mega_moe.mega_moe_config import ( + MegaMoEConfig, + Stage1Config, + Stage2Config, + ) + + seed = min( + (row for row in seed_records if row["source_m"] == tokens), + key=lambda row: row["median_max_rank_ms"], + )["config"] + base = MegaMoEConfig( + stage1=Stage1Config(**seed["stage1"]), + stage2=Stage2Config(**seed["stage2"]), + p2p_quant=seed["p2p_quant"], + combine_blocks=seed.get("combine_blocks", 128), + combine_waves=seed.get("combine_waves", 8), + ) + baseline_ms = None + if mori_op is not None: + + def baseline_body(): + return run_online_mori( + mori_op, + x, + weights, + ids, + expert_mask, + (w1, s1, w2, s2), + 64, + 4, + tokens * world, + ) + + baseline_output = baseline_body().clone() + error = (reference.float() - baseline_output.float()).square().sum() + norm = baseline_output.float().square().sum() + dist.all_reduce(error) + dist.all_reduce(norm) + if float((error / norm).sqrt()) >= 0.1: + raise AssertionError("MegaMoE vs MoRI+AITER exceeds relL2=0.1") + baseline_graph = capture(baseline_body) + baseline_ms = time_graph( + baseline_graph, args.iters * 2, device, cache_graph + )[1] + if args.paired_k: + base = replace(base, stage1=replace(base.stage1, pair_k=True)) + results = [] + active = base + + def select(_): + mega._active_config = active + return active + + mega._select_config = select + + def body(): + return mega(x, weights, ids, config_tokens=tokens) + + def evaluate(configs, phase): + nonlocal active + configs = list(dict.fromkeys(configs)) + random.Random(tokens + len(results)).shuffle(configs) + graphs = [] + for config in configs: + active = config + candidate = body() + error = (candidate.float() - reference.float()).square().sum() + norm = reference.float().square().sum() + dist.all_reduce(error) + dist.all_reduce(norm) + rel_l2 = float((error / norm).sqrt()) + if rel_l2 >= 0.001: + raise AssertionError( + f"EP{world} M{tokens} {config}: relL2={rel_l2}" + ) + graphs.append((config, capture(body), rel_l2)) + # Reverse the order on the second pass to share clock/thermal drift. + timings = {config: [] for config, _, _ in graphs} + for order in (graphs, list(reversed(graphs))): + for config, graph, _ in order: + active = config + timings[config].append( + time_graph(graph, args.iters, device, cache_graph)[1] + ) + for config, _, rel_l2 in graphs: + bm = config.stage1.sort_block_m + tiles = (local_counts + bm - 1) // bm + weight_bytes = float(tiles.sum()) * 3 * 5120 * 2304 * (0.5 + 1 / 32) + lower_bytes = ( + float((local_counts > 0).sum()) * 3 * 5120 * 2304 * (0.5 + 1 / 32) + ) + ms_max = sum(timings[config]) / len(timings[config]) + record = { + "ep": world, + "source_m": tokens, + "mtpr": args.mtpr, + "fixed": mega._s1_fixed_slot, + "route": args.route, + "received_unique": received.tolist(), + "phase": phase, + "config": asdict(config), + "rel_l2": rel_l2, + "median_max_rank_ms": ms_max, + "passes": timings[config], + "baseline_ms": baseline_ms, + "rank0_row_utilization": float(local_counts.sum()) + / max(1, int(tiles.sum()) * bm), + "rank0_ideal_weight_bytes": lower_bytes, + "rank0_tile_weight_bytes": weight_bytes, + "rank0_ideal_weight_gbps": lower_bytes / (ms_max * 1e6), + } + results.append((ms_max, config)) + if rank == 0: + records.append(record) + Path(args.output).write_text(json.dumps(records, indent=2) + "\n") + print( + f"[TUNE] EP{world} M={tokens} {phase} {ms_max * 1000:.1f}us " + f"W13={bm}x{config.stage1.tile_n}/{config.stage1.num_waves}w " + f"W2={config.stage2.block_m}x{config.stage2.block_n} " + f"rows={record['rank0_row_utilization']:.3f} relL2={rel_l2:.6f}", + flush=True, + ) + + evaluate( + [base] if args.transport or args.selected_only else layouts(base), "layout" + ) + if args.selected_only: + if args.profile_dir: + active = base + profile_graph( + capture(body), + f"mega_m{tokens}", + rank, + args.profile_dir, + cache_graph=cache_graph, + ) + barrier() + return + if args.transport: + evaluate( + [ + replace( + base, + stage1=replace( + base.stage1, payload_chunk_rows=chunk, prepare_quant_cu=cu + ), + ) + for chunk in (128, 256, 512, 1024) + for cu in (32, 64, 128) + ], + "prepare_payload", + ) + best = min(results, key=lambda item: item[0])[1] + evaluate( + [ + replace(best, stage2=replace(best.stage2, block_m=bm, block_n=bn)) + for bm in (16, 32, 64) + if bm <= best.stage1.sort_block_m + for bn in (128, 256, 512) + ], + "W2_rows", + ) + best = min(results, key=lambda item: item[0])[1] + evaluate( + [ + replace( + best, + stage2=replace( + best.stage2, + block_n=n, + persist_cu=cu, + use_nt=nt, + persist_strided=strided, + ), + ) + for n, cu, nt, strided in ( + (128, 256, True, False), + (256, 128, True, False), + (256, 192, True, False), + (256, 256, False, False), + (256, 256, True, True), + (512, 256, True, False), + (512, 128, True, False), + ) + ], + "W2", + ) + if best.stage1.sort_block_m >= 64: + evaluate( + [ + replace(best, stage2=replace(best.stage2, block_m=bm, block_n=bn)) + for bm in (32, 64) + for bn in (128, 256) + ], + "W2_rows", + ) + if args.schedule: + best = min(results, key=lambda item: item[0])[1] + dispatch = (4, 8, 16, 32, 64, 128) + evaluate( + [ + replace(best, stage1=replace(best.stage1, num_dispatch_cu=d)) + for d in dispatch + if d % world == 0 + ], + "dispatch", + ) + best = min(results, key=lambda item: item[0])[1] + evaluate( + [ + replace( + best, + stage1=replace( + best.stage1, + grid_mult=g, + work_shards=s, + b_nt=nt, + waves_per_eu_hint=wpe, + ), + ) + for g, s, nt, wpe in ( + (1, 1, 3, 2), + (1, 2, 3, 2), + (1, 8, 3, 2), + (2, 4, 3, 2), + (3, 4, 3, 2), + (1, 4, 0, 2), + (1, 4, 3, 1), + (1, 4, 3, 4), + ) + ], + "W13_schedule", + ) + if args.kernel_knobs: + best = min(results, key=lambda item: item[0])[1] + evaluate( + [ + replace(best, stage1=replace(best.stage1, **updates)) + for updates in ( + {"mfma_amajor": not best.stage1.mfma_amajor}, + {"async_a_copy": False}, + {"pipe_weights": False}, + {"swizzle_a": False}, + ) + ], + "W13_pipeline", + ) + best = min(results, key=lambda item: item[0])[1] + evaluate( + [ + replace(best, stage2=replace(best.stage2, **updates)) + for updates in ( + {"b_hoist": False}, + {"ascale_prefetch": False}, + {"spatial_partition": 0}, + {"spatial_partition": 102}, + {"spatial_partition": 202}, + {"spatial_partition": 802}, + {"bf16_lds": True}, + ) + ], + "W2_pipeline", + ) + best = min(results, key=lambda item: item[0])[1] + combine_results = [] + for blocks, waves in ( + (16, 4), + (32, 4), + (64, 4), + (128, 4), + (128, 8), + (256, 4), + (256, 8), + ): + evaluate( + [replace(best, combine_blocks=blocks, combine_waves=waves)], + f"combine_{blocks}_{waves}", + ) + combine_results.append((results[-1][0], blocks, waves)) + if rank == 0: + print(f"[COMBINE_WINNER] M={tokens} {min(combine_results)}", flush=True) + mega.comb_op.cfg = mega.comb_cfg + if rank == 0: + best_ms, best = min(results, key=lambda item: item[0]) + print( + f"[WINNER] EP{world} M={tokens} {best_ms * 1000:.1f}us {best}", + flush=True, + ) + if baseline_ms is not None: + print( + f"[BASELINE] M={tokens} MoRI+AITER={baseline_ms * 1000:.1f}us " + f"ratio={baseline_ms / best_ms:.3f}", + flush=True, + ) + barrier() + + for tokens in map(int, args.tokens.split(",")): + tune_case(tokens) + if args.selected_only: + from aiter.ops.flydsl.kernels.mega_moe.mega_moe_stage1 import ( + compile_mega_moe_stage1_bundle, + ) + from aiter.ops.flydsl.kernels.mega_moe.mega_moe_stage2 import _G2_LAUNCH_CACHE + + launcher = compile_mega_moe_stage1_bundle( + model_dim=mega.model_dim, + inter_dim=mega.inter_dim, + rank=mega.rank, + experts_per_rank=mega.epr, + fuse_npes=mega.world_size, + fuse_topk=mega.topk, + fuse_cap=mega._s1_cap, + fuse_mtpr=mega.mtpr, + fuse_scale_dim=mega._s1_scale_dim, + fixed_slot_dispatch=mega._s1_fixed_slot, + num_cu=mega._s1_num_cu, + swiglu_limit=mega.swiglu_limit, + tile_state_stride=mega._s1_tile_state_stride, + variants=mega._bundle_plan.stage1_variants, + ) + resources = {} + artifact_index = 0 + for launch in (launcher, *_G2_LAUNCH_CACHE.values()): + for artifact in launch._mem_cache.values(): + ir = artifact._ir_text + for binary in re.finditer(r'bin = "((?:[^"\\]|\\.)*)"', ir): + encoded = binary[1] + data = bytearray() + pos = 0 + while pos < len(encoded): + if encoded[pos] == "\\": + if encoded[pos + 1] in ('"', "\\"): + data.append(ord(encoded[pos + 1])) + pos += 2 + else: + data.append(int(encoded[pos + 1 : pos + 3], 16)) + pos += 3 + else: + data.append(ord(encoded[pos])) + pos += 1 + Path( + args.output + f".rank{rank}.{artifact_index}.hsaco" + ).write_bytes(data) + artifact_index += 1 + for match in re.finditer(r'gpu.kernel_metadata<"([^"]+)"', ir): + start = ir.find("metadata = {", match.end()) + end = ir.find("}", start) + resources[match[1]] = { + key: int(value) + for key, value in re.findall( + r"(\w+) = (\d+) : i64", ir[start:end] + ) + } + Path(args.output + f".rank{rank}.resources.json").write_text( + json.dumps(resources, indent=2) + "\n" + ) + ms.shmem_finalize() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() From c105e46ddfeca13689ee9e5b3a6ebb79fc62e385 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Mon, 5 Oct 2026 15:53:32 +0000 Subject: [PATCH 2/4] [FlyDSL] Keep Flash exhaustive tuning harness local Co-authored-by: Codex Signed-off-by: Fangzhou Ai --- op_tests/multigpu_tests/tune_mega_moe_v2.py | 516 -------------------- 1 file changed, 516 deletions(-) delete mode 100644 op_tests/multigpu_tests/tune_mega_moe_v2.py diff --git a/op_tests/multigpu_tests/tune_mega_moe_v2.py b/op_tests/multigpu_tests/tune_mega_moe_v2.py deleted file mode 100644 index f52e6bfb3b6..00000000000 --- a/op_tests/multigpu_tests/tune_mega_moe_v2.py +++ /dev/null @@ -1,516 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. -"""Layout and scheduling ablations for Flash EP2/EP4, with cold-cache timing.""" - -import argparse -import json -import os -import random -import re -from dataclasses import asdict, replace -from pathlib import Path - -from bench_mega_moe_v2 import ( - MegaMoEV2, - barrier, - capture, - dist, - make_inputs, - make_weights, - mori, - ms, - profile_graph, - run_online_mori, - setup_dist, - time_graph, - torch, -) - - -def layouts(base): - result = [base] - for bm, bn, waves in ( - (16, 128, 2), - (16, 256, 2), - (16, 256, 4), - (16, 512, 4), - (32, 128, 4), - (32, 64, 2), - (32, 128, 2), - (32, 256, 4), - (32, 256, 8), - (32, 512, 8), - (64, 256, 4), - (64, 128, 2), - (64, 128, 4), - (64, 256, 8), - (64, 512, 8), - (128, 256, 8), - (128, 128, 4), - ): - result.append( - replace( - base, - stage1=replace( - base.stage1, - sort_block_m=bm, - tile_n=bn, - num_waves=waves, - mfma_amajor=bm >= 32, - ), - stage2=replace(base.stage2, block_m=min(bm, 64)), - ) - ) - return list(dict.fromkeys(result)) - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--tokens", default="1,8,32,64,128,256,512,1024") - parser.add_argument("--mtpr", type=int, default=1024) - parser.add_argument("--iters", type=int, default=15) - parser.add_argument("--route", choices=("uniform", "hot-rank0"), default="uniform") - parser.add_argument("--output", required=True) - parser.add_argument("--schedule", action="store_true") - parser.add_argument("--paired-k", action="store_true") - parser.add_argument("--baseline", action="store_true") - parser.add_argument("--seed-config-json") - parser.add_argument("--kernel-knobs", action="store_true") - parser.add_argument("--transport", action="store_true") - parser.add_argument("--selected-only", action="store_true") - parser.add_argument("--profile-dir", default="") - args = parser.parse_args() - seed_records = ( - json.loads(Path(args.seed_config_json).read_text()) - if args.seed_config_json - else [] - ) - os.environ["MORI_EP_LAUNCH_CONFIG_MODE"] = "MANUAL" - os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0" - rank, world, device = setup_dist() - if world not in (2, 4): - raise ValueError("Flash requires EP2 or EP4") - w1, s1, w2, s2 = make_weights(384 // world, 5120, 2304, rank, device) - mega = MegaMoEV2( - rank=rank, - world_size=world, - model_dim=5120, - inter_dim=2304, - experts=384, - topk=6, - quant="a8w4", - w1=w1, - w1_scale=s1, - w2=w2, - w2_scale=s2, - max_tok_per_rank=args.mtpr, - swiglu_limit=10.0, - ) - default_select = mega._select_config - mori_op = None - expert_mask = torch.zeros(385, dtype=torch.int32, device=device) - expert_mask[rank * mega.epr : (rank + 1) * mega.epr] = 1 - if args.baseline: - mori_op = mori.ops.EpDispatchCombineOp( - mori.ops.EpDispatchCombineConfig( - data_type=torch.bfloat16, - rank=rank, - world_size=world, - hidden_dim=5120, - scale_dim=0, - scale_type_size=4, - max_token_type_size=2, - max_num_inp_token_per_rank=args.mtpr, - num_experts_per_rank=mega.epr, - num_experts_per_token=6, - warp_num_per_block=16, - block_num=80, - rdma_block_num=0, - kernel_type=mori.ops.EpDispatchCombineKernelType.IntraNode, - gpu_per_node=world, - ) - ) - cache = torch.zeros(512 * 1024 * 1024, dtype=torch.uint8, device=device) - cache_graph = capture(lambda: cache.add_(1)) - records = [] - - def tune_case(tokens): - x, weights, ids = make_inputs( - tokens, rank, world, 5120, 384, 6, args.route, 8.0, device - ) - counts = torch.bincount(ids.flatten().long(), minlength=384) - dist.all_reduce(counts) - local_counts = counts[rank * mega.epr : (rank + 1) * mega.epr] - received = ( - (ids[:, :, None] // mega.epr == torch.arange(world, device=device)) - .any(1) - .sum(0) - ) - dist.all_reduce(received) - mega._select_config = default_select - reference = mega(x, weights, ids, config_tokens=tokens).clone() - barrier() - base = default_select(tokens) - if seed_records: - from aiter.ops.flydsl.kernels.mega_moe.mega_moe_config import ( - MegaMoEConfig, - Stage1Config, - Stage2Config, - ) - - seed = min( - (row for row in seed_records if row["source_m"] == tokens), - key=lambda row: row["median_max_rank_ms"], - )["config"] - base = MegaMoEConfig( - stage1=Stage1Config(**seed["stage1"]), - stage2=Stage2Config(**seed["stage2"]), - p2p_quant=seed["p2p_quant"], - combine_blocks=seed.get("combine_blocks", 128), - combine_waves=seed.get("combine_waves", 8), - ) - baseline_ms = None - if mori_op is not None: - - def baseline_body(): - return run_online_mori( - mori_op, - x, - weights, - ids, - expert_mask, - (w1, s1, w2, s2), - 64, - 4, - tokens * world, - ) - - baseline_output = baseline_body().clone() - error = (reference.float() - baseline_output.float()).square().sum() - norm = baseline_output.float().square().sum() - dist.all_reduce(error) - dist.all_reduce(norm) - if float((error / norm).sqrt()) >= 0.1: - raise AssertionError("MegaMoE vs MoRI+AITER exceeds relL2=0.1") - baseline_graph = capture(baseline_body) - baseline_ms = time_graph( - baseline_graph, args.iters * 2, device, cache_graph - )[1] - if args.paired_k: - base = replace(base, stage1=replace(base.stage1, pair_k=True)) - results = [] - active = base - - def select(_): - mega._active_config = active - return active - - mega._select_config = select - - def body(): - return mega(x, weights, ids, config_tokens=tokens) - - def evaluate(configs, phase): - nonlocal active - configs = list(dict.fromkeys(configs)) - random.Random(tokens + len(results)).shuffle(configs) - graphs = [] - for config in configs: - active = config - candidate = body() - error = (candidate.float() - reference.float()).square().sum() - norm = reference.float().square().sum() - dist.all_reduce(error) - dist.all_reduce(norm) - rel_l2 = float((error / norm).sqrt()) - if rel_l2 >= 0.001: - raise AssertionError( - f"EP{world} M{tokens} {config}: relL2={rel_l2}" - ) - graphs.append((config, capture(body), rel_l2)) - # Reverse the order on the second pass to share clock/thermal drift. - timings = {config: [] for config, _, _ in graphs} - for order in (graphs, list(reversed(graphs))): - for config, graph, _ in order: - active = config - timings[config].append( - time_graph(graph, args.iters, device, cache_graph)[1] - ) - for config, _, rel_l2 in graphs: - bm = config.stage1.sort_block_m - tiles = (local_counts + bm - 1) // bm - weight_bytes = float(tiles.sum()) * 3 * 5120 * 2304 * (0.5 + 1 / 32) - lower_bytes = ( - float((local_counts > 0).sum()) * 3 * 5120 * 2304 * (0.5 + 1 / 32) - ) - ms_max = sum(timings[config]) / len(timings[config]) - record = { - "ep": world, - "source_m": tokens, - "mtpr": args.mtpr, - "fixed": mega._s1_fixed_slot, - "route": args.route, - "received_unique": received.tolist(), - "phase": phase, - "config": asdict(config), - "rel_l2": rel_l2, - "median_max_rank_ms": ms_max, - "passes": timings[config], - "baseline_ms": baseline_ms, - "rank0_row_utilization": float(local_counts.sum()) - / max(1, int(tiles.sum()) * bm), - "rank0_ideal_weight_bytes": lower_bytes, - "rank0_tile_weight_bytes": weight_bytes, - "rank0_ideal_weight_gbps": lower_bytes / (ms_max * 1e6), - } - results.append((ms_max, config)) - if rank == 0: - records.append(record) - Path(args.output).write_text(json.dumps(records, indent=2) + "\n") - print( - f"[TUNE] EP{world} M={tokens} {phase} {ms_max * 1000:.1f}us " - f"W13={bm}x{config.stage1.tile_n}/{config.stage1.num_waves}w " - f"W2={config.stage2.block_m}x{config.stage2.block_n} " - f"rows={record['rank0_row_utilization']:.3f} relL2={rel_l2:.6f}", - flush=True, - ) - - evaluate( - [base] if args.transport or args.selected_only else layouts(base), "layout" - ) - if args.selected_only: - if args.profile_dir: - active = base - profile_graph( - capture(body), - f"mega_m{tokens}", - rank, - args.profile_dir, - cache_graph=cache_graph, - ) - barrier() - return - if args.transport: - evaluate( - [ - replace( - base, - stage1=replace( - base.stage1, payload_chunk_rows=chunk, prepare_quant_cu=cu - ), - ) - for chunk in (128, 256, 512, 1024) - for cu in (32, 64, 128) - ], - "prepare_payload", - ) - best = min(results, key=lambda item: item[0])[1] - evaluate( - [ - replace(best, stage2=replace(best.stage2, block_m=bm, block_n=bn)) - for bm in (16, 32, 64) - if bm <= best.stage1.sort_block_m - for bn in (128, 256, 512) - ], - "W2_rows", - ) - best = min(results, key=lambda item: item[0])[1] - evaluate( - [ - replace( - best, - stage2=replace( - best.stage2, - block_n=n, - persist_cu=cu, - use_nt=nt, - persist_strided=strided, - ), - ) - for n, cu, nt, strided in ( - (128, 256, True, False), - (256, 128, True, False), - (256, 192, True, False), - (256, 256, False, False), - (256, 256, True, True), - (512, 256, True, False), - (512, 128, True, False), - ) - ], - "W2", - ) - if best.stage1.sort_block_m >= 64: - evaluate( - [ - replace(best, stage2=replace(best.stage2, block_m=bm, block_n=bn)) - for bm in (32, 64) - for bn in (128, 256) - ], - "W2_rows", - ) - if args.schedule: - best = min(results, key=lambda item: item[0])[1] - dispatch = (4, 8, 16, 32, 64, 128) - evaluate( - [ - replace(best, stage1=replace(best.stage1, num_dispatch_cu=d)) - for d in dispatch - if d % world == 0 - ], - "dispatch", - ) - best = min(results, key=lambda item: item[0])[1] - evaluate( - [ - replace( - best, - stage1=replace( - best.stage1, - grid_mult=g, - work_shards=s, - b_nt=nt, - waves_per_eu_hint=wpe, - ), - ) - for g, s, nt, wpe in ( - (1, 1, 3, 2), - (1, 2, 3, 2), - (1, 8, 3, 2), - (2, 4, 3, 2), - (3, 4, 3, 2), - (1, 4, 0, 2), - (1, 4, 3, 1), - (1, 4, 3, 4), - ) - ], - "W13_schedule", - ) - if args.kernel_knobs: - best = min(results, key=lambda item: item[0])[1] - evaluate( - [ - replace(best, stage1=replace(best.stage1, **updates)) - for updates in ( - {"mfma_amajor": not best.stage1.mfma_amajor}, - {"async_a_copy": False}, - {"pipe_weights": False}, - {"swizzle_a": False}, - ) - ], - "W13_pipeline", - ) - best = min(results, key=lambda item: item[0])[1] - evaluate( - [ - replace(best, stage2=replace(best.stage2, **updates)) - for updates in ( - {"b_hoist": False}, - {"ascale_prefetch": False}, - {"spatial_partition": 0}, - {"spatial_partition": 102}, - {"spatial_partition": 202}, - {"spatial_partition": 802}, - {"bf16_lds": True}, - ) - ], - "W2_pipeline", - ) - best = min(results, key=lambda item: item[0])[1] - combine_results = [] - for blocks, waves in ( - (16, 4), - (32, 4), - (64, 4), - (128, 4), - (128, 8), - (256, 4), - (256, 8), - ): - evaluate( - [replace(best, combine_blocks=blocks, combine_waves=waves)], - f"combine_{blocks}_{waves}", - ) - combine_results.append((results[-1][0], blocks, waves)) - if rank == 0: - print(f"[COMBINE_WINNER] M={tokens} {min(combine_results)}", flush=True) - mega.comb_op.cfg = mega.comb_cfg - if rank == 0: - best_ms, best = min(results, key=lambda item: item[0]) - print( - f"[WINNER] EP{world} M={tokens} {best_ms * 1000:.1f}us {best}", - flush=True, - ) - if baseline_ms is not None: - print( - f"[BASELINE] M={tokens} MoRI+AITER={baseline_ms * 1000:.1f}us " - f"ratio={baseline_ms / best_ms:.3f}", - flush=True, - ) - barrier() - - for tokens in map(int, args.tokens.split(",")): - tune_case(tokens) - if args.selected_only: - from aiter.ops.flydsl.kernels.mega_moe.mega_moe_stage1 import ( - compile_mega_moe_stage1_bundle, - ) - from aiter.ops.flydsl.kernels.mega_moe.mega_moe_stage2 import _G2_LAUNCH_CACHE - - launcher = compile_mega_moe_stage1_bundle( - model_dim=mega.model_dim, - inter_dim=mega.inter_dim, - rank=mega.rank, - experts_per_rank=mega.epr, - fuse_npes=mega.world_size, - fuse_topk=mega.topk, - fuse_cap=mega._s1_cap, - fuse_mtpr=mega.mtpr, - fuse_scale_dim=mega._s1_scale_dim, - fixed_slot_dispatch=mega._s1_fixed_slot, - num_cu=mega._s1_num_cu, - swiglu_limit=mega.swiglu_limit, - tile_state_stride=mega._s1_tile_state_stride, - variants=mega._bundle_plan.stage1_variants, - ) - resources = {} - artifact_index = 0 - for launch in (launcher, *_G2_LAUNCH_CACHE.values()): - for artifact in launch._mem_cache.values(): - ir = artifact._ir_text - for binary in re.finditer(r'bin = "((?:[^"\\]|\\.)*)"', ir): - encoded = binary[1] - data = bytearray() - pos = 0 - while pos < len(encoded): - if encoded[pos] == "\\": - if encoded[pos + 1] in ('"', "\\"): - data.append(ord(encoded[pos + 1])) - pos += 2 - else: - data.append(int(encoded[pos + 1 : pos + 3], 16)) - pos += 3 - else: - data.append(ord(encoded[pos])) - pos += 1 - Path( - args.output + f".rank{rank}.{artifact_index}.hsaco" - ).write_bytes(data) - artifact_index += 1 - for match in re.finditer(r'gpu.kernel_metadata<"([^"]+)"', ir): - start = ir.find("metadata = {", match.end()) - end = ir.find("}", start) - resources[match[1]] = { - key: int(value) - for key, value in re.findall( - r"(\w+) = (\d+) : i64", ir[start:end] - ) - } - Path(args.output + f".rank{rank}.resources.json").write_text( - json.dumps(resources, indent=2) + "\n" - ) - ms.shmem_finalize() - dist.destroy_process_group() - - -if __name__ == "__main__": - main() From ca048428946feb075fee7c8043b518fb5fdedcae Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Mon, 5 Oct 2026 16:24:20 +0000 Subject: [PATCH 3/4] Tune Flash EP2 speculative decode shapes and remove fixed-slot override Co-authored-by: Codex Signed-off-by: Fangzhou Ai --- aiter/configs/mega_moe_flash_gfx950.json | 837 +++++++++++++++++- aiter/ops/flydsl/kernels/mega_moe/envs.py | 8 - .../kernels/mega_moe/mega_moe_config.py | 39 +- 3 files changed, 848 insertions(+), 36 deletions(-) diff --git a/aiter/configs/mega_moe_flash_gfx950.json b/aiter/configs/mega_moe_flash_gfx950.json index 1388f02f79a..450f94d459f 100644 --- a/aiter/configs/mega_moe_flash_gfx950.json +++ b/aiter/configs/mega_moe_flash_gfx950.json @@ -1379,6 +1379,191 @@ "b_nt": 0, "work_shards": 4 }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 0, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "ascale_prefetch": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "b_hoist": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 32, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, "stage2": { "block_m": 16, "block_n": 256, @@ -1392,6 +1577,33 @@ "combine_blocks": 64, "combine_waves": 4 }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "b_hoist": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, { "stage1": { "sort_block_m": 16, @@ -1425,22 +1637,581 @@ "num_waves": 4, "grid_mult": 1, "num_dispatch_cu": 128, - "mfma_amajor": false, + "mfma_amajor": true, "async_a_copy": true, "use_tile_resource": true, "b_nt": 0, + "swizzle_a": false, "work_shards": 4 }, "stage2": { "block_m": 16, - "block_n": 128, + "block_n": 256, "persist": true, "persist_cu": 256, "use_nt": true, - "b_hoist": false, + "persist_strided": true, + "ascale_prefetch": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "ascale_prefetch": false, "spatial_partition": 202 }, "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "ascale_prefetch": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "b_hoist": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 512, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "b_hoist": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "b_hoist": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "b_hoist": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", "combine_blocks": 128, "combine_waves": 4 }, @@ -1535,11 +2306,63 @@ "62": 51, "63": 44, "64": 52, + "66": 53, + "72": 54, + "78": 20, + "84": 44, + "90": 55, "96": 48, - "128": 53, - "192": 54, - "256": 55, - "512": 56 + "102": 48, + "108": 48, + "114": 56, + "120": 57, + "126": 58, + "128": 59, + "132": 59, + "138": 59, + "144": 60, + "150": 61, + "156": 59, + "162": 62, + "168": 63, + "174": 62, + "180": 64, + "186": 62, + "192": 62, + "198": 62, + "204": 62, + "210": 65, + "216": 62, + "222": 62, + "228": 66, + "234": 67, + "240": 68, + "246": 67, + "252": 69, + "256": 70, + "258": 70, + "264": 71, + "270": 72, + "276": 71, + "282": 70, + "288": 73, + "294": 74, + "300": 75, + "306": 70, + "312": 76, + "318": 77, + "324": 72, + "330": 72, + "336": 78, + "342": 79, + "348": 74, + "354": 70, + "360": 80, + "366": 81, + "372": 82, + "378": 83, + "384": 84, + "512": 85 } }, "4": { diff --git a/aiter/ops/flydsl/kernels/mega_moe/envs.py b/aiter/ops/flydsl/kernels/mega_moe/envs.py index a1f2df14438..d27c4985435 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/envs.py +++ b/aiter/ops/flydsl/kernels/mega_moe/envs.py @@ -16,16 +16,8 @@ def _flag(value: str) -> bool: return value == "1" -def fixed_slot_max_mtpr(default: int) -> int: - return int(os.environ.get("AITER_MEGA_FIXED_SLOT_MAX_MTPR", str(default))) - - # name -> (parser, default) environment_variables: dict[str, tuple[Callable[[str], Any], str]] = { - # Largest MTPR still dispatched fixed-slot (direct expert slots, no count - # exchange). 511 / 1023 / 2047 admit the 256 / 512 / 1024-token capacities, - # validated on the EP8 v4_pro layout. Read once, at import. - "AITER_MEGA_FIXED_SLOT_MAX_MTPR": (int, "255"), # Default of forward(mask_invalid_slots=...): combine skips top-k slots with # id -1. Off because the id check costs prefill time; callers that pad with # -1 ask for it per call. diff --git a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py index ba01a2a7ca2..2d7ec8ba542 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py +++ b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py @@ -8,8 +8,6 @@ from functools import cache from pathlib import Path -from . import envs - TOKEN_BUCKETS = ( 1, 4, @@ -28,11 +26,12 @@ 32768, ) FLASH_TOKEN_BUCKETS = tuple(range(1, 65)) + (96, 128, 192) + TOKEN_BUCKETS[7:] +FLASH_EP2_TOKEN_BUCKETS = tuple( + sorted(set(FLASH_TOKEN_BUCKETS) | set(range(6, 385, 6))) +) FIXED_GRID_MULT_VALUES = (1, 2, 3, 4, 6, 8, 12, 16, 24, 32) P2P_FP8_MIN_MTPR = 1024 -# Fixed-slot (direct expert slots, no count exchange) is the default up to MTPR -# 128; AITER_MEGA_FIXED_SLOT_MAX_MTPR extends it (see envs.py). -FIXED_SLOT_MAX_MTPR = envs.AITER_MEGA_FIXED_SLOT_MAX_MTPR +FIXED_SLOT_MAX_MTPR = 128 MAX_MTPR_CLASS = 32768 # Source-indexed payload storage cuts the maximum-capacity activation buffer # by roughly ``topk``. Keep every smaller capacity on the historical layout. @@ -194,12 +193,12 @@ def mtpr_config_class(mtpr: int) -> int: def _use_fixed_slots(mtpr: int, world_size: int, experts_per_rank: int) -> bool: - geometry = (world_size, experts_per_rank) - if geometry in ((2, 192), (4, 96)): - limit = envs.fixed_slot_max_mtpr(512 if world_size == 2 else 1024) - else: - limit = FIXED_SLOT_MAX_MTPR - return mtpr <= limit and geometry in ((8, 48), (2, 192), (4, 96)) + limit = { + (8, 48): FIXED_SLOT_MAX_MTPR, + (2, 192): 512, + (4, 96): 1024, + }.get((world_size, experts_per_rank), 0) + return mtpr <= limit def _fixed_dispatch_cu(bucket: int) -> int: @@ -571,11 +570,9 @@ def select_mega_moe_config( "MegaMoE v2 fanout pair ids support at most " f"{MAX_FANOUT_EXPERTS_PER_RANK} experts per rank" ) - buckets = ( - FLASH_TOKEN_BUCKETS - if (model_dim, inter_dim, experts_per_rank * world_size) == (5120, 2304, 384) - else TOKEN_BUCKETS - ) + buckets = TOKEN_BUCKETS + if (model_dim, inter_dim, experts_per_rank * world_size) == (5120, 2304, 384): + buckets = FLASH_EP2_TOKEN_BUCKETS if world_size == 2 else FLASH_TOKEN_BUCKETS bucket = nearest_token_bucket(tokens, buckets) mtpr_class = mtpr_config_class(mtpr) fixed_slot_dispatch = _use_fixed_slots(mtpr_class, world_size, experts_per_rank) @@ -608,11 +605,11 @@ def build_mega_moe_bundle_plan( """Deduplicate variants while keeping Stage1/Stage2 selection atomic.""" if mtpr <= 0 or mtpr & (mtpr - 1): raise ValueError(f"mtpr={mtpr} must be a positive power of two") - token_buckets = ( - FLASH_TOKEN_BUCKETS - if (model_dim, inter_dim, experts_per_rank * world_size) == (5120, 2304, 384) - else TOKEN_BUCKETS - ) + token_buckets = TOKEN_BUCKETS + if (model_dim, inter_dim, experts_per_rank * world_size) == (5120, 2304, 384): + token_buckets = ( + FLASH_EP2_TOKEN_BUCKETS if world_size == 2 else FLASH_TOKEN_BUCKETS + ) buckets = tuple(bucket for bucket in token_buckets if bucket <= mtpr) if not buckets or buckets[-1] != mtpr: raise ValueError(f"mtpr={mtpr} has no exact token bucket") From 154b0a6d204dcbdf375cfc830c043f2e0b0ee69c Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Mon, 5 Oct 2026 17:43:46 +0000 Subject: [PATCH 4/4] Tune Flash EP2 and EP4 speculative decode through concurrency 192 Preserve measured compact workspace classes and extend per-M configuration coverage for both EP sizes. Co-authored-by: Codex Signed-off-by: Fangzhou Ai --- aiter/configs/mega_moe_flash_gfx950.json | 8775 ++++++++++++++++- .../kernels/mega_moe/mega_moe_config.py | 28 +- 2 files changed, 8745 insertions(+), 58 deletions(-) diff --git a/aiter/configs/mega_moe_flash_gfx950.json b/aiter/configs/mega_moe_flash_gfx950.json index 450f94d459f..0e9f0492b6e 100644 --- a/aiter/configs/mega_moe_flash_gfx950.json +++ b/aiter/configs/mega_moe_flash_gfx950.json @@ -2215,6 +2215,182 @@ "combine_blocks": 128, "combine_waves": 4 }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, { "stage1": { "sort_block_m": 32, @@ -2362,7 +2538,28 @@ "372": 82, "378": 83, "384": 84, - "512": 85 + "390": 85, + "396": 86, + "402": 87, + "408": 88, + "414": 86, + "420": 89, + "426": 88, + "432": 87, + "438": 88, + "444": 89, + "450": 84, + "456": 85, + "462": 90, + "468": 87, + "474": 84, + "480": 90, + "486": 90, + "492": 84, + "498": 90, + "504": 90, + "510": 91, + "512": 92 } }, "4": { @@ -3736,11 +3933,11 @@ }, "stage2": { "block_m": 16, - "block_n": 256, + "block_n": 128, "persist": true, - "persist_cu": 192, + "persist_cu": 256, "use_nt": true, - "spatial_partition": 802 + "spatial_partition": 102 }, "p2p_quant": "none", "combine_blocks": 128, @@ -3763,20 +3960,19 @@ "block_m": 16, "block_n": 256, "persist": true, - "persist_cu": 256, + "persist_cu": 128, "use_nt": true, - "persist_strided": true, - "spatial_partition": 802 + "spatial_partition": 0 }, "p2p_quant": "none", - "combine_blocks": 64, - "combine_waves": 4 + "combine_blocks": 128, + "combine_waves": 8 }, { "stage1": { - "sort_block_m": 32, + "sort_block_m": 16, "tile_n": 256, - "num_waves": 8, + "num_waves": 4, "grid_mult": 1, "num_dispatch_cu": 128, "mfma_amajor": false, @@ -3786,49 +3982,48 @@ "work_shards": 4 }, "stage2": { - "block_m": 32, + "block_m": 16, "block_n": 256, "persist": true, - "persist_cu": 256, + "persist_cu": 192, "use_nt": true, - "persist_strided": true, "spatial_partition": 802 }, "p2p_quant": "none", - "combine_blocks": 64, - "combine_waves": 4 + "combine_blocks": 128, + "combine_waves": 8 }, { "stage1": { - "sort_block_m": 32, + "sort_block_m": 16, "tile_n": 256, - "num_waves": 8, + "num_waves": 4, "grid_mult": 1, "num_dispatch_cu": 128, - "mfma_amajor": false, + "mfma_amajor": true, "async_a_copy": true, "use_tile_resource": true, "b_nt": 0, - "swizzle_a": false, "work_shards": 4 }, "stage2": { - "block_m": 32, + "block_m": 16, "block_n": 256, "persist": true, - "persist_cu": 128, + "persist_cu": 256, "use_nt": true, + "persist_strided": true, "spatial_partition": 802 }, "p2p_quant": "none", - "combine_blocks": 128, - "combine_waves": 4 + "combine_blocks": 256, + "combine_waves": 8 }, { "stage1": { - "sort_block_m": 32, - "tile_n": 512, - "num_waves": 8, + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, "grid_mult": 1, "num_dispatch_cu": 128, "mfma_amajor": false, @@ -3838,41 +4033,2192 @@ "work_shards": 4 }, "stage2": { - "block_m": 32, + "block_m": 16, "block_n": 256, "persist": true, "persist_cu": 256, - "use_nt": false, - "spatial_partition": 0 + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 }, "p2p_quant": "none", - "combine_blocks": 128, - "combine_waves": 8 + "combine_blocks": 64, + "combine_waves": 4 }, { "stage1": { - "sort_block_m": 64, - "tile_n": 512, - "num_waves": 8, + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, "grid_mult": 1, "num_dispatch_cu": 128, - "mfma_amajor": true, - "async_a_copy": false, + "mfma_amajor": false, + "async_a_copy": true, "use_tile_resource": true, "b_nt": 0, "work_shards": 4 }, "stage2": { - "block_m": 64, - "block_n": 128, + "block_m": 16, + "block_n": 256, "persist": true, "persist_cu": 256, "use_nt": true, - "spatial_partition": 802, - "bf16_lds": true + "persist_strided": true, + "spatial_partition": 802 }, "p2p_quant": "none", - "combine_blocks": 256, + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "ascale_prefetch": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 16, + "tile_n": 256, + "num_waves": 4, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 16, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "b_hoist": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 64, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "swizzle_a": false, + "work_shards": 1 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "swizzle_a": false, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 2 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 3, + "waves_per_eu_hint": 1, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 2 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "waves_per_eu_hint": 1, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 2 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 102 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 802 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 256, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": false, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 2 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 64, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 64, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 3, + "waves_per_eu_hint": 1, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 2 + }, + "stage2": { + "block_m": 64, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1 + }, + "stage2": { + "block_m": 64, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 64, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 64, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 32, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 64, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 64, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 128, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4 + }, + "stage2": { + "block_m": 64, + "block_n": 128, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802, + "bf16_lds": true + }, + "p2p_quant": "none", + "combine_blocks": 256, "combine_waves": 8 } ], @@ -3941,12 +6287,6345 @@ "62": 48, "63": 42, "64": 52, - "96": 53, - "128": 54, - "192": 55, - "256": 56, - "512": 57, - "1024": 58 + "66": 24, + "72": 25, + "78": 25, + "84": 53, + "90": 54, + "96": 55, + "102": 25, + "108": 30, + "114": 56, + "120": 55, + "126": 33, + "128": 57, + "132": 58, + "138": 59, + "144": 60, + "150": 61, + "156": 62, + "162": 63, + "168": 64, + "174": 65, + "180": 66, + "186": 67, + "192": 68, + "198": 64, + "204": 69, + "210": 66, + "216": 66, + "222": 70, + "228": 71, + "234": 72, + "240": 73, + "246": 74, + "252": 72, + "256": 75, + "258": 76, + "264": 77, + "270": 78, + "276": 76, + "282": 79, + "288": 80, + "294": 81, + "300": 82, + "306": 83, + "312": 84, + "318": 81, + "324": 84, + "330": 85, + "336": 86, + "342": 87, + "348": 84, + "354": 81, + "360": 81, + "366": 88, + "372": 81, + "378": 89, + "384": 90, + "390": 90, + "396": 91, + "402": 92, + "408": 92, + "414": 90, + "420": 93, + "426": 92, + "432": 94, + "438": 92, + "444": 95, + "450": 95, + "456": 96, + "462": 97, + "468": 98, + "474": 99, + "480": 100, + "486": 101, + "492": 102, + "498": 102, + "504": 102, + "510": 102, + "512": 103, + "516": 104, + "522": 105, + "528": 106, + "534": 107, + "540": 106, + "546": 108, + "552": 109, + "558": 110, + "564": 111, + "570": 112, + "576": 109, + "582": 112, + "588": 109, + "594": 109, + "600": 113, + "606": 114, + "612": 115, + "618": 114, + "624": 116, + "630": 117, + "636": 118, + "642": 118, + "648": 115, + "654": 117, + "660": 118, + "666": 119, + "672": 120, + "678": 119, + "684": 115, + "690": 121, + "696": 122, + "702": 123, + "708": 115, + "714": 124, + "720": 115, + "726": 125, + "732": 113, + "738": 119, + "744": 115, + "750": 126, + "756": 127, + "762": 109, + "768": 128, + "774": 129, + "780": 130, + "786": 131, + "792": 130, + "798": 132, + "804": 132, + "810": 132, + "816": 133, + "822": 134, + "828": 135, + "834": 128, + "840": 136, + "846": 133, + "852": 129, + "858": 133, + "864": 137, + "870": 138, + "876": 137, + "882": 139, + "888": 139, + "894": 137, + "900": 138, + "906": 138, + "912": 137, + "918": 140, + "924": 137, + "930": 137, + "936": 137, + "942": 140, + "948": 137, + "954": 137, + "960": 140, + "966": 137, + "972": 137, + "978": 137, + "984": 141, + "990": 137, + "996": 137, + "1002": 137, + "1008": 138, + "1014": 140, + "1020": 138, + "1024": 142 + } + } + }, + "compact": { + "2": { + "8192": { + "configs": [ + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 384 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "persist_strided": true, + "skew_cu": 96, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 384 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 240, + "use_nt": false, + "skew_cu": 96, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 384 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 240, + "use_nt": false, + "skew_cu": 112 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + } + ], + "tokens": { + "516": 0, + "522": 1, + "528": 2, + "534": 3, + "540": 4, + "546": 5, + "552": 6, + "558": 7, + "564": 8, + "570": 9, + "576": 0, + "582": 10, + "588": 11, + "594": 12, + "600": 4, + "606": 13, + "612": 5, + "618": 14, + "624": 15, + "630": 16, + "636": 17, + "642": 18, + "648": 19, + "654": 20, + "660": 21, + "666": 7, + "672": 6, + "678": 22, + "684": 23, + "690": 24, + "696": 25, + "702": 26, + "708": 27, + "714": 28, + "720": 4, + "726": 29, + "732": 30, + "738": 31, + "744": 17, + "750": 32, + "756": 33, + "762": 0, + "768": 34, + "774": 35, + "780": 36, + "786": 37, + "792": 34, + "798": 38, + "804": 34, + "810": 39, + "816": 36, + "822": 40, + "828": 41, + "834": 42, + "840": 36, + "846": 40, + "852": 43, + "858": 43, + "864": 44, + "870": 43, + "876": 45, + "882": 46, + "888": 47, + "894": 48, + "900": 45, + "906": 49, + "912": 46, + "918": 50, + "924": 45, + "930": 51, + "936": 52, + "942": 53, + "948": 54, + "954": 55, + "960": 46, + "966": 56, + "972": 57, + "978": 58, + "984": 59, + "990": 60, + "996": 61, + "1002": 61, + "1008": 62, + "1014": 63, + "1020": 64, + "1024": 65, + "1026": 57, + "1032": 57, + "1038": 66, + "1044": 67, + "1050": 51, + "1056": 68, + "1062": 69, + "1068": 70, + "1074": 71, + "1080": 61, + "1086": 54, + "1092": 72, + "1098": 70, + "1104": 73, + "1110": 61, + "1116": 63, + "1122": 74, + "1128": 51, + "1134": 75, + "1140": 76, + "1146": 77, + "1152": 57, + "2048": 78, + "4096": 79, + "8192": 80 + } + }, + "16384": { + "configs": [ + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 192, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 64, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "persist_strided": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 32, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 16, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 3, + "work_shards": 1, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 384 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "persist_strided": true, + "skew_cu": 96, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 384 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 240, + "use_nt": false, + "skew_cu": 96, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 384 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 240, + "use_nt": false, + "skew_cu": 112 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 384, + "prepare_quant_cu": 96 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 128, + "use_nt": true, + "skew_cu": 96 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + } + ], + "tokens": { + "516": 0, + "522": 1, + "528": 0, + "534": 2, + "540": 3, + "546": 4, + "552": 4, + "558": 5, + "564": 6, + "570": 7, + "576": 2, + "582": 8, + "588": 9, + "594": 10, + "600": 11, + "606": 12, + "612": 13, + "618": 14, + "624": 15, + "630": 16, + "636": 17, + "642": 8, + "648": 18, + "654": 19, + "660": 18, + "666": 2, + "672": 20, + "678": 21, + "684": 22, + "690": 23, + "696": 24, + "702": 22, + "708": 25, + "714": 26, + "720": 27, + "726": 28, + "732": 29, + "738": 30, + "744": 31, + "750": 32, + "756": 29, + "762": 33, + "768": 34, + "774": 35, + "780": 36, + "786": 36, + "792": 37, + "798": 38, + "804": 34, + "810": 39, + "816": 40, + "822": 41, + "828": 42, + "834": 43, + "840": 44, + "846": 45, + "852": 34, + "858": 45, + "864": 46, + "870": 47, + "876": 48, + "882": 49, + "888": 50, + "894": 48, + "900": 51, + "906": 52, + "912": 53, + "918": 54, + "924": 49, + "930": 55, + "936": 56, + "942": 57, + "948": 51, + "954": 58, + "960": 59, + "966": 60, + "972": 61, + "978": 61, + "984": 62, + "990": 63, + "996": 64, + "1002": 65, + "1008": 66, + "1014": 67, + "1020": 68, + "1024": 69, + "1026": 70, + "1032": 71, + "1038": 72, + "1044": 73, + "1050": 74, + "1056": 75, + "1062": 76, + "1068": 77, + "1074": 78, + "1080": 79, + "1086": 80, + "1092": 81, + "1098": 72, + "1104": 66, + "1110": 82, + "1116": 83, + "1122": 74, + "1128": 71, + "1134": 84, + "1140": 85, + "1146": 57, + "1152": 80, + "2048": 86, + "4096": 87, + "8192": 88, + "16384": 89 + } + } + }, + "4": { + "8192": { + "configs": [ + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "ascale_prefetch": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 512, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "persist_strided": true, + "skew_cu": 96, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 384 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 240, + "use_nt": false, + "skew_cu": 96, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 384 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 240, + "use_nt": false, + "skew_cu": 112, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + } + ], + "tokens": { + "1026": 0, + "1032": 1, + "1038": 2, + "1044": 3, + "1050": 4, + "1056": 5, + "1062": 6, + "1068": 7, + "1074": 8, + "1080": 9, + "1086": 10, + "1092": 11, + "1098": 12, + "1104": 13, + "1110": 14, + "1116": 15, + "1122": 16, + "1128": 17, + "1134": 18, + "1140": 19, + "1146": 20, + "1152": 21, + "2048": 22, + "4096": 23, + "8192": 24 + } + }, + "16384": { + "configs": [ + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "b_hoist": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "ascale_prefetch": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "ascale_prefetch": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 802 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 128, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 4 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 128, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 1024, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 256, + "use_nt": true, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 256, + "prepare_quant_cu": 32 + }, + "stage2": { + "block_m": 32, + "block_n": 256, + "persist": true, + "persist_cu": 256, + "use_nt": false, + "persist_strided": true, + "skew_cu": 96, + "spatial_partition": 102 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": true, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 384 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 240, + "use_nt": false, + "skew_cu": 96, + "spatial_partition": 202 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 384 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 240, + "use_nt": false, + "skew_cu": 112 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + }, + { + "stage1": { + "sort_block_m": 64, + "tile_n": 512, + "num_waves": 8, + "grid_mult": 1, + "num_dispatch_cu": 32, + "mfma_amajor": true, + "async_a_copy": false, + "use_tile_resource": true, + "b_nt": 0, + "work_shards": 4, + "payload_chunk_rows": 384, + "prepare_quant_cu": 96 + }, + "stage2": { + "block_m": 64, + "block_n": 512, + "persist": true, + "persist_cu": 192, + "use_nt": false, + "skew_cu": 96, + "spatial_partition": 0 + }, + "p2p_quant": "fp8_blockwise_1x32", + "combine_blocks": 256, + "combine_waves": 8 + } + ], + "tokens": { + "1026": 0, + "1032": 1, + "1038": 2, + "1044": 3, + "1050": 4, + "1056": 5, + "1062": 6, + "1068": 7, + "1074": 8, + "1080": 9, + "1086": 10, + "1092": 11, + "1098": 12, + "1104": 3, + "1110": 13, + "1116": 11, + "1122": 14, + "1128": 15, + "1134": 16, + "1140": 17, + "1146": 16, + "1152": 18, + "2048": 19, + "4096": 20, + "8192": 21, + "16384": 22 + } } } } diff --git a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py index 2d7ec8ba542..bf4f94d2b55 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py +++ b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py @@ -26,8 +26,8 @@ 32768, ) FLASH_TOKEN_BUCKETS = tuple(range(1, 65)) + (96, 128, 192) + TOKEN_BUCKETS[7:] -FLASH_EP2_TOKEN_BUCKETS = tuple( - sorted(set(FLASH_TOKEN_BUCKETS) | set(range(6, 385, 6))) +FLASH_SPEC_DECODE_TOKEN_BUCKETS = tuple( + sorted(set(FLASH_TOKEN_BUCKETS) | set(range(6, 1153, 6))) ) FIXED_GRID_MULT_VALUES = (1, 2, 3, 4, 6, 8, 12, 16, 24, 32) P2P_FP8_MIN_MTPR = 1024 @@ -380,8 +380,13 @@ def _flash_tuning_data(): def _select_flash_config(bucket, mtpr_class, world_size, fixed): - if fixed: - policy = _flash_tuning_data()["fixed"][str(world_size)] + tuning = _flash_tuning_data() + policy = ( + tuning["fixed"][str(world_size)] + if fixed + else tuning.get("compact", {}).get(str(world_size), {}).get(str(mtpr_class)) + ) + if policy is not None: config_id = policy["tokens"].get(str(bucket)) if config_id is not None: config = policy["configs"][config_id] @@ -570,11 +575,16 @@ def select_mega_moe_config( "MegaMoE v2 fanout pair ids support at most " f"{MAX_FANOUT_EXPERTS_PER_RANK} experts per rank" ) + is_flash = (model_dim, inter_dim, experts_per_rank * world_size) == ( + 5120, + 2304, + 384, + ) buckets = TOKEN_BUCKETS - if (model_dim, inter_dim, experts_per_rank * world_size) == (5120, 2304, 384): - buckets = FLASH_EP2_TOKEN_BUCKETS if world_size == 2 else FLASH_TOKEN_BUCKETS + if is_flash: + buckets = FLASH_SPEC_DECODE_TOKEN_BUCKETS bucket = nearest_token_bucket(tokens, buckets) - mtpr_class = mtpr_config_class(mtpr) + mtpr_class = mtpr if is_flash else mtpr_config_class(mtpr) fixed_slot_dispatch = _use_fixed_slots(mtpr_class, world_size, experts_per_rank) total_segments = world_size * experts_per_rank + world_size if total_segments > MAX_FANOUT_SEGMENTS: @@ -607,9 +617,7 @@ def build_mega_moe_bundle_plan( raise ValueError(f"mtpr={mtpr} must be a positive power of two") token_buckets = TOKEN_BUCKETS if (model_dim, inter_dim, experts_per_rank * world_size) == (5120, 2304, 384): - token_buckets = ( - FLASH_EP2_TOKEN_BUCKETS if world_size == 2 else FLASH_TOKEN_BUCKETS - ) + token_buckets = FLASH_SPEC_DECODE_TOKEN_BUCKETS buckets = tuple(bucket for bucket in token_buckets if bucket <= mtpr) if not buckets or buckets[-1] != mtpr: raise ValueError(f"mtpr={mtpr} has no exact token bucket")