Add CuTeDSL MXFP8 pad/unpad kernels

[alexsamardzic]

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

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_fused_pad_token_groups \
    test/prototype/moe_training/test_kernels.py::test_fused_unpad_token_groups

Benchmarking results, for input dims derived from some realistic models

Tokens x Dim	Groups	Pad CUDA us	Pad CuTeDSL us	Pad CuTeDSL GB/s	Pad CuTeDSL speedup	Unpad CUDA us	Unpad CuTeDSL us	Unpad CuTeDSL GB/s	Unpad CuTeDSL speedup
131072x2048	8	302.72	299.41	3589.8	1.01	161.67	161.64	6649.1	1.00
131072x4096	8	601.51	598.63	3590.8	1.00	324.53	324.80	6618.2	1.00
131072x5120	8	747.95	745.83	3602.7	1.00	402.61	403.78	6654.6	1.00
131072x7168	8	1045.47	1043.78	3604.0	1.00	564.16	565.70	6649.7	1.00
131072x8192	8	1197.17	1197.00	3591.6	1.00	647.42	649.40	6620.2	1.00
131072x10240	8	1490.91	1488.15	3611.2	1.00	808.43	810.08	6633.9	1.00
131072x12288	8	1785.78	1785.47	3611.8	1.00	971.25	974.48	6617.6	1.00
131072x16384	8	N/A	2384.09	3606.5	N/A	1298.22	1302.70	6600.4	1.00

Benchmarking command to produce the table above

python benchmarks/prototype/moe_training/mxfp8/bench_pad_unpad_token_groups.py \
    131072x2048x8 \
    131072x4096x8 \
    131072x5120x8 \
    131072x7168x8 \
    131072x8192x8 \
    131072x10240x8 \
    131072x12288x8 \
    131072x16384x8 \
    --alignment-size 32 --multiple-of 1 --cuda-graph-bench --graph-iters 1000

Benchmarking results, for input dims being powers of two

Tokens x Dim	Groups	Pad CUDA us	Pad CuTeDSL us	Pad CuTeDSL GB/s	Pad CuTeDSL speedup	Unpad CUDA us	Unpad CuTeDSL us	Unpad CuTeDSL GB/s	Unpad CuTeDSL speedup
128x32	8	6.15	4.10	8.0	1.50	4.10	4.10	8.0	1.00
128x64	8	6.15	4.10	16.0	1.50	4.10	4.10	16.0	1.00
128x128	8	6.15	4.11	31.9	1.50	4.10	4.10	32.0	1.00
128x256	8	6.15	4.10	63.9	1.50	4.10	4.10	64.0	1.00
128x512	8	6.15	6.15	85.3	1.00	4.10	4.10	127.9	1.00
128x1024	8	6.15	6.14	170.7	1.00	4.10	4.10	255.9	1.00
128x2048	8	8.19	6.15	341.2	1.33	4.10	4.10	511.7	1.00
128x4096	8	8.20	6.15	682.0	1.33	6.14	4.85	864.6	1.27
128x8192	8	10.25	8.33	1007.6	1.23	6.15	6.15	1364.5	1.00
128x16384	8	14.47	12.30	1364.5	1.18	9.87	8.33	2014.5	1.19
256x32	8	6.15	4.10	12.0	1.50	4.10	4.10	12.0	1.00
256x64	8	6.15	5.67	17.3	1.09	4.10	4.10	24.0	1.00
256x128	8	6.15	6.14	32.0	1.00	4.10	4.10	48.0	1.00
256x256	8	6.42	4.10	95.9	1.57	4.10	4.10	95.9	1.00
256x512	8	8.19	6.15	128.0	1.33	4.10	4.10	191.9	1.00
256x1024	8	8.19	6.15	255.9	1.33	4.10	4.10	383.7	1.00
256x2048	8	8.19	6.15	511.9	1.33	4.10	4.10	767.1	1.00
256x4096	8	8.43	6.15	1022.8	1.37	6.15	4.10	1534.6	1.50
256x8192	8	10.33	10.24	1228.9	1.01	6.15	6.15	2045.9	1.00
256x16384	8	16.39	14.33	1755.7	1.14	10.24	10.24	2457.8	1.00
512x32	8	6.15	6.15	13.3	1.00	4.10	4.10	20.0	1.00
512x64	8	6.15	6.15	26.7	1.00	4.10	4.10	40.0	1.00
512x128	8	7.27	4.10	79.9	1.77	4.10	4.10	79.9	1.00
512x256	8	6.96	4.10	159.7	1.70	4.10	4.10	159.9	1.00
512x512	8	8.19	6.15	213.3	1.33	4.10	4.10	319.7	1.00
512x1024	8	8.19	6.15	426.5	1.33	4.10	4.10	639.6	1.00
512x2048	8	8.20	6.15	852.5	1.33	4.10	4.10	1278.6	1.00
512x4096	8	10.24	8.19	1280.0	1.25	6.15	6.14	1706.7	1.00
512x8192	8	12.30	10.25	2046.5	1.20	6.15	6.15	3410.2	1.00
512x16384	8	18.44	16.39	2558.3	1.12	10.24	10.24	4097.9	1.00
1024x32	8	6.15	6.15	24.0	1.00	4.10	4.10	36.0	1.00
1024x64	8	8.19	4.10	71.9	2.00	4.10	4.10	72.0	1.00
1024x128	8	8.19	4.10	143.8	2.00	4.10	4.10	143.9	1.00
1024x256	8	8.19	6.15	191.9	1.33	4.10	4.10	287.9	1.00
1024x512	8	8.19	6.15	383.9	1.33	4.10	4.10	575.3	1.00
1024x1024	8	8.19	6.15	767.8	1.33	4.10	4.10	1151.3	1.00
2048x32	8	8.08	6.15	45.3	1.32	4.10	4.10	68.0	21:35 [1/17223]
2048x128	8	8.19	6.14	181.3	1.33	4.10	4.10	271.8	1.00
2048x256	8	6.15	6.15	362.5	1.00	4.10	4.10	543.7	1.00
2048x512	8	8.19	6.15	725.1	1.33	4.10	4.10	1087.2	1.00
2048x1024	8	8.19	6.15	1449.4	1.33	4.10	4.10	2174.5	1.00
2048x2048	8	10.25	8.20	2174.3	1.25	6.09	6.06	2943.6	1.01
2048x4096	8	14.33	12.29	2900.8	1.17	6.15	6.15	5797.3	1.00
2048x8192	8	22.52	20.47	3482.6	1.10	10.24	10.24	6963.6	1.00
2048x16384	8	47.03	45.04	3166.2	1.04	26.10	25.01	5701.2	1.04
4096x32	8	8.19	6.15	88.0	1.33	4.10	4.10	131.9	1.00
4096x64	8	8.20	6.15	176.0	1.33	4.10	4.10	263.8	1.00
4096x128	8	6.16	6.15	351.9	1.00	4.10	4.10	527.6	1.00
4096x256	8	8.19	6.15	703.9	1.33	4.10	4.10	1053.9	1.00
4096x512	8	8.20	6.15	1406.6	1.33	4.10	4.10	2109.2	1.00
4096x1024	8	10.24	8.20	2110.4	1.25	5.01	6.14	2817.0	0.82
4096x2048	8	14.33	12.29	2816.3	1.17	6.15	6.15	5627.9	1.00
4096x4096	8	21.88	20.48	3380.0	1.07	10.24	10.24	6758.7	1.00
4096x8192	8	43.15	41.01	3375.1	1.05	22.62	22.36	6190.8	1.01
4096x16384	8	88.35	85.39	3242.0	1.03	51.06	50.40	5492.5	1.01
8192x32	8	8.19	6.15	173.2	1.33	4.10	4.10	259.6	1.00
8192x64	8	8.19	6.15	346.3	1.33	4.10	4.10	519.4	1.00
8192x128	8	8.19	6.15	692.7	1.33	4.10	4.10	1039.0	1.00
8192x256	8	8.19	8.19	1040.0	1.00	4.10	4.10	2077.6	1.00
8192x512	8	10.24	8.20	2078.4	1.25	6.03	6.15	2772.4	0.98
8192x1024	8	14.33	12.29	2773.7	1.17	6.15	6.15	5541.7	1.00
8192x2048	8	20.49	18.44	3695.7	1.11	10.24	8.20	8310.6	1.25
8192x4096	8	43.00	40.55	3361.5	1.06	22.54	22.52	6052.6	1.00
8192x8192	8	86.03	83.94	3247.8	1.02	49.17	49.13	5549.1	1.00
8192x16384	8	166.25	164.14	3321.9	1.01	95.08	96.28	5663.5	0.99
16384x32	8	8.22	8.20	257.8	1.00	6.15	6.15	343.8	1.00
16384x64	8	8.20	8.20	515.6	1.00	6.15	6.15	687.6	1.00
16384x128	8	10.24	9.20	919.2	1.11	6.15	6.15	1374.7	1.00
16384x256	8	10.25	10.24	1651.2	1.00	6.15	6.15	2749.4	1.00
16384x512	8	14.33	12.41	2725.8	1.16	6.15	7.08	4778.4	0.87
16384x1024	8	20.49	18.43	3669.0	1.11	8.20	8.25	8202.6	0.99
16384x2048	8	41.45	39.03	3465.8	1.06	21.55	20.53	6588.9	1.05
16384x4096	8	84.77	82.02	3298.2	1.03	46.22	45.87	5898.1	1.01
16384x8192	8	161.87	159.24	3397.7	1.02	88.15	87.67	6171.5	1.01
16384x16384	8	316.10	312.69	3460.8	1.01	171.73	171.48	6310.5	1.00

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 78e70ef with merge base 9d0c494 ():
💚 Looks good so far! There are no failures yet. 💚

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

[vkuzo]

split to test and benchmarks

[vkuzo]

are there tests we can extend instead of adding new ones?

[vkuzo]

prefix with _

[vkuzo]

prefix with _

[alexsamardzic]

All comments addressed.