[Common] Support compact unswizzled scales in MXFP8 quantization

[XFDG]

Description

MXFP8Tensor.quantize_() accepts caller-provided unswizzled scale_inv buffers, including compact ones that hold one scale per block ((rows, cols/32) rowwise, (rows/32, cols) columnwise) instead of TE's 128x4-padded shape. #3518 asks GEMM to accept such scales too (#3519). The quantize kernels, however, assumed the padded layout, so quantizing into compact buffers wrote outside them or produced wrong scales.

Reproduced on current main (eb3c860) on B200: the reporter's script prints 256 256 instead of 0 0. A sweep over 9 shapes × FP32/BF16/FP16 inputs × rowwise/columnwise/both × 1D/2D scaling, plus E5M2, passed 53 of 180 cases. Compact views 1 or 4 bytes past 16-byte alignment raised cudaErrorMisalignedAddress. The FP8 data was always correct; the scales were wrong in three kernels:

Kernel	Used for	Failure with compact scales
Generic quantize_mxfp8_kernel	FP32 input, 2D scaling, columnwise-only, partial tiles	Rowwise stores are bounded by column only and columnwise stores not at all. They write past the buffer when rows % 64 != 0 and into the next scale row when cols % 64 != 0.
Staged bidimensional cast (CastTraitsSwizzle)	BF16/FP16, rowwise + columnwise	Rowwise scales are flushed in 4- or 8-byte groups, so rows that are not a multiple of four scales long are dropped or misplaced. Columnwise scales are stored for whole 64-column iterations.
Register-resident rowwise and bidimensional casts	BF16, cols % 128 == 0 / % 256 == 0	64- and 128-bit stores fault on views that are not 16-byte aligned.

So BF16 inputs stay in bounds only for some shapes.

Review of the first commit found that views whose scale rows hold more entries than blocks also fail in the specialized kernels. The register-resident bidimensional kernel faults when columnwise scale rows are not a multiple of 16 bytes long. The rowwise-only kernels reject such rows (register-resident) or leave some of them unwritten (staged).

Fixes #3689

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Generic kernel: bound rowwise scale stores by the scale tensor's rows, and columnwise stores by its rows and columns. TE's padded tensors cover every chunk, so their stores are unchanged.
• Kernel selection: use the specialized kernels only when the scale pointers are 16-byte aligned, the rowwise scale stride is a multiple of four, rowwise-only casts write packed scale rows (one entry per block), and, for the staged bidimensional kernel, each columnwise scale row spans whole 64-column iterations. Other layouts use the generic kernel. The register-resident bidimensional kernel also needs columnwise scale rows that are a multiple of 16 bytes long; other row lengths take the staged kernel. TE's padded tensors always qualify, so their kernel selection is unchanged.
• Validation: reject unswizzled scale tensors that are not 2D or hold fewer than one scale per block, which would otherwise leave scales unwritten. GEMM-swizzled scales are not affected.
• Tests: tests/pytorch/mxfp8/test_mxfp8_quantize_compact_scales.py and CompactScalesCastMXFP8TestSuite in tests/cpp/operator/test_cast_mxfp8.cu quantize into compact views, and into views whose scale rows hold more entries than blocks. Canary bytes surround every view, and the views are aligned and unaligned. The tests compare the results byte for byte with TE's padded output. The C++ suite also covers FP16 input and the dbias variant.

GEMM's handling of compact scales (#3519) is unchanged.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

Testing on B200 (CUDA 13.1, PyTorch 2.11.0a0)

PyTorch-only editable build for sm_100a with NVTE_WITH_NCCL_EP=0 and NVTE_WITH_CUTEDSL=0. Each run checked that the imported libraries came from this build.

Check	main (eb3c860)	First commit (9cae9e1)	This PR
Reporter's script	256 256	0 0	0 0
Sweep with aligned / +4 B / +1 B views	53/180 / misaligned address / misaligned address	180/180 each	180/180 each
PyTorch test: compact views, 218 cases	174 failed	218 passed	218 passed
PyTorch test: longer scale rows, 108 cases		11 failed	108 passed
C++ suite, 1008 cases	130 of the first commit's 144 cases failed	13 failed	1008 passed
All suites in test_cast_mxfp8		940 passed, 488 skipped	1804 passed, 488 skipped
tests/pytorch/mxfp8		665 passed, 1035 skipped	773 passed, 1035 skipped
test_mxfp8_2d_quantize.py		135 passed	135 passed
test_quantized_tensor.py		390 passed, 200 skipped	390 passed, 200 skipped
test_numerics.py -k mxfp8		81 passed	81 passed
test_fusible_ops.py -k mxfp8		364 passed, 445 skipped	364 passed, 445 skipped
test_weight_swizzle_in_layers.py		52 passed	52 passed
test_partial_cast.py		1 passed	1 passed

Not covered: JAX (not installed; JAX passes padded MXFP8 scales, which keep the same stores and kernels), the CuTeDSL backend (not built), and performance (not measured).

The C++ operator tests on main do not compile with CUDA 13.1. Since #3325, DType::kFloat8UE5M3 exists for every CUDA version, but TypeInfo::types lists fp8ue5m3 only for CUDA 13.4 and later. I built test_cast_mxfp8.cu against a local copy of test_common.h that ends the type lookup there; that workaround is not part of this PR.

🤖 Generated with Claude Code

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; the previous alignment issue is fixed and no new actionable issues were found.

Summary

Adds support for compact, unswizzled MXFP8 scale buffers.

• MXFP8 quantization now writes scales into compact unswizzled buffers.

The previous alignment finding is fixed. No new actionable issues were found. GPU tests were not run during this review.

_{Reviews (2) · Last reviewed commit: "[Common] Check MXFP8 scale row lengths b..." · Reviewed by Greptile}

[wujingyue]

cc @vthumbe1503

FYI, this PR looks a bit vibey but you might want to steer it?