[moe training] refactor to share benchmarking and profiling utils

[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] refactor to share benchmarking and profiling utils

MoE layer FSDP fwd + bwd bench

rm -rf /tmp/torchinductor_danvm; torchrun --nproc-per-node=8 --local-ranks-filter=0 benchmarks/prototype/moe_training/benchmark_moe_fsdp.py --compile  --recipe="fp8_rowwwise"

...
BF16 time: 100860.36600000001 us
Scaled time: 110353.22200000001 us
Speedup: 0.914x

Scaled grouped mm with dynamic quant fwd + bwd bench

rm -rf /tmp/torchinductor_danvm; torchrun --nproc-per-node=8 --local-ranks-filter=0 benchmarks/prototype/moe_training/benchmark_scaled_grouped_mm_dq.py --compile 

...

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/2767

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