Wire MoE and MXFP8Linear training paths to CuTeDSL by default

[alexsamardzic]

(Replaces #4620 due to incorrect ghstack base branch targeting.)

This PR wires the MXFP8 MoE grouped-mm and MXFP8Linear training paths to use the CuTeDSL backend by default, while keeping the legacy CUDA/Triton backend selectable for validation and benchmarking. It adds end-to-end correctness tests for grouped-mm, padded grouped-mm, prequantized MXTensor grouped-mm, and MXFP8Linear fwd/bwd, plus benchmark support for comparing CuTeDSL against the legacy backend. This prepares the next PR to delete torchao._C_mxfp8 and the now-obsolete legacy CUDA/Triton paths.

---

To run tests:

pytest -q \
    test/prototype/moe_training/test_mxfp8_grouped_mm.py::test_mxfp8_grouped_gemm_backend_fwd_bwd \
    test/prototype/moe_training/test_mxfp8_grouped_mm.py::test_mxfp8_grouped_gemm_padded_backend_fwd_bwd \
    test/prototype/moe_training/test_mxfp8_grouped_mm.py::test_mxfp8_grouped_gemm_mxtensor_backend_forward \
    test/prototype/moe_training/test_mxfp8_grouped_mm.py::test_mxfp8_grouped_gemm_with_dq_fwd_bwd \
    test/prototype/moe_training/test_mxfp8_linear.py::test_mxfp8_linear_fwd_bwd_sqnr

---

To run MoE grouped-mm fwd/bwd recipe benchmark:

python benchmarks/prototype/moe_training/mxfp8/roofline_unified.py \
    --backend=both --K=4096 --N=4096 --G=8 --breakdown_M=16384 \
    --outfile_speedup=roofline_speedup_K4096_N4096.csv \
    --outfile_quant_2d=roofline_quant_2d_K4096_N4096.csv \
    --outfile_quant_3d=roofline_quant_3d_K4096_N4096.csv \
    --plot_file=roofline_K4096_N4096.png

Results of benchmark above:

---

To run MXFP8Linear fwd/bwd recipe benchmark:

python scripts/validate_mxfp8_linear_cutedsl.py \
    --M 4096 --N 4096 --K 2048 \
    --cuda-graph-bench --graph-iters 100

---

Would be good to run on multi-GPU machine:

torchrun --nproc-per-node=2 --local-ranks-filter=0 \
    benchmarks/prototype/moe_training/mxfp8/bench_ep_pipeline.py --backend=legacy

torchrun --nproc-per-node=2 --local-ranks-filter=0 \
    benchmarks/prototype/moe_training/mxfp8/bench_ep_pipeline.py --backend=cutedsl

[alexsamardzic]

Stack from ghstack (oldest at bottom):

• -> Wire MoE and MXFP8Linear training paths to CuTeDSL by default #4705
• Add CuTeDSL MXFP8 K-groups scale rearrange kernel #4704
• Add CuTeDSL MXFP8 pad/unpad kernels #4703
• Fix fake metadata for CuTeDSL 32x1 MXFP8 quantization #4702
• Add CuTeDSL MXFP8 M-groups scale rearrange kernel #4701

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [ALERT ONLY] - System Maintenance

❌ 2 New Failures, 1 Unrelated Failure

As of commit da0c29d with merge base 33ba9f7 ():

NEW FAILURES - The following jobs have failed:

• Run 1xH100 Tests / test (H100 2.13, mt-l-x86iamx-22-225-h100, torch==2.13.0 torchvision==0.28.0 mslk==1.3.0 --index-... / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1
• Run 1xH100 Tests / test (H100, mt-l-x86iamx-22-225-h100, --pre torch torchvision mslk --index-url https://download.p... / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

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

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

• Run PT2E GPU Tests / test-rocm (ROCM Nightly, linux.rocm.gpu.gfx942.1, --pre torch torchvision --index-url https://dow... / linux-job (gh) (trunk failure)

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