[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] use custom ops instead of wrap_triton 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.

❌ 1 Cancelled Job

As of commit 43ffe7a with merge base 317179e ():

CANCELLED JOB - The following job was cancelled. Please retry:

• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https... / linux-job (gh)
  ##[error]The operation was canceled.

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