[moe training] add fp8 rowwise kernels for expert weights

[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] add fp8 rowwise kernels for expert weights

Summary

• torch.compile is too slow for quantizing expert weights along dim1 for backward (see Inductor codegen for float8 dynamic quantization ops for scaled_grouped_mm backward pass is slow pytorch#159769)
• I wrote a triton kernel to do this with better perf. We can remove it once torch.compile is better.
• This PR is focused on numerical accuracy. Perf benchmarking is in next PR in stack.

Test plan

• pytest test/prototype/moe_training/test_kernels.py

[pytorch-bot]

🔗 Helpful Links

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

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

[vkuzo]

this seems expensive, can we just extract the bits?

[vkuzo]

this is confusing because it sounds like Float8TrainingTensor, maybe name it float8_data?