[moe training] add memory bandwidth calculations to kernel benchmarking scripts

[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] add memory bandwidth calculations to kernel benchmarking scripts

documenting example torch.compile codegen for the 3d expert weigh transpose + quant kernel: https://www.internalfb.com/phabricator/paste/view/P1904683419

input_shape            torch_time_us    triton_time_us    torch_mem_bw_gbps    triton_mem_bw_gbps  triton_speedup
-------------------  ---------------  ----------------  -------------------  --------------------  ----------------
(1, (8192, 5120))            114.976           131.392              1823.99               1596.1   0.88x
(1, (5120, 8192))            145.44            131.456              1441.94               1595.33  1.11x
(16, (8192, 5120))          2182.32           1629.92               1537.56               2058.66  1.34x
(16, (5120, 8192))          2081.39           1745.89               1612.12               1921.91  1.19x
(128, (8192, 5120))        21692.4           23648.6                1237.46               1135.1   0.92x
(128, (5120, 8192))        18088.4           23502.3                1484.02               1142.17  0.77x

[pytorch-bot]

🔗 Helpful Links

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

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