[moe training] add benchmark script for moe layer

[danielvegamyhre]

Stacked PRs:

• [moe training] integrate rowwise expert quant kernel #2698
• [moe training] add bench script for fp8 rowwise kernels and update autotune configs #2697
• [moe training] add fp8 rowwise kernels for expert weights #2696
• ->[moe training] add benchmark script for moe layer #2671

---

[moe training] add benchmark script for moe layer

Summary

• Add benchmarking script (with profiling) for a standalone MoE layer. This isolation makes perf analysis + optimization easier.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 1 Cancelled Job, 1 Pending, 1 Unrelated Failure

As of commit 3a73857 with merge base 1c96994 ():

NEW FAILURE - The following job has failed:

• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh)
  Process completed with exit code 1.

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

• Run TorchAO Experimental Tests / test-cpu-ops (linux.arm64.2xlarge) (gh)
  ##[error]The operation was canceled.

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https... / linux-job (gh) (trunk failure)
  test/dtypes/test_affine_quantized_float.py::TestAffineQuantizedFloat8Compile::test_expected_kernels_on_gpu_granularity1_torch_compile_mode_reduce-overhead

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

[danielvegamyhre]

cc @vkuzo @drisspg can you review this stack please

[vkuzo]

nit: move to benchmarks/prototype/moe_training, to match existing location of benchmarks

[danielvegamyhre]

done, moved all moe training benchmark scripts in this PR stack there

[danielvegamyhre]

confirmed test failures unrelated to this change