[bvolpato]

Description

Honor caller-provided scale strides in bidirectional BF16-to-MXFP8 casts. The 2048- and 4096-column kernel specializations bake in TE's default padded scale strides. When an output tensor supplies a larger rowwise or columnwise pitch, scale values land in the wrong logical rows even though the FP8 payload is correct.

For example, BF16 input shaped (64, 2048) with unswizzled scale shapes (64, 68) and (2, 2064) reproduces the corruption. Use the K-specialized kernel only when both strides match its constants. Other layouts use the existing runtime-stride kernel.

Related to #3710.

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

• Keep K-specialized launches for matching default scale strides and use runtime strides for custom layouts.
• Add GPU regression coverage for both FP8 formats, custom pitches in either direction, and default-layout/runtime-width controls. Compare FP8 data and scale bytes with canonical quantization output.
• Run the regression on Blackwell quantization hardware, including SM120, without requiring MXFP8 GEMM support.

Validation

Built the common library and PyTorch extension for an RTX 5070 Ti (SM120), using CUDA 13.2 and PyTorch 2.12.1 with NCCL EP and CuTeDSL disabled. Verified the imported libraries came from the tested checkout. The regression fails before the fix and passes afterward. Direct CUDA launcher checks, Compute Sanitizer memcheck, scoped pre-commit, C++ L0 lint, and the license check passed.

Rerun the focused check with:

python -m pytest tests/pytorch/test_quantized_tensor.py -k pitched_scales

The full L0 runtime and JAX suites were not run.

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