Fix fake metadata for CuTeDSL 32x1 MXFP8 quantization

[alexsamardzic]

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

To test:

pytest -q test/prototype/moe_training/test_kernels.py -k "cutedsl and 32x1"

This PR also contains a script to validate/benchmark the CuTeDSL 2D 32x1 dim1 quantization path against the existing CUDA path inside the wgrad workflow; the script is temporary and will be deleted when the torchao._C_mxfp8 extension removed.

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

• 📄 Preview Python docs built from this PR

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

❌ 1 Unclassified Failure

As of commit 0c9947a with merge base 9d0c494 ():

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• Run PT2E GPU Tests (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)

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

[vkuzo]

split logic to test and benchmarks

[vkuzo]

@claude review

[claude]

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

---

Claude Code Review

• Read the PR diff
• Review torchao/prototype/moe_training/kernels/mxfp8/quant.py
• Review scripts/validate_wgrad_dim1_cutedsl.py
• Post review

---

I reviewed the fake-metadata fix against the real kernel in cutedsl_quantize_2d_32x1.py. The core fix is correct — the fake op now faithfully mirrors the real op's output metadata.

✅ Correctness of the fake metadata

Field	Real kernel (mxfp8_quantize_cutedsl_2d_32x1)	Fake (this PR)	Match
q_data strides	empty_strided((M,K), (1,M)) — column-major	empty_strided((m,k), (1,m))	✅ (was (k,1), i.e. row-major — this was the bug)
q_data dtype	float8_e4m3fn	float8_e4m3fn	✅
blocked scales	flat (ceil_div(K,128)*128 * ceil_div(m_blocks,4)*4,)	flat (ceil_div(k,128)*128 * ceil_div(m_blocks,4)*4,)	✅ (product is commutative)
non-blocked scales	(K, m_blocks)	(k, m_blocks)	✅ (previously the fake had no non-blocked branch at all)
scales dtype	float8_e8m0fnu	float8_e8m0fnu	✅

The two real problems are addressed: (1) the q_data stride was (k, 1) (row-major) but the kernel produces column-major (1, M), and (2) the fake unconditionally assumed blocked_scale_output=True, so it returned wrong metadata whenever the op was called with blocked_scale_output=False (as the validation script and wgrad path do). Both are now fixed.

🔵 Minor: swapped variable names in the blocked branch

quant.py:1153-1154

padded_scale_rows = ceil_div(m_blocks, 4) * 4
padded_scale_cols = ceil_div(k, 128) * 128

The final tensor is 1-D (padded_scale_cols * padded_scale_rows), so the size is correct regardless. But the naming is inverted relative to the real kernel, where padded_scale_rows is derived from K (rounded to 128) and padded_scale_cols from m_blocks (rounded to 4). Aligning the names with the real kernel would make it easier to keep the two in sync and avoid future confusion. Purely cosmetic — no behavioral impact.

Fix this →

🔵 Minor: dropped explanatory comments

The previous version had short comments explaining the padding rationale ("Only scales tensor needs padding for 4x128 scale factor tiles", "output data has same dimensions as input"). Those were removed. Since the real kernel keeps equivalent comments (cutedsl_quantize_2d_32x1.py:1033-1051), it'd be nice to keep a one-liner here too so the fake stays self-documenting.

🔵 Script: scripts/validate_wgrad_dim1_cutedsl.py

• This matches @vkuzo's inline suggestion to split the logic into test/ and benchmarks/. The file mixes numeric validation (torch.equal, SQNR reporting) with benchmarking (benchmark_function, CUDA-graph timing). Given benchmarks/utils.py already exists and is imported here, moving the timing path under benchmarks/ and the equality check into a test under test/prototype/moe_training/ would fit the repo layout better. The PR body notes the script is temporary (to be deleted when torchao._C_mxfp8 is removed), so this is a judgment call on whether it's worth splitting a throwaway file.
• Minor: run() computes grad_weight.transpose(-2, -1) and the diff-coordinate flattening at lines 209-213 assumes a 3-D output — worth a brief comment that this is the grouped (per-expert) 3-D result, since it's not obvious from the 2-D inputs.

Summary

The metadata fix is correct and well-targeted — the stride and the missing blocked_scale_output=False branch are exactly the discrepancies that would have caused wrong behavior under torch.compile/fake-tensor tracing. Only cosmetic nits remain (variable naming, dropped comments) plus @vkuzo's test/benchmark split suggestion for the script. Nothing blocking on the kernel change itself.
· branch gh/alexsamardzic/11/head

[alexsamardzic]

All comments addressed.