[moe training] integrate rowwise expert quant kernel

[danielvegamyhre]

Stacked PRs:

• ->[moe training] integrate rowwise expert quant kernel #2698
• [moe training] add bench script for fp8 rowwise kernels and update autotune configs #2697
• [moe training] add fp8 rowwise kernels for expert weights #2696
• [moe training] add benchmark script for moe layer #2671
• [moe training] add llama4 benchmarking script #2669
• [moe training] use smaller block sizes for per group scaling kernels to improve perf #2668

---

[moe training] integrate rowwise expert quant kernel

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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:

• ghstack-mergeability-check and Check labels failing with 'Resource not accessible by integration'

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

[danielvegamyhre]

@zou3519 here is a different problem occurring with wrap_triton in this PR. I added a new triton kernel _triton_fp8_rowwise_3d_transpose_scales_rhs_kernel and wrapped with wrap_triton, and wrapped that in a custom triton op triton_fp8_rowwise_3d_transpose_rhs.

In eager it works, but when I try to use torch.compile now I get an error about using a FunctionalTensor. This doesn't happen with the other triton kernels used for this feature.

Repro

1. Check out this PR
2. Run unit test using compile: pytest test/prototype/moe_training/test_training.py -k test_moe_float8_training[True-target_fqns0]

Error

• link to output