[alexsamardzic]

Stack from ghstack (oldest at bottom):

• -> Wire MoE and MXFP8Linear training paths to CuTeDSL by default #4620
• Add CuTeDSL MXFP8 K-groups scale rearrange kernel #4604
• Add CuTeDSL MXFP8 pad/unpad kernels #4594
• Fix fake metadata for CuTeDSL 32x1 MXFP8 quantization #4589
• Add CuTeDSL MXFP8 M-groups scale rearrange kernel #4585

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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, 5 Unrelated Failures

As of commit 864b33b with merge base 3818749 ():

NEW FAILURE - The following job has 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

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

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

• Run 1xH100 Tests / test (H100, mt-l-x86iamx-22-225-h100, --pre torch torchvision mslk --index-url https://download.p... / linux-job (gh) (trunk failure)
  test/dtypes/test_affine_quantized_float.py::TestAffineQuantizedFloat8Compile::test_per_row_with_float32
• Run 1xL4 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh) (trunk failure)
  test/dtypes/test_affine_quantized_float.py::TestAffineQuantizedFloat8Compile::test_per_row_with_float32
• Run Regression Tests / test-nightly (CPU Nightly (Clang), linux.4xlarge, --pre torch torchvision --index-url https://dow... / linux-job (gh) (trunk failure)
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch torchvision --index-url https://download.py... / linux-job (gh) (trunk failure)
  test/prototype/test_int4_opaque_tensor.py::TestInt4OpaqueTensor::test_module_path_float32_use_hqq_True
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch torchvision --index-url htt... / linux-job (gh) (trunk failure)
  test/prototype/test_quantized_training.py::TestQuantizedTraining::test_int8_weight_only_training_compile_True_device_cuda

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

[alexsamardzic]

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]

Added --calibrate_roofline option to roofline_unified.py. When turned on, the script measures achievable copy bandwidth plus BF16/FP8 GEMM throughput on the current GPU and uses those calibrated values for the roofline model, instead of relying only on the hard-coded B200 assumptions. The output and plot labels now refer to these values as achievable roofline ceilings.

[andrewor14]

Hi @alexsamardzic, is this the PR that's incorrectly wired? Should we close it then?

[alexsamardzic]

Replaced by #4705.

[alexsamardzic]

Hi @alexsamardzic, is this the PR that's incorrectly wired? Should we close it then?

All five on the stack were wrong, I used plain ghstack command to create them. I had to close all of them, and replaced them with another stack, created through ghstack --direct, #4705 is replacement for this one, and the rest is in the same new stack. I apologize for the mess - previously I was creating proper branches for my PRs here, and in this case I used ghstakh out of habit from mostly working with PyTorch repo.

Btw: All five of these PRs are reasonably simple, I've stated test plans, and also benchmarking results showing that the perfomance is unchanged, or better - so I hope we could merge them quickly. The top one, #4705, is the only one to be careful about, as this one switches the default to new CuTeDSL kernels. After we merge this, and then if we see for couple weeks that there are no issues, I'll post another PR to delete legacy kernels and the C++ extension.