Feature/unswizzle

[int-smart]

Description

This PR adds unswizzle support for scaling factors and extends the swizzle module so scaling tensors can be converted from GEMM-swizzled layout back to compact layout, including multi-tensor paths. It also adds round-trip and standalone tests to validate unswizzle correctness.

Fixes # (issue)

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

Please list the changes introduced in this PR:

• Added unswizzle APIs and implementation in transformer_engine/common/swizzle/swizzle.cu and declarations in transformer_engine/common/include/transformer_engine/swizzle.h
• Added multi-tensor unswizzle support with swizzle-like validation assumptions (homogeneous scaling mode/layout, swizzled input and compact output expectations)
• Refactored multi-tensor unswizzle launch/kernels to mirror swizzle structure (split row-wise and column-wise kernels) for easier readability
• Added/extended tests in tests/cpp/operator/test_swizzle.cu, including standalone unswizzle and swizzle→unswizzle round-trip coverage

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

[greptile-apps]

Greptile Summary

This PR introduces unswizzle support for MXFP8/NVFP4 scaling factors, adding single-tensor and multi-tensor nvte_unswizzle_scaling_factors / nvte_multi_tensor_unswizzle_scaling_factors APIs that reverse the GEMM-swizzled layout back to the compact row-major format. New inverse kernels (unswizzle_row_scaling_kernel_impl, unswizzle_col_scaling_kernel_impl) mirror the existing swizzle kernels, and a suite of standalone and round-trip tests validates correctness for both aligned and padded shapes.

Key concerns:

• Missing input-size validation in unswizzle_scaling_factors (swizzle.cu lines ~1159–1200): The function derives m and k from output->scale_inv.shape, then asserts m * k == output->scale_inv.numel() — a check that is always true for a well-formed 2D tensor and thus provides no safety. Critically, input->scale_inv.numel() is never validated against m * k. This is asymmetric with swizzle_scaling_factors (which validates output->numel()) and with multi_tensor_unswizzle_scaling_factors (which derives m/k from input's shape and validates the output). A caller passing a swizzled input that was built from differently-padded dimensions will trigger silent OOB reads. The same gap exists in both the rowwise and columnwise MXFP8 and NVFP4 paths.

• unswizzle_col_scaling_kernel_impl SLM load relies on undocumented layout invariant (line ~461): The contiguous read of SF_TILE_SIZE_I32 * k_tiles_in_tb int32s from input_i32[i] is correct only because the swizzle kernel writes exactly that many contiguous bytes for each (M-tile, K-tile-block) pair. This invariant should be documented or guarded with an assertion; the existing tests do not stress non-power-of-two K-tile counts large enough to expose any edge case in this region.

• Dual-scale asymmetry (single-tensor and multi-tensor unswizzle both reject tensors with both rowwise and columnwise scale factors), which breaks the round-trip for tensors that carry both — this is an ongoing discussion in prior threads and remains unresolved.

• Roundtrip test only exercises aligned shapes: performTestSwizzleUnswizzleRoundtrip uses num_tiles vectors that produce M/K exact multiples of 128, so the padded-tail correctness validated by the standalone unswizzle tests is not covered by the round-trip suite.

Confidence Score: 2/5

• Not safe to merge: missing input-size validation in unswizzle_scaling_factors can lead to silent OOB reads, and the dual-scale asymmetry (confirmed by a senior developer to need fixing) remains unresolved.
• The single-tensor unswizzle_scaling_factors skips any check that input->scale_inv.numel() == m * k (where m/k come from the output tensor), making it possible to silently operate with a mis-sized input buffer. The multi-tensor path handles this correctly, highlighting the inconsistency. Additionally, the dual-scale restriction that prevents a tensor with both rowwise and columnwise scales from being unswizzled was flagged by a senior developer as needing to be removed, but the current PR still retains it. Together these represent correctness and API-completeness issues that should be addressed before merging.
• transformer_engine/common/swizzle/swizzle.cu — specifically the unswizzle_scaling_factors function (input size validation) and unswizzle_col_scaling_kernel_impl (SLM load invariant documentation).

Important Files Changed

Filename	Overview
transformer_engine/common/swizzle/swizzle.cu	Adds unswizzle kernels and host-side dispatch (single- and multi-tensor). Key issues: unswizzle_scaling_factors derives m/k from the output tensor shape but never validates that the input tensor has the same element count — the output-only check is tautological. The unswizzle_col_scaling_kernel_impl SLM load relies on a layout invariant inherited from the swizzle kernel that is not explicitly documented. The dual-scale restriction (`!has_rowwise
tests/cpp/operator/test_swizzle.cu	Adds standalone unswizzle tests and round-trip tests. Previously reported UB (uninitialized SF_MODE_X/SF_MODE_Y) and missing spaces in skip messages appear addressed in the visible code. Roundtrip test compares only unpadded scale bytes (scale_shape[0]*scale_shape[1]), and all roundtrip shapes are exact multiples of 128, so padding-region correctness is not exercised in round-trips. The reference unswizzle implementation appears logically correct as the exact inverse of compute_ref_swizzle.
transformer_engine/common/include/transformer_engine/swizzle.h	Adds nvte_unswizzle_scaling_factors and nvte_multi_tensor_unswizzle_scaling_factors declarations with matching docstrings. No issues found.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[nvte_unswizzle_scaling_factors] --> B[unswizzle_scaling_factors]
    B --> C{scaling_mode?}
    C -->|MXFP8_1D| D[Derive m,k from output->scale_inv.shape]
    C -->|NVFP4_1D| E[Derive m,k from output scale shape]
    D --> F{rowwise or columnwise?}
    E --> F
    F -->|rowwise| G[Validate m%128==0, k%4==0]
    F -->|columnwise| G
    G --> H{vec_load_size selection\n based on num_tiles}
    H --> I[Launch unswizzle_scaling_kernel\nrow or col path]
    I -->|row_scaling=true| J[unswizzle_row_scaling_kernel_impl\nSLM load → regs_unshuffle → store compact]
    I -->|row_scaling=false| K[unswizzle_col_scaling_kernel_impl\nSLM load → regs_unshuffle_with_bit_shifts → store compact]

    L[nvte_multi_tensor_unswizzle_scaling_factors] --> M[multi_tensor_unswizzle_scaling_factors]
    M --> N{rowwise_unswizzle?}
    N -->|yes| O[Batch tensors into MultiSwizzleArgs\nm,k from input shape\nvalidate output.numel==m*k]
    N -->|no| P[Batch columnwise tensors]
    O --> Q[launch_multi_tensor_unswizzle_scaling_factors\nrowwise=true]
    P --> R[launch_multi_tensor_unswizzle_scaling_factors\nrowwise=false]
    Q --> S[multi_tensor_unswizzle_row_scaling_kernel]
    R --> T[multi_tensor_unswizzle_col_scaling_kernel]

    style D fill:#ffcccc,stroke:#cc0000
    style E fill:#ffcccc,stroke:#cc0000

Loading

_{Last reviewed commit: "Refactor unswizzling..."}

[ptrendx]

@int-smart Please address the comments from Greptile and ideally also add the test case with the input not already padded to 128,128.

[int-smart]

@ptrendx Will look into these

[int-smart]

@ptrendx From what I am understanding then, there is no relevance of padding to the unswizzle kernel. Since the padding is already done during the swizzling operation I can just mirror it back to compact layout with the zero pads correctly in the compact layout and that should do. Is that assumption correct. Initially I was thinking of removing the padding from the scale_inv itself since this would be used for checkpointing

[ptrendx]

@int-smart I'm not sure I follow, I think that what you are saying is probably correct, but let me try to clarify just in case:

• the scaling factors, irrespective of the compact or gemm-ready layout, are zero-padded to the multiple of [128,4] (or the transpose in case of compact and columnwise).
• So for the unswizzle, you should just use the same size of the output unswizzled tensor as the original swizzled one. You don't even need to zero it before unswizzling, since the swizzled tensor already has 0s in the right places so unswizzling it will put 0s in the pad positions.

[int-smart]

@ptrendx Makes sense. I added that in the last commit.

[ptrendx]

The code is pretty convoluted here and it doesn't have to be. There are some pieces there that you could do at the beginning without looking at the scaling factor (like checking whether the input has scale_inv/columnwise_scale_inv and checking if the output has them too). For the rest I would say that avoiding code duplication here is not worth breaking of the flow of NVFP4/MXFP8 specific logic, so I would probably just have a larger switch with 2 completely separate code paths rather than multiple switch statements.

[ptrendx]

I would say that the logic here is a little backwards, even though I understand how here it is not obvious. Ultimately it is the output that tells you what to do in the function - think about the quantize function where the input does not know anything about the format to which it is quantized and it is the output that controls scaling mode and whether we need rowwise or columnwise quantization. Therefore here I would also treat the output as a "source of truth" on what we need to do and then check that the input tensor provides the right data (as opposed to this code which looks to input to know what to do and then checks the output).

[int-smart]

Chaned this for single tensor. Let me know if that makes sense. Can you tell me how this would be called so that I can check the input and output and how they are allocated. Currently I am assuming from your comment above that the output would have all the necessary information to decide between rowwise, columnwise, scaling_mode and data pointers along with dimensions such as m and k. If this is fine then I can make these changes to multi tensor version as well.

[ptrendx]

Yes, the changes look good. Please update the multitensor version accordingly.

[greptile-apps]

original_m_list/original_k_list set but unused by unswizzle kernels

Inside multi_tensor_unswizzle_scaling_factors, both the rowwise path (lines ~955–956 and ~960–961) and the columnwise path (lines ~1006–1007) populate kernel_args.original_m_list[pos] and kernel_args.original_k_list[pos]. However, neither multi_tensor_unswizzle_row_scaling_kernel nor multi_tensor_unswizzle_col_scaling_kernel reads these fields — they only consume m_list and k_list. The swizzle kernels need the original (unpadded) dimensions to zero-fill padding, but the unswizzle kernels always operate on already-padded swizzled input and produce padded compact output, so no masking is required.

Setting unused struct fields is harmless today but adds noise and could mislead a reader into thinking the unswizzle kernels honour padding boundaries the same way the swizzle kernels do. Consider either removing these assignments or adding a comment explaining why they are intentionally populated (e.g., "kept for future per-element padding masking").

[greptile-apps]

Roundtrip test only covers aligned matrix dimensions

performTestSwizzleUnswizzleRoundtrip is instantiated exclusively with the existing num_tiles vector, which always produces M = num_tiles_M * MAT_TILE_DIM_M — values that are exact multiples of 128 (the scale-M alignment). The standalone performTestUnswizzle1D intentionally adds padded shapes (e.g., M=160, K=96) via unswizzle_data_shapes, but no equivalent padded cases exist for the roundtrip.

If the output-size validation or padding-mask logic ever diverges between the swizzle and unswizzle paths for non-aligned M/K, the roundtrip test would pass while standalone tests fail (or vice-versa). Consider adding a few padded shapes (e.g., {4, 3} tile-count pairs or raw {160, 96} shapes) to num_tiles or creating a separate data-shape vector for the roundtrip suite.

[ptrendx]

@int-smart Is there a reason for that difference between the tests?

[greptile-apps]

unswizzle_col_scaling_kernel_impl: SLM load stride mismatch for multi-block K dimension

The SLM load for each M-tile reads SF_TILE_SIZE_I32 * k_tiles_in_tb contiguous int32s from input_i32[i]:

const int4* input_v4i = reinterpret_cast<const int4*>(input_i32[i]);
for (int j = linear_id; j < SF_TILE_SIZE_I32 * k_tiles_in_tb / 4; j += ...)
    slm_v4i[j] = input_v4i[j];

input_i32[i] is set to base + bid_x * TB_DIM * SF_TILE_SIZE_I32 + mt * SF_TILE_DIM_M_I32 * K_i32, where the stride between adjacent M-tiles in the swizzled layout is SF_TILE_DIM_M_I32 * K_i32.

For a full K-tile block (k_tiles_in_tb == TB_DIM == 32), the write size is 32 * SF_TILE_SIZE_I32 = 32 * 128 = 4096 int32s. The M-tile stride is 32 * K_i32. When K_i32 > 128 (e.g., K = 132 K-tiles), the M-tile stride 32 * 132 = 4224 > 4096, so there is a gap that the read does not cross—it is safe.

However, for the last K-tile block (when bid_x == grid_dim_x - 1) and k_tiles_in_tb < TB_DIM, the following K-tile block for the same M-tile starts at offset (bid_x+1) * TB_DIM * 128 + mt * 32 * K_i32, which is beyond the current read range. This appears correct in isolation, but the stride chosen by the swizzle for that region may leave uninitialised bytes between consecutive partial K-tile writes for the same M-tile.

Consider tracing through with K = 8 K-tiles (K_i32 = 8), M = 128:

• num_tiles_k = 8 / SF_TILE_DIM_K_I32 = 8 / 4 = 2; grid_dim_x = DIVUP(2, TB_DIM) = 1
• So bid_x=0 is the last block; k_tiles_in_tb = (2-1) % 32 + 1 = 2
• Read: 2 * 128 = 256 int32s from 0 + 0 * 32 * 8 = 0
• Swizzle stored: 2 * 128 = 256 int32s at offset 0

Layouts agree here. But consider K_i32 = 132, grid_dim_x = 2, bid_x = 1 (last), mt = 1:

• input_i32[1] = 1 * 32 * 128 + 1 * 32 * 132 = 4096 + 4224 = 8320
• k_tiles_in_tb = (33-1) % 32 + 1 = 1
• Read: 1 * 128 = 128 int32s from 8320

Swizzle for (bid_x=1, mt=1) wrote 128 int32s starting at 1 * 32 * 128 + 1 * 32 * 132 = 8320. This matches.

After more careful analysis, the reads do appear correct for the cases tested by the test suite (all aligned shapes). However, the correctness relies on k_tiles_in_tb being computed identically in the swizzle and unswizzle kernels. Please add an assertion or comment clarifying the layout invariant assumed by this contiguous read, and add a test covering non-power-of-two K-tile counts (e.g., K = 132 * 4 = 528 with M = 256) to catch any latent mismatch.

[greptile-apps]

NVTE_CHECK validation is derived from output, not input — input dimensions are never validated

This same pattern applies to the NVFP4 rowwise and columnwise paths: m/k are extracted from output->...shape and the NVTE_CHECK only tests that m * k == output->...numel() (trivially true), while input->...numel() is never checked against m/k. Adding the symmetrical input-side validation here and in the NVFP4 columnwise block (lines ~1254–1270) will close the gap.

// Example for NVFP4 rowwise path:
NVTE_CHECK(input->scale_inv.numel() == static_cast<size_t>(m) * k,
           "Input NVFP4 row-wise scaling factor size mismatch: got ",
           input->scale_inv.numel(), ", expected ", static_cast<size_t>(m) * k, ".");