[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] 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.

✅ No Failures

As of commit f4c270f with merge base 143c3a6 ():
💚 Looks good so far! There are no failures yet. 💚

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

[zou3519]

cc @davidberard98 too