[moe training] add bench script for fp8 rowwise kernels and update autotune configs

[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 bench script for fp8 rowwise kernels and update autotune configs

Performance vs torch.compile

It's faster for llama4 shape (16, 5120, 4*5120), but slower for skinny shapes.

input_shape            torch_time_us    triton_time_us
-------------------  ---------------  ----------------
(8, (4096, 1024))              92.32           113.248
(16, (20480, 5120))          5746.62          4075.38

There is more we can do, for example writing row major outputs was roughly 2x faster, but we need the outputs in col-major. I can probably look with NCU and figure out what's going on but for now this is a start.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[vkuzo]

we have so many of these, maybe reuse in a separate PR?