support swapAB for m_grouped_fp8_gemm_nt_masked

[Wangzheee]

SwapAB: Significantly improve the performance for M%64<32

Description

• Significantly improve the performance for BLOCK_M = 32 or M%64<32
• Swap A B: WGMMA::wgmma(desc_b, desc_a, accum, k)
• Use multi-math_warp_groups and multi-wave MMAs support BLOCK_N(64, 128, 256): BLOCK_N=256 for H20

How to use

• export ENABLE_SWAPAB=1

Improvements (H20)

Aligned M, desired state: masked_m[j] = int(expected_m_per_group * random.uniform(1, 1))

• Original

Perf (num_groups=16, expected_m_per_group= 32, n=4096, k=7168, 1D2D): 257 us | 117 TFLOPS | 1862 GB/s
Perf (num_groups=16, expected_m_per_group= 32, n=7168, k=2048, 1D2D): 127 us | 118 TFLOPS | 1914 GB/s
Perf (num_groups=16, expected_m_per_group= 96, n=4096, k=7168, 1D2D): 497 us | 181 TFLOPS | 993 GB/s
Perf (num_groups=16, expected_m_per_group= 96, n=7168, k=2048, 1D2D): 240 us | 188 TFLOPS | 1085 GB/s
Perf (num_groups=16, expected_m_per_group= 160, n=4096, k=7168, 1D2D): 921 us | 163 TFLOPS | 554 GB/s
Perf (num_groups=16, expected_m_per_group= 160, n=7168, k=2048, 1D2D): 472 us | 159 TFLOPS | 587 GB/s

• SwapAB

Perf (num_groups=16, expected_m_per_group= 32, n=4096, k=7168, 1D2D): 154 us | 195 TFLOPS | 3100 GB/s
Perf (num_groups=16, expected_m_per_group= 32, n=7168, k=2048, 1D2D): 84 us | 180 TFLOPS | 2906 GB/s
Perf (num_groups=16, expected_m_per_group= 96, n=4096, k=7168, 1D2D): 340 us | 266 TFLOPS | 1454 GB/s
Perf (num_groups=16, expected_m_per_group= 96, n=7168, k=2048, 1D2D): 191 us | 236 TFLOPS | 1364 GB/s
Perf (num_groups=16, expected_m_per_group= 160, n=4096, k=7168, 1D2D): 572 us | 263 TFLOPS | 891 GB/s
Perf (num_groups=16, expected_m_per_group= 160, n=7168, k=2048, 1D2D): 303 us | 248 TFLOPS | 915 GB/s

Other case (original test): masked_m[j] = int(expected_m_per_group * random.uniform(0.7, 1.3))

• Original

Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D): 349 us | 214 TFLOPS | 141 GB/s
Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D): 159 us | 215 TFLOPS | 213 GB/s
Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D): 348 us | 178 TFLOPS | 216 GB/s
Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D): 159 us | 191 TFLOPS | 292 GB/s
Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D): 347 us | 174 TFLOPS | 384 GB/s
Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D): 217 us | 146 TFLOPS | 353 GB/s
Perf (num_groups=16, expected_m_per_group= 64, n=4096, k=7168, 1D2D): 406 us | 153 TFLOPS | 1199 GB/s
Perf (num_groups=16, expected_m_per_group= 64, n=7168, k=2048, 1D2D): 172 us | 164 TFLOPS | 1455 GB/s
Perf (num_groups=16, expected_m_per_group= 128, n=4096, k=7168, 1D2D): 740 us | 165 TFLOPS | 678 GB/s
Perf (num_groups=16, expected_m_per_group= 128, n=7168, k=2048, 1D2D): 354 us | 168 TFLOPS | 758 GB/s

• SwapAB

Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D): 304 us | 246 TFLOPS | 162 GB/s
Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D): 145 us | 235 TFLOPS | 233 GB/s
Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D): 271 us | 229 TFLOPS | 278 GB/s
Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D): 136 us | 224 TFLOPS | 342 GB/s
Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D): 271 us | 223 TFLOPS | 493 GB/s
Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D): 146 us | 217 TFLOPS | 523 GB/s
Perf (num_groups=16, expected_m_per_group= 64, n=4096, k=7168, 1D2D): 308 us | 201 TFLOPS | 1580 GB/s
Perf (num_groups=16, expected_m_per_group= 64, n=7168, k=2048, 1D2D): 161 us | 175 TFLOPS | 1556 GB/s
Perf (num_groups=16, expected_m_per_group= 128, n=4096, k=7168, 1D2D): 539 us | 227 TFLOPS | 931 GB/s
Perf (num_groups=16, expected_m_per_group= 128, n=7168, k=2048, 1D2D): 272 us | 218 TFLOPS | 986 GB/s

TODO

• In LowLatency, expected_m does not represent the maximum possible value of actual shape M in GroupsGEMM. How to choose BLOCK_M and BLOCK_N?
• Manually set BLOCK_M and BLOCK_N: support swapAB for m_grouped_fp8_gemm_nt_masked #192 (comment)
• Set BLOCK_M and BLOCK_N to improve performance for H100/H800

[Wangzheee]

Manually set one of them. Experiments have found that in most cases, 256 performs the best

[LyricZhao]

Thanks! Merging it later.

[Wangzheee]

Thanks! Merging it later.

Hi～ Do you have a plan for code merge? We have already used this PR on our online service for H20.

[whybeyoung]

LGTM

[Chtholly-Boss]

为什么在other case里

num_groups=1, expected_m_per_group=1024, n=4096, k=7168

这个case也能有提升？num groups＝1时实际上相当于是一个1024×4096×7168的矩阵乘吧？SwapAB在这里的优势是什么

[LyricZhao]

Do you have a plan for code merge?

Sorry, we will try to merge this by the end of Oct. As swap AB will introduce non-batch-invariant and deterministic issues, we will consider it more carefully and do some refactors before merging. Also, as most the code can be reused, we will also refactor the epilogue part to make this feature less change for the code.

Thanks for your contribution! We will refactor for you, no change request👍🏻

cc @zheanxu

[Wangzheee]

为什么在other case里

num_groups=1, expected_m_per_group=1024, n=4096, k=7168

这个case也能有提升？num groups＝1时实际上相当于是一个1024×4096×7168的矩阵乘吧？SwapAB在这里的优势是什么

The swapAB variant “swap” the WGMMA tile usage, mapping the original problem’s M dimension onto WGMMA’s N dimension (which must be a multiple of 8). This enables smaller BLOCK_M (32). The performance advantage primarily comes from finer tiling granularity and better resource utilization.

[Wangzheee]

Do you have a plan for code merge?

Sorry, we will try to merge this by the end of Oct. As swap AB will introduce non-batch-invariant and deterministic issues, we will consider it more carefully and do some refactors before merging. Also, as most the code can be reused, we will also refactor the epilogue part to make this feature less change for the code.

Thanks for your contribution! We will refactor for you, no change request👍🏻

cc @zheanxu

Thanks～ Looking forward to the release of the new version.