[TRITON]: Moe tuning mi350

[Chi-Chu319]

Authors: @Chi-Chu319 @juuso-oskari

WARNING: Don't merge this yet. It fails testing small shapes with our num_stage=3 change. Investigation on going. Reproducer

pytest   op_tests/triton_tests/test_moe.py::test_fused_moe[True-True-dtype0-False-False-True-64-14336-4096-2-8]

We change the num_stage to 2. We will investigate this outside scope of this pr. If its compiler related we will raise attention from the compiler ppl.

• Fix config loading
• configs for mi350
• Added various support for benchmark for future tuning
• EVEN_K loading of a and b
• 405B expert matrix loading int32 overflow

Fixed xcd remapping: Previously we used GRID_MN = EM // BLOCK_SIZE_M * num_pid_n for the xcd_remap(pid, GRID_MN,NUM_XCD=8). This results in the last XCD receiving all the dummy tokens (as EM=num_tokens_post_padded + dummy tokens) which have no work and have early return -> the last XCD is idling while other XCDs still have work. Fixed the GRID_MN to be calculated as GRID_MN = num_tokens_post_padded // BLOCK_SIZE_M * num_pid_n.

Perf comparison to main

main:

Profiling Summary (kernel times in ns, TFLOPs, and Bandwidth in GB/s):
+------------------+-------------------+------------------------+--------+------------------+-----------------+
|    M | model     | _fused_moe_kernel | Other Kernels Sum (ns) | TFLOPs | Bandwidth (GB/s) | AI (Flops/Byte) |
+------------------+-------------------+------------------------+--------+------------------+-----------------+
| 128 llama3-405B  |    13150217.33    |       1164348.65       | 33.97  |     1063.87      |      31.93      |
| 256 llama3-405B  |    14980420.00    |       1185801.70       | 59.63  |      935.99      |      63.71      |
| 512 llama3-405B  |    17507314.60    |       1252937.57       | 102.05 |      804.49      |      126.86     |
| 1024 llama3-405B |    23240926.78    |       1388565.29       | 153.76 |      611.43      |      251.47     |
| 2048 llama3-405B |    33828796.57    |       1304083.53       | 211.26 |      427.50      |      494.18     |
+------------------+-------------------+------------------------+--------+------------------+-----------------+

moe-tuning-mi350:

Profiling Summary (kernel times in ns, TFLOPs, and Bandwidth in GB/s):
+------------------+-------------------+------------------------+--------+------------------+-----------------+
|    M | model     | _fused_moe_kernel | Other Kernels Sum (ns) | TFLOPs | Bandwidth (GB/s) | AI (Flops/Byte) |
+------------------+-------------------+------------------------+--------+------------------+-----------------+
| 128 llama3-405B  |     2838261.55    |       1094917.86       | 157.38 |     4929.11      |      31.93      |
| 256 llama3-405B  |     2931990.21    |       1110832.19       | 304.69 |     4782.27      |      63.71      |
| 512 llama3-405B  |     4357953.36    |       1192987.41       | 409.99 |     3231.90      |      126.86     |
| 1024 llama3-405B |     6993911.33    |       1260398.59       | 510.93 |     2031.81      |      251.47     |
| 2048 llama3-405B |    12335273.92    |       1241415.39       | 579.38 |     1172.41      |      494.18     |
+------------------+-------------------+------------------------+--------+------------------+-----------------+

[rahulbatra85]

@Chi-Chu319 Please address black CI issues