Add CuTeDSL MXFP8 M-groups scale rearrange kernel

[alexsamardzic]

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

To test:

pytest -q test/prototype/moe_training/test_kernels.py::test_cutedsl_mx_block_rearrange_2d_M_groups

This PR adds a CuTeDSL implementation of the existing CUDA mx_block_rearrange_2d_M_groups kernel, as well as validation/benchmark script that compares it against the current CUDA and Triton implementations. The production call sites are unchanged in this PR; the CUDA and Triton variants will be deactivated and deleted in a follow-up.

Benchmarking results

Scales RxC	Groups	Triton us	CUDA us	CuTeDSL us	CuTeDSL GB/s	CuTeDSL speedup
131072x37	8	92.06	14.33	10.25	988.9	8.98
131072x45	8	67.62	14.34	12.29	995.5	5.50
131072x64	8	102.35	12.28	10.24	1644.9	10.00
131072x160	8	110.65	22.64	20.47	2056.6	5.40
131072x224	8	131.01	26.75	24.77	2379.5	5.29
131072x320	8	170.82	36.88	32.77	2570.2	5.21
131072x512	8	221.46	55.36	53.24	2530.9	4.16
131072x768	8	298.79	80.75	79.75	2534.3	3.75
131072x1024	8	362.33	103.43	103.82	2595.7	3.49
131072x1025	8	297.41	202.61	159.34	1695.3	1.87
131072x1536	8	460.57	148.23	151.47	2668.7	3.04
131072x2048	8	561.30	194.48	200.57	2687.2	2.80
131072x2049	8	878.67	378.79	307.05	1757.5	2.86
131072x3072	8	860.29	288.97	295.60	2735.0	2.91
131072x4096	8	1039.48	380.25	392.23	2748.2	2.65
131072x4097	8	1034.57	760.03	602.83	1789.2	1.72

Benchmarking command to produce the table above

python benchmarks/prototype/moe_training/mxfp8/bench_mx_block_rearrange_2d_M_groups.py \
    131072x37x8 \
    131072x45x8 \
    131072x64x8 \
    131072x160x8 \
    131072x224x8 \
    131072x320x8 \
    131072x512x8 \
    131072x768x8 \
    131072x1024x8 \
    131072x1025x8 \
    131072x1536x8 \
    131072x2048x8 \
    131072x2049x8 \
    131072x3072x8 \
    131072x4096x8 \
    131072x4097x8 \
    --multiple-of 32 --cuda-graph-bench --graph-iters 1000 --bench-only triton cuda cutedsl

Edit: Some additional benchmarking results:

Benchmarking results, for R/C being powers of two

Scales RxC	Groups	Triton us	CUDA us	CuTeDSL us	CuTeDSL GB/s	CuTeDSL speedup
256x32	8	4.10	6.15	4.10	12.0	1.00
256x64	8	4.10	6.15	6.14	16.0	0.67
256x128	8	4.10	6.15	4.10	47.9	1.00
256x256	8	4.10	6.15	4.10	95.9	1.00
256x512	8	4.10	6.15	6.14	128.0	0.67
256x1024	8	6.14	8.19	6.14	256.0	1.00
256x2048	8	6.15	8.19	6.15	511.8	1.00
256x4096	8	10.24	8.44	8.19	768.1	1.25
256x8192	8	16.38	12.63	10.24	1228.5	1.60
256x16384	8	26.61	22.53	16.38	1536.5	1.62
512x32	8	4.10	6.15	4.10	16.0	1.00
512x64	8	4.10	6.15	4.10	32.0	1.00
512x128	8	4.10	6.15	6.14	42.7	0.67
512x256	8	6.14	6.15	6.14	85.3	1.00
512x512	8	6.14	6.15	6.14	170.7	1.00
512x1024	8	6.15	8.19	6.14	341.3	1.00
512x2048	8	8.19	8.20	6.14	682.6	1.33
512x4096	8	10.25	10.22	8.19	1024.1	1.25
512x8192	8	18.42	14.33	10.25	1637.3	1.80
512x16384	8	32.36	24.57	18.43	1821.1	1.76
1024x32	8	4.10	8.19	4.10	24.0	1.00
1024x64	8	6.14	8.20	6.15	32.0	1.00
1024x128	8	4.12	8.19	6.11	64.4	0.67
1024x256	8	6.14	8.20	6.14	128.0	1.00
1024x512	8	6.15	10.24	6.14	256.0	1.00
1024x1024	8	8.19	8.20	6.14	512.0	1.33
1024x2048	8	10.24	10.24	6.15	1023.3	1.67
1024x4096	8	14.36	12.29	8.20	1535.1	1.75
1024x8192	8	24.57	16.39	14.33	1756.0	1.71
1024x16384	8	45.07	28.70	22.62	2225.1	1.99
2048x32	8	6.14	8.20	4.10	40.0	1.50
2048x64	8	8.19	10.24	6.14	53.4	1.33
2048x128	8	6.14	10.24	6.14	106.7	1.00
2048x256	8	8.19	10.24	6.15	213.3	1.33
2048x512	8	6.15	10.24	6.15	426.6	1.00
2048x1024	8	8.60	10.25	6.15	852.8	1.40
2048x2048	8	14.43	12.29	8.19	1280.1	1.76
2048x4096	8	23.14	14.33	12.29	1707.1	1.88
2048x8192	8	40.56	22.55	18.44	2274.2	2.20
2048x16384	8	73.00	39.19	34.84	2407.9	2.10
4096x32	8	8.19	8.19	4.10	71.9	2.00
4096x64	8	6.15	10.24	6.14	96.0	1.00
4096x128	8	9.97	10.24	6.15	192.0	1.62
4096x256	8	10.24	10.24	6.14	383.9	1.67
4096x512	8	10.24	10.25	6.15	767.5	1.67
4096x1024	8	14.35	10.28	8.19	1152.1	1.75
4096x2048	8	22.54	14.32	10.28	1836.2	2.19
4096x4096	8	40.93	20.49	16.39	2302.9	2.50
4096x8192	8	70.09	34.86	30.71	2458.6	2.28
4096x16384	8	135.22	66.78	63.50	2377.8	2.13
8192x32	8	8.20	8.20	4.10	135.9	2.00
8192x64	8	10.24	10.24	6.14	181.3	1.67
8192x128	8	10.25	10.24	6.14	362.6	1.67
8192x256	8	14.33	10.25	6.14	725.3	2.33
8192x512	8	20.49	10.30	8.19	1088.1	2.50
8192x1024	8	22.54	12.29	10.24	1740.3	2.20
8192x2048	8	41.53	18.44	16.38	2176.6	2.54
8192x4096	8	72.00	30.75	28.66	2487.8	2.51
8192x8192	8	131.69	61.87	59.37	2401.9	2.22
8192x16384	8	251.60	112.87	112.87	2526.9	2.23
16384x32	8	12.28	8.20	6.14	176.0	2.00
16384x64	8	12.29	10.25	6.14	352.0	2.00
16384x128	8	18.43	10.24	6.14	703.9	3.00
16384x256	8	22.54	10.25	8.19	1056.0	2.75
16384x512	8	26.61	12.95	10.24	1689.8	2.60
16384x1024	8	53.30	18.29	16.37	2114.4	3.26
16384x2048	8	73.66	28.81	27.18	2545.7	2.71
16384x4096	8	144.68	58.05	56.09	2467.7	2.58
16384x8192	8	253.98	106.86	108.48	2551.8	2.34
16384x16384	8	489.53	206.37	209.55	2642.0	2.34

Benchmarking results, for R/C derived from some realistic (Mixtral, Qwen, Deepseek, Llama) models

Scales RxC	Groups	Triton us	CUDA us	CuTeDSL us	CuTeDSL GB/s	CuTeDSL speedup
4096x64	8	6.15	10.24	6.14	96.0	1.00
4096x128	8	6.15	10.24	6.14	192.0	1.00
4096x160	8	8.19	10.24	6.14	240.0	1.33
4096x224	8	8.20	10.25	6.14	335.9	1.33
4096x256	8	8.19	10.24	6.14	384.0	1.33
4096x384	8	8.20	10.25	6.14	575.9	1.33
4096x512	8	10.24	10.25	6.15	767.4	1.67
16384x64	8	18.43	10.24	6.14	352.0	3.00
16384x128	8	23.77	10.24	6.15	703.7	3.87
16384x160	8	16.39	10.25	6.15	879.6	2.67
16384x224	8	20.47	10.25	6.31	1198.8	3.24
16384x256	8	22.53	10.25	6.58	1315.1	3.43
16384x384	8	20.50	12.29	8.25	1572.6	2.48
16384x512	8	28.66	14.28	10.25	1688.5	2.80
65536x64	8	59.37	10.25	6.15	1375.1	9.66
65536x128	8	72.74	12.29	10.24	1651.4	7.10
65536x160	8	43.28	14.34	12.29	1720.3	3.52
65536x224	8	69.59	16.38	14.34	2062.8	4.85
65536x256	8	63.47	16.39	14.35	2356.1	4.42
65536x384	8	116.04	24.57	20.49	2475.6	5.66
65536x512	8	110.75	28.67	26.61	2541.3	4.16
131072x64	8	126.89	12.30	10.24	1645.1	12.39
131072x128	8	156.86	16.40	14.35	2347.9	10.93
131072x160	8	110.54	22.54	20.47	2057.0	5.40
131072x224	8	177.65	26.64	24.59	2397.3	7.22
131072x256	8	125.14	28.67	26.61	2531.4	4.70
131072x384	8	342.02	42.95	39.29	2572.0	8.70
131072x512	8	280.07	55.30	53.23	2531.4	5.26
262144x64	8	149.58	16.39	14.39	2336.8	10.40
262144x128	8	212.52	28.65	26.61	2526.5	7.99
262144x160	8	649.90	37.23	36.85	2280.8	17.64
262144x224	8	344.17	51.21	50.96	2308.9	6.75
262144x256	8	557.62	55.31	53.22	2526.7	10.48
262144x384	8	590.93	80.08	77.95	2587.8	7.58
262144x512	8	520.05	100.44	101.73	2643.8	5.11

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

• 📄 Preview Python docs built from this PR

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

❌ 2 Cancelled Jobs

As of commit a2acc22 with merge base 9d0c494 ():

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• Run 1xL4 Tests / test (SM-89 2.13, mt-l-x86aavx2-11-41-l4, torch==2.13.0 --index-url https://download.pytorch.org/... / linux-job (gh)
• Run 1xL4 Tests / test (SM-89, mt-l-x86aavx2-11-41-l4, --pre torch --index-url https://download.pytorch.org/whl/nig... / linux-job (gh)

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

[vkuzo]

can we split the logic here to test and benchmarks

[vkuzo]

@claude review

[claude]

Claude finished @vkuzo's task in 1m 47s —— View job

---

Code Review

• Read changed files
• Analyze CuTeDSL kernel implementation
• Analyze test additions
• Review the validate/benchmark script (re: @vkuzo's split comment)
• Post review feedback

---

Overall this is a well-structured CuTeDSL port with solid benchmarking evidence — the kernel is on par with or faster than the CUDA baseline across a wide sweep, and the correctness handling of the partial-column / non-16-multiple paths is careful. A few points below.

Re: @vkuzo's comment — split scripts/validate_m_groups_cutedsl.py into test + benchmark

I agree with the split, and I'd note the two halves already have natural homes:

• Correctness is already covered by test_cutedsl_mx_block_rearrange_2d_M_groups in test/prototype/moe_training/test_kernels.py (compares against the to_blocked reference). The bench_only is None validation block in the script (lines 110-168) is largely redundant with that test. The elaborate --deterministic diff-decoding block (lines 134-168) reads like debugging scaffolding — worth dropping rather than shipping.
• Benchmarking — the timing logic (lines 23-51, 170-197) fits benchmarks/ alongside the benchmarks.utils helper it already imports (see bench_matmul referenced in a recent commit for the established pattern). Consider moving it to something like benchmarks/prototype/moe_training/bench_m_groups_rearrange.py.

So the concrete suggestion is: delete the script, keep the test that already exists, and add a benchmark-only file under benchmarks/.

Kernel correctness / robustness

• _load4 unaligned 4-byte path (lines 124-133): the fast path guards on input_offset % 4 == 0 but does a recast_ptr(..., dtype=Uint32).load() with no .align(4) hint. It works because of the runtime modulo check, but unlike the 16-byte (.align(16)) and 8-byte (.align(8)) paths, the alignment is left implicit. Worth an .align(4) for consistency and to let the compiler emit an aligned load.

• torch.zeros on every call (line 585): allocating + zeroing (padded_rows, padded_cols) per invocation. The kernel writes zeros into partial-column and intra-chunk row-padding regions itself, but row chunks in the over-estimated num_row_chunks range (ceil_div(rows,128) + num_groups) are never launched for the trailing padded rows, so the memset does appear necessary for those. If that's the only reason, a one-line comment saying "zero-init required for unwritten padded rows" would save the next reader the analysis.

• num_row_chunks = ceil_div(rows, 128) + num_groups (line 605): this is a safe over-estimate (each group's ceil-rounding adds at most one extra chunk), and extra blocks early-out via active == 0. Fine, just launches up to num_groups idle blocks. Not a bug, noting it in case a tighter grid is worth it later.

Test coverage

• test_cutedsl_mx_block_rearrange_2d_M_groups (test_kernels.py:284-330) parametrizes only n_groups=8 and never exercises chunk_width other than the auto default. The default picks 64/32/16 based on cols, but the explicit chunk_width in (16, 32, 64, 128) code paths and the 128 branch in particular go untested. Consider adding a couple of cases with an explicit chunk_width (including 128) since the branching logic (ALL_COL_CHUNKS_FULL vs COLS_MULTIPLE_OF_16 vs byte-wise) differs meaningfully per width. Fix this →
• torch.randint(0, 255, ...) in both the test and the script uses an exclusive upper bound of 255, so the e8m0 0xFF (NaN) byte is never generated. Since this kernel only moves bytes it doesn't matter for correctness, but 256 would give full coverage of the byte range.

Minor

• The kernel module is ~650 lines with substantial repetition across the three store paths (ALL_COL_CHUNKS_FULL / full_col_chunk / partial in _process_chunk, lines 241-397). The three blocks writing value0..value3 to smem_offset0..3 are near-identical. Given CuTeDSL's @cute.jit constraints this may be hard to factor, but a small helper for the "store 4 values to the 4 smem offsets" step would cut the duplication noticeably.
• Nice touch using @functools.cache on the compile step keyed by the specialization params — that's the right way to amortize cute.compile.

Nothing here is blocking; the split requested by @vkuzo is the main actionable item, and the .align(4) and explicit-chunk_width test coverage are the two I'd most encourage.
· branch gh/alexsamardzic/10/head

[vkuzo]

is there a test we can extend instead of building a new one?

[vkuzo]

prefix with _

[alexsamardzic]

All comments addressed. Kept the Triton and CuTeDSL tests separate, as the PyTorch reference can't express the CuTeDSL test's partial-column (non-multiple-of-4) cases.