Correctly pad scaling factor inverses to satisfy cuteDSL requirements

[ksivaman]

Description

Fix grouped MXFP8 swizzle when per-expert rows aren't a multiple of 128 and pad each expert's scales to (128, 4).

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

• Make sure scaling factor inverses are 128x4 padded per tensor.

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

[ksivaman]

/te-ci

[greptile-apps]

Greptile Summary

This PR fixes grouped MXFP8 swizzle when per-expert rows are not a multiple of 128. The core change introduces a "compact" vs "per-tensor-padded" layout distinction: the quantize kernel writes a compact buffer (no padding between experts), while the swizzle output must be padded to (roundup(M,128), roundup(DIVUP(K,32),4)) per tensor for cuDNN consumption. The fix detects the input layout by comparing buffer sizes, sets separate input_stride_bytes/output_stride_bytes for the grouped kernels, and adds IS_PADDED_K/IS_PADDED_M compile-time template flags to prevent out-of-bounds loads past the compact per-tensor extent.

Confidence Score: 5/5

Safe to merge; no P0/P1 issues found; logic is correct and well-tested across edge cases.

All findings are P2 or below. The compact-layout detection, OOB-load prevention, and output buffer allocation are logically correct and consistent between swizzle.cu and swizzle.cpp. The test suite covers aligned, unaligned, and mixed shapes including the originally-failing workload shape.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/common/swizzle/swizzle.cu	Core fix: adds compact-layout detection, separate input/output strides, and IS_PADDED_K/IS_PADDED_M template specializations to avoid OOB loads in the grouped uniform-shape kernels. Logic is correct; earlier padding_m/padding_k compound checks are cleanly separated into orthogonal guards.
transformer_engine/pytorch/csrc/extensions/swizzle.cpp	Python-facing layer: allocates the output buffer in the correct per-tensor padded shape (roundup(M,128), roundup(DIVUP(K,32),4)) instead of using the compact input shape, so cuDNN sees the right strides between experts. The compute_padded_grouped_scale_shape lambda correctly mirrors the swizzle.cu padding formulas for both rowwise and colwise directions.
tests/cpp/operator/test_swizzle.cu	Adds SwizzleGroupedCompactInputTestSuite covering aligned/unaligned M, unaligned K, and combinations; includes the originally-failing shape (2, 2880, 2880). The gather_compact_grouped_scale helper faithfully replicates the quantize kernel's layout, including the trailing group-level alignment.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["maybe_swizzle_grouped_tensor (swizzle.cpp)"]
    A -->|allocate output| B["compute_padded_grouped_scale_shape\nnum_tensors × roundup(M,128) × roundup(⌈K/32⌉,4)"]
    A --> C["nvte_swizzle_grouped_scaling_factors (swizzle.cu)"]

    C --> D{Detect input layout}
    D -->|numel == num_tensors × padded_scale_elems| E["input_is_compact = false\ninput_stride = padded_m × padded_k"]
    D -->|numel == compact_total_scale_elems| F["input_is_compact = true\ninput_stride = m × padded_k (rowwise)\nor ⌈M/32⌉ × padded_m (colwise)"]
    D -->|mismatch| G[NVTE_ERROR]

    E --> H[dispatch_swizzle_*_kernel_impl]
    F --> H

    H -->|IS_PADDED_M=true, row ≥ original_M| I[Zero register, skip __ldg]
    H -->|IS_PADDED_K=true, k_coord ≥ original_K| J[Zero register, skip __ldg]
    H -->|in-bounds| K["__ldg + per-byte boundary zeroing"]

    I --> L[Output: per-tensor padded layout\noutput_stride = padded_m × padded_k]
    J --> L
    K --> L

Loading

_{Reviews (3): Last reviewed commit: "Add test for swizzle + padding fusion" | Re-trigger Greptile}

[greptile-apps]

Silent truncation when logical_shape_nvte.data[0] is not divisible by num_tensors

per_tensor_first_dim is computed with plain integer division. If logical_shape_nvte.data[0] is not an exact multiple of num_tensors (e.g. due to a caller bug or unexpected grouped layout), the result is silently truncated, causing padded_m to be underestimated and the output buffer to be too small. A divisibility assertion would catch this much earlier with a clear error message.

Suggested change

	const size_t per_tensor_first_dim = logical_shape_nvte.data[0] / num_tensors;
	const size_t per_tensor_first_dim = logical_shape_nvte.data[0] / num_tensors;
	NVTE_CHECK(logical_shape_nvte.data[0] % num_tensors == 0,
	"Grouped GEMM swizzle expects logical_shape first dim to be divisible by num_tensors.");

[greptile-apps]

Implicit contract on compact-buffer alignment is not validated

The compact_total_scale_elems formula assumes the upstream quantize kernel allocates the compact scale buffer with its total first dim rounded up to 128 (rowwise) or 4 (colwise). If a caller passes a "plain compact" buffer of size exactly num_tensors * m * padded_k (without trailing alignment slack), neither branch matches and NVTE_ERROR fires with a size-mismatch message that may be hard to diagnose.

Consider also accepting num_tensors * compact_scale_elems as a valid compact size, or documenting this alignment requirement in the error message.

[ptrendx]

@ksivaman Could you add a test exercising the change?

[ksivaman]

/te-ci