Accept compact MXFP8 scales in GEMM preparation

[wujingyue]

general_gemm() rejects compact MXFP8 scales even when with_gemm_swizzled_scales=False: a 64×64 tensor requires rowwise scales of (128, 4) instead of (64, 2), or columnwise scales of (4, 128) instead of (2, 64).

Accept exact compact MXFP8 scale shapes in the single-tensor swizzle API and read them directly into a padded, swizzled output. A compact-input CUDA kernel uses independent input strides and fills padding with zeros, without an intermediate padding allocation or copy. The PyTorch helper allocates the required output shape. Already-padded inputs retain the existing kernel; GEMM preparation preserves the caller's buffers and unswizzled flag.

Fixes #3518 for the PyTorch single-GEMM path. Grouped swizzle is unchanged.

Validation:

• Built matching TE common and PyTorch libraries from this checkout on NVIDIA GB200.
• All 30 focused tests passed: 18 GEMM cases covering both operands, TN/NN/NT, square/rectangular shapes, compact/padded storage, and input preservation; 12 byte-layout cases covering both orientations and padding contents.
• Compute Sanitizer memcheck: all 12 byte-layout cases passed with zero errors.
• Captured GPU microbenchmarks measured 1.1–4.3 µs for direct swizzling versus 2.9–11.1 µs for padding followed by the existing swizzle, across shapes 64×64, 96×160, 4096×4128, and 16384×4128 in both orientations (about 2.5–4.2× faster). These are preparation timings, not end-to-end training speedups.
• cpplint, isort, Black, clang-format, and whitespace checks passed. Full TE L0 suites have not been run.

The GEMM tests construct compact inputs from independently quantized reference planes to isolate swizzle/GEMM behavior. Compact in-place quantization showed separate rectangular numerical mismatches during investigation; this patch does not address those.

[vthumbe1503]

@wujingyue Thanks for the PR!

I wanted to understand what is the use-case where we have unpadded scales in the mxfp8 tensor?

TE quantizer always produces padded scales in the MXFP8 Tensor. And so we should never face this issue where the scales are not padded to (128,4) or (4,128). If external framework is constructing/manipulating the MXFP8Tensor, then I would say, it is the responsibility of that framework and not TE, to pad the scales tensor.

If the external framework cannot sacrifice the performance of additional padding. We can potentially fuse the padding within the swizzling kernel in a way swizzling kernel accepts unpadded scales and fills zero manually in those unpadded locations. We have done something similar for the grouped swizzling kernel.

[wujingyue]

I wanted to understand what is the use-case where we have unpadded scales in the mxfp8 tensor?

Sure. Have you seen the issue this PR links to? I've been adding context there because this PR is just for demo at this moment. Shall we follow up over there?

TE quantizer always produces padded scales in the MXFP8 Tensor.

Not really -- according to #3518, quantizer itself works great with unpadded scales. But I could certainly be missing something!

cc @vthumbe1503

[vthumbe1503]

In the issue you mentioned, custom MXFP8Tensor is being created for weight like this.

weight = MXFP8Tensor(
    shape=(64, 64),
    dtype=torch.bfloat16,
    device="cuda",
    rowwise_data=empty((64, 64)),
    columnwise_data=empty((64, 64)),
    rowwise_scale_inv=empty((128, 4) if padded else (64, 2)),
    columnwise_scale_inv=empty((4, 128) if padded else (2, 64)),
    fp8_dtype=tex.DType.kFloat8E4M3,
    quantizer=quantizer,
    with_gemm_swizzled_scales=False,
)

I see that this is used in Megatron FSDP v2. So I see that you are dealing with weight sharding in FSDP which can be tricky since sharded tensors might not be padded?

quantizer itself works great with unpadded scales

I understand quantization works with unpadded scales. So the problem seems to be that GEMMS need swizzled scales padded.

So 2 solutions can be

1. Pad the compact scales in megatron itself before constructing the MXFP8Tensor shard for the weights. We do something similar in FSD2P.

   TransformerEngine/transformer_engine/pytorch/tensor/mxfp8_tensor.py

   Line 718 in 796346c

   if pad_dim0 > 0:

2. Or instead of padding and creating a copy for the compact scales. How about we handle padding directly in the swizzling kernel, to avoid padding performance penalty? We do something similar in grouped swizzling kernel

   TransformerEngine/transformer_engine/common/swizzle/swizzle.cu

   Lines 376 to 377 in 796346c

   	__device__ __forceinline__ void dispatch_swizzle_row_scaling_kernel_impl(
   	const void* input, void* output, const int M, const int K, const int original_M,

   . Essentially swizzling kernel will be aware that the scales are not padded already. So manually fills them with zeros.

With the current solution, we are silently introducing a performance penalty in the framework by doing the padding, which I am hoping to avoid.

[wujingyue]

How about we handle padding directly in the swizzling kernel, to avoid padding performance penalty

I asked my agent to one-shot it. I really don't understand what it's doing. Could you or someone from the TE team take over? @vthumbe1503