[moe training] fix scaling type bug; refactor distributed tests

[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] fix scaling type bug; refactor distributed tests

Summary

• MoEScalingType was not getting plumbed through properly in cases where torch_dispatch args param consists of a list of lists, rather than a list of tensors/ScaledGroupedMMTensors. This PR fixes the issue.
• Refactor our distributed tests for MoE training to align with the refactor of MoEs in torchtitan (unify moe implementation for llama4 and deepseek_v3 torchtitan#1534)

Test plan

• pytest test/prototype/moe_training/test_training.py -s -k test_moe_float8_training -s
• torchrun --local-ranks-filter=0 --nproc_per_node=8 -m pytest test/prototype/moe_training/test_fsdp.py -s
• torchrun --local-ranks-filter=0 --nproc_per_node=8 -m pytest test/prototype/moe_training/test_tp.py -s
• torchrun --local-ranks-filter=0 --nproc_per_node=8 -m pytest test/prototype/moe_training/test_fsdp_tp.py -s

[pytorch-bot]

🔗 Helpful Links

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

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

[danielvegamyhre]

cc @vkuzo @drisspg for review of this stack