Add CuTeDSL MXFP8 pad/unpad kernels

[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/4594

• 📄 Preview Python docs built from this PR

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

❌ 5 New Failures, 3 Unrelated Failures

As of commit 3db0406 with merge base 3818749 ():

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 2
• Run Regression Tests / test (CPU 2.13, linux.4xlarge, torch==2.13.0 torchvision==0.28.0 --index-url https://download.pyt... / linux-job (gh)
  RuntimeError: Command docker exec -t dfae9894d0da980fd918fc8cf09f0a5ee5aa99a994c1eb357fb99267126aa740 /exec failed with exit code 2
• Run Regression Tests / test (CUDA 2.13, linux.g5.12xlarge.nvidia.gpu, torch==2.13.0 torchvision==0.28.0 --index-url http... / linux-job (gh)
  RuntimeError: Command docker exec -t 05de08557bdfcb88109ba28738075cc010c97ff9f3b8a4045f88f8b74041861b /exec failed with exit code 2
• Run Regression Tests / test-nightly (CPU Nightly (Clang), linux.4xlarge, --pre torch torchvision --index-url https://dow... / linux-job (gh)
  RuntimeError: Command docker exec -t 306f54180fc9f6d7cd6f7f2a1bdc8f1bc05273d211749dfbc2e3bb2129c5a0ee /exec failed with exit code 2
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch torchvision --index-url https://download.py... / linux-job (gh)
  RuntimeError: Command docker exec -t 77bfaa3b38a71309a98d5605284070067fe409151bf187632d340ee313e37c8e /exec failed with exit code 2

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch torchvision --index-url htt... / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

BROKEN TRUNK - The following jobs failed but was 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

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

[alexsamardzic]

PR contains CuTeDSL implementations of the MXFP8 pad/unpad kernels and correspondign tests and validate/benchmark script comparing them against the existing CUDA and PyTorch paths. The kernels performance is on par with CUDA kernels performance, see below.

To test:

pytest -q \
    test/prototype/moe_training/test_kernels.py::test_cutedsl_pad_token_groups \
    test/prototype/moe_training/test_kernels.py::test_cutedsl_unpad_token_groups

Benchmarking results, for input dims derived from some realistic models

Tokens x dim	Pad CUDA us	Pad CuTeDSL us	Pad CuTeDSL speedup	Unpad CUDA us	Unpad CuTeDSL us	Unpad CuTeDSL speedup
131072x2048	301.53	299.05	1.01	161.72	161.49	1.00
131072x4096	599.37	594.71	1.01	324.08	323.91	1.00
131072x5120	744.63	740.33	1.01	402.78	402.60	1.00
131072x7168	1039.60	1035.83	1.00	564.11	563.94	1.00
131072x8192	1190.37	1186.81	1.00	647.58	647.18	1.00
131072x10240	1484.48	1481.30	1.00	808.21	807.37	1.00
131072x12288	1781.94	1778.70	1.00	971.44	970.20	1.00

Benchmarking command to produce the table above

{
  for s in 131072,2048 131072,4096 131072,5120 131072,7168 131072,8192 131072,10240 131072,12288 131072,16384; do
    IFS=, read -r t d <<< "$s"
    PYTHONPATH=$PWD python scripts/validate_pad_unpad_cutedsl.py \
      --tokens "$t" --dim "$d" --groups 8 --multiple-of 1 \
      --alignment-size 32 --mode both \
      --cuda-graph-bench --graph-iters 1000 2>&1
  done
} | perl -ne '
BEGIN {
  print "| Tokens x dim | Pad CUDA us | Pad CuTeDSL us | Pad CuTeDSL speedup | Unpad CUDA us | Unpad CuTeDSL us | Unpad CuTeDSL speedup |\n";
  print "|---:|---:|---:|---:|---:|---:|---:|\n";
}
if (/=False/) { die "correctness failure: $_"; }
if (/shape=\((\d+),\s*(\d+)\)/) { $shape = "$1x$2"; }
elsif (/^pad_cuda_us=([\d.]+)/) { $pad_cuda = $1; }
elsif (/^pad_cutedsl_us=([\d.]+)/) { $pad_cutedsl = $1; }
elsif (/^unpad_cuda_us=([\d.]+)/) { $unpad_cuda = $1; }
elsif (/^unpad_cutedsl_us=([\d.]+)/) {
  $unpad_cutedsl = $1;
  printf "| %s | %.2f | %.2f | %.2f | %.2f | %.2f | %.2f |\n",
    $shape,
    $pad_cuda, $pad_cutedsl, $pad_cuda / $pad_cutedsl,
    $unpad_cuda, $unpad_cutedsl, $unpad_cuda / $unpad_cutedsl;
}
'

Benchmarking results, for input dims being powers of two

Tokens x dim	Pad CUDA us	Pad CuTeDSL us	Pad CuTeDSL speedup	Unpad CUDA us	Unpad CuTeDSL us	Unpad CuTeDSL speedup
128x32	6.16	4.10	1.50	4.10	4.10	1.00
128x64	6.17	4.10	1.50	4.10	4.10	1.00
128x128	6.19	4.11	1.51	4.02	4.08	0.99
128x256	6.18	4.12	1.50	3.75	4.00	0.94
128x512	6.15	6.15	1.00	4.10	4.10	1.00
128x1024	6.15	6.14	1.00	4.10	4.10	1.00
128x2048	8.20	6.15	1.33	4.11	4.13	1.00
128x4096	8.20	6.15	1.33	4.10	4.10	1.00
128x8192	10.25	8.20	1.25	6.15	6.15	1.00
128x16384	14.35	12.30	1.17	8.22	8.20	1.00
256x32	6.16	4.37	1.41	4.12	4.10	1.00
256x64	6.16	5.80	1.06	4.08	4.00	1.02
256x128	6.16	5.73	1.08	4.07	4.01	1.01
256x256	6.15	4.56	1.35	4.05	3.84	1.05
256x512	6.16	6.15	1.00	4.10	4.08	1.00
256x1024	6.15	6.15	1.00	4.10	4.10	1.00
256x2048	8.20	6.15	1.33	4.10	4.10	1.00
256x4096	8.20	6.15	1.33	4.13	4.11	1.00
256x8192	10.32	10.24	1.01	6.15	6.15	1.00
256x16384	16.38	14.35	1.14	10.20	8.96	1.14
512x32	6.16	6.15	1.00	4.06	4.02	1.01
512x64	6.16	6.15	1.00	4.07	4.06	1.00
512x128	6.16	6.14	1.00	4.07	4.06	1.00
512x256	6.15	6.15	1.00	4.10	4.10	1.00
512x512	6.15	6.15	1.00	4.10	4.10	1.00
512x1024	8.21	6.15	1.33	4.09	4.10	1.00
512x2048	8.19	6.15	1.33	4.10	4.10	1.00
512x4096	10.24	8.19	1.25	6.10	5.97	1.02
512x8192	12.29	10.25	1.20	6.17	6.15	1.00
512x16384	18.43	16.40	1.12	10.24	10.19	1.00
1024x32	6.18	6.14	1.01	4.11	4.09	1.00
1024x64	6.16	6.15	1.00	4.10	4.09	1.00
1024x128	6.15	6.15	1.00	4.10	4.10	1.00
1024x256	6.15	6.17	1.00	4.08	4.10	1.00
1024x512	8.20	6.15	1.33	4.10	4.10	1.00
1024x1024	8.20	6.15	1.33	4.10	4.10	1.00
1024x2048	8.20	6.15	1.33	4.11	4.10	1.00
1024x4096	10.27	8.20	1.25	6.15	6.15	1.00
1024x8192	14.35	14.01	1.02	6.15	6.15	1.00
1024x16384	24.74	23.81	1.04	10.24	10.25	1.00
2048x32	6.16	6.15	1.00	4.01	4.07	0.99
2048x64	6.15	6.15	1.00	4.10	4.10	1.00
2048x128	6.15	6.15	1.00	4.11	4.09	1.00
2048x256	6.15	6.15	1.00	4.07	4.08	1.00
2048x512	8.20	6.17	1.33	4.10	4.10	1.00
2048x1024	8.19	6.15	1.33	4.10	4.10	1.00
2048x2048	10.24	8.20	1.25	4.50	4.10	1.10
2048x4096	14.34	12.30	1.17	6.15	6.15	1.00
2048x8192	22.52	20.48	1.10	10.24	10.24	1.00
2048x16384	46.87	45.04	1.04	25.70	25.79	1.00
4096x32	6.15	6.15	1.00	4.10	4.10	1.00
4096x64	6.15	6.15	1.00	4.10	4.13	0.99
4096x128	6.18	6.15	1.00	4.10	4.10	1.00
4096x256	8.20	6.15	1.33	4.11	4.10	1.00
4096x512	8.20	6.15	1.33	4.12	4.10	1.00
4096x1024	10.24	8.21	1.25	4.62	6.14	0.75
4096x2048	14.34	12.29	1.17	6.15	6.15	1.00
4096x4096	21.89	20.47	1.07	10.24	10.24	1.00
4096x8192	43.15	40.98	1.05	22.56	22.54	1.00
4096x16384	87.13	84.71	1.03	50.93	50.78	1.00
8192x32	8.19	6.15	1.33	4.10	4.10	1.00
8192x64	8.19	6.15	1.33	4.10	4.10	1.00
8192x128	8.19	6.15	1.33	4.10	4.10	1.00
8192x256	8.20	8.19	1.00	4.10	4.10	1.00
8192x512	10.24	8.20	1.25	4.27	6.14	0.70
8192x1024	14.33	12.29	1.17	6.17	6.15	1.00
8192x2048	20.50	18.45	1.11	10.21	8.83	1.16
8192x4096	42.66	40.95	1.04	22.54	22.52	1.00
8192x8192	86.00	83.75	1.03	49.30	49.15	1.00
8192x16384	165.54	163.85	1.01	95.37	95.73	1.00
16384x32	8.20	8.20	1.00	6.15	6.15	1.00
16384x64	8.20	8.20	1.00	6.15	6.16	1.00
16384x128	8.24	8.20	1.00	6.15	6.15	1.00
16384x256	10.25	10.25	1.00	6.17	6.15	1.00
16384x512	14.34	12.31	1.16	6.15	6.15	1.00
16384x1024	20.48	18.45	1.11	8.20	8.25	0.99
16384x2048	41.00	39.19	1.05	22.24	21.57	1.03
16384x4096	84.73	82.38	1.03	45.61	45.39	1.00
16384x8192	161.36	159.03	1.01	88.20	87.84	1.00
16384x16384	314.99	312.02	1.01	171.34	170.79	1.00

[alexsamardzic]

Replaced by #4703.