[mxfp8 moe training] integrate mxfp8 dim0 triton kernel

[danielvegamyhre]

Stacked PRs:

• [mxfp8 moe training] use mxfp8 dim0 kernel in mxfp8 all2all; simplify mxfp8 dim0 kernel #3131
• ->[mxfp8 moe training] integrate mxfp8 dim0 triton kernel #3129
• [mxfp8 moe training] add triton kernel for mxfp8 quantization along dim0 #3128

---

[mxfp8 moe training] integrate mxfp8 dim0 triton kernel

Test plan

• pytest test/prototype/moe_training/test_scaled_grouped_mm.py -k dq -s

Benchmarks

M,N,K,G recipe bf16_fwd_bwd_us scaled_fwd_bwd_us scaled_fwd_bwd_speedup bf16_fwd_us scaled_fwd_us scaled_fwd_speedup ----------------------- -------------------- ----------------- ------------------- ------------------------ ------------- --------------- -------------------- (16384, 8192, 5120, 1) MoEScalingType.MXFP8 4239.74 2978.88 1.423x 1229.87 758.8 1.621x (16384, 8192, 5120, 2) MoEScalingType.MXFP8 4192.19 3381.7 1.24x 1229.5 1079.33 1.139x (16384, 8192, 5120, 4) MoEScalingType.MXFP8 3920.91 3419.1 1.147x 1093.42 820.416 1.333x (16384, 8192, 5120, 8) MoEScalingType.MXFP8 4309.06 3633.54 1.186x 1093.73 932.128 1.173x (128000, 8192, 5120, 1) MoEScalingType.MXFP8 50533.1 23270.9 2.172x 12149.6 6208.59 1.957x (128000, 8192, 5120, 2) MoEScalingType.MXFP8 57250.8 23629.6 2.423x 10176.2 6408.26 1.588x (128000, 8192, 5120, 4) MoEScalingType.MXFP8 35872.8 25179.2 1.425x 10041.7 5813.94 1.727x (128000, 8192, 5120, 8) MoEScalingType.MXFP8 50138.8 23592 2.125x 18598.9 6110.18 3.044x (16384, 1536, 5120, 1) MoEScalingType.MXFP8 808 987.136 0.819x 246.816 261.12 0.945x (16384, 1536, 5120, 2) MoEScalingType.MXFP8 855.072 914.56 0.935x 224.496 263.2 0.853x (16384, 1536, 5120, 4) MoEScalingType.MXFP8 824.4 1034.11 0.797x 287.744 273.28 1.053x (16384, 1536, 5120, 8) MoEScalingType.MXFP8 847.968 1033.44 0.821x 220.384 283.712 0.777x (128000, 1536, 5120, 1) MoEScalingType.MXFP8 6480.8 7623.65 0.85x 2100.29 2025.28 1.037x (128000, 1536, 5120, 2) MoEScalingType.MXFP8 6530.54 7277.41 0.897x 2112.32 1929.28 1.095x (128000, 1536, 5120, 4) MoEScalingType.MXFP8 7770.05 6168.45 1.26x 2020.35 1638.43 1.233x (128000, 1536, 5120, 8) MoEScalingType.MXFP8 7438.1 6244.24 1.191x 1847.2 1786.94 1.034x (16384, 2048, 7168, 1) MoEScalingType.MXFP8 1739.78 1519.78 1.145x 452.512 392.224 1.154x (16384, 2048, 7168, 2) MoEScalingType.MXFP8 1628.64 1522.7 1.07x 468 402.432 1.163x (16384, 2048, 7168, 4) MoEScalingType.MXFP8 1564.16 1437.3 1.088x 398.272 392.448 1.015x (16384, 2048, 7168, 8) MoEScalingType.MXFP8 1478.3 1647.55 0.897x 416.8 420.032 0.992x (128000, 2048, 7168, 1) MoEScalingType.MXFP8 13811.2 11483.7 1.203x 3793.09 3032.96 1.251x (128000, 2048, 7168, 2) MoEScalingType.MXFP8 12086.2 11340.2 1.066x 3795.1 3009.82 1.261x (128000, 2048, 7168, 4) MoEScalingType.MXFP8 12410.9 10389.6 1.195x 3529.3 2807.25 1.257x (128000, 2048, 7168, 8) MoEScalingType.MXFP8 14126 9803.76 1.441x 3377.52 2585.7 1.306x 

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3129

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• An error occurred (AccessDenied) when calling the ... operation

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3129

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• An error occurred (AccessDenied) when calling the ... operation

❌ 1 New Failure

As of commit 168d4b7 with merge base cd21d0e ():

NEW FAILURE - The following job has failed:

• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.