[moe training] update bench script to compare fp8 dynamic quant scaled_grouped_mm fwd+bwd against bf16

[danielvegamyhre]

Stacked PRs:

• [moe training] add memory bandwidth calculations to kernel benchmarking scripts #2769
• [moe training] refactor to share benchmarking and profiling utils #2767
• ->[moe training] update bench script to compare fp8 dynamic quant scaled_grouped_mm fwd+bwd against bf16 #2765
• [moe training] remove duplicate benchmark script #2762
• [moe training] use llama4 shapes for kernel benchmarks #2756
• [moe training] fix scaling type bug; refactor distributed tests #2749
• [moe training] work around wrap_triton bug by using normal custom ops instead for fp8 rowwise kernels  #2734
• [moe training] update tests for torchtitan moe refactor #2733

---

[moe training] update bench script to compare fp8 dynamic quant scaled_grouped_mm fwd+bwd against bf16

Summary

• Previously this benchmarking script only measured fp8 rowwise perf, without comparing against a baseline.
• This script now compares against the bf16 baseline so we can easily see the relative speedup/slowdown.

Example output:

A_shape        B_shape             bf16_time_us    scaled_time_us  fp8_speedup
-------------  ----------------  --------------  ----------------  -------------
(16640, 5120)  (16, 8192, 5120)         10523.7           12889.3  0.816x

[pytorch-bot]

🔗 Helpful Links

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

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