[mxfp8 moe training] bench and profile mxfp8 a2a fwd and bwd separately

[danielvegamyhre]

Stacked PRs:

• ->[mxfp8 moe training] bench and profile mxfp8 a2a fwd and bwd separately #3203
• [mxfp8 moe training] improve unit tests #3201
• [mxfp8 moe training] integrate triton quant/dequant kernels into mxfp8 all to all #3197
• [mxfp8 moe training] add triton kernel for mxfp8 dequantization #3195

---

[mxfp8 moe training] bench and profile mxfp8 a2a fwd and bwd separately

Single node benchmark on 4xB200 over NVLink shows perf is flat or slightly slower in this setting.

I expect speedups for multi-node EP though, where IB bandwidth is much more constrained than NVLink.

input_shape         num_splits    fwd_bf16_ms    fwd_mxfp8_ms    bwd_bf16_ms    bwd_mxfp8_ms
----------------  ------------  -------------  --------------  -------------  --------------
(1, 8192, 5120)              4       0.269697        0.479122       0.695028        0.993789
(2, 8192, 5120)              4       0.347715        0.468697       0.791324        0.872646
(4, 8192, 5120)              4       0.593996        0.585684       1.28176         1.24674
(8, 8192, 5120)              4       1.53808         1.03233        2.38809         2.3224
(16, 8192, 5120)             4       1.77031         1.8789         4.36899         4.46669

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 1 Cancelled Job, 1 Unrelated Failure

As of commit 19a4679 with merge base b644211 ():

NEW FAILURES - The following jobs have failed:

• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https... / linux-job (gh)
  RuntimeError: Command docker exec -t 8d8aa7758b140bfd2af2e7e9a6fa4622c1ea6670814ce864443273b74cc094e8 /exec failed with exit code 1
• Run Regression Tests (aarch64) / test-cpu-ops (macos-14) (gh)
  /Users/runner/miniconda3/envs/venv/lib/python3.10/site-packages/executorch/share/cmake/../../include/executorch/runtime/core/memory_allocator.h:143:21: error: no matching function for call to 'mul_overflows'

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

• Run Regression Tests (aarch64) / 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 Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/dtypes/test_nf4.py::TestComm::test_comm

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