Conversation
- Introduced `nvte_unswizzle_scaling_factors` to convert swizzled scaling factors back to row-major format. - Implemented `regs_unshuffle_with_bit_shifts` and `regs_unshuffle` for unshuffling operations in CUDA kernels. - Added `unswizzle_row_scaling_kernel_impl` and `unswizzle_col_scaling_kernel_impl` for handling unswizzling in row and column scaling respectively. These changes enhance the functionality of the swizzle module, enabling better handling of scaling factors in tensor operations. Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
These enhancements tests the changes introduced for unswizzling Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
- Introduced `compute_ref_unswizzle` to handle the conversion of swizzled scaling factors back to their original format. - Added `performTestUnswizzle1D` to validate the unswizzling process with various scaling modes. - Created `UnswizzleTestSuite` for comprehensive testing of unswizzling operations. Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
- Moved the definition of `swizzle_row_scaling_kernel` to a new location for better organization. - Ensured the kernel implementation is now properly defined and accessible for scaling operations in the swizzle module. Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
- Introduced `multi_tensor_unswizzle_scaling_factors` to convert swizzled scaling factors back to their original row-major format. - Implemented CUDA kernels for unswizzling in both row and column scaling, enhancing the swizzle module's functionality. - Updated the launch function to handle multiple tensor unswizzling operations efficiently. These changes improve the handling of scaling factors in tensor operations, ensuring better performance and organization within the swizzle module. Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
for more information, see https://pre-commit.ci
Greptile SummaryThis PR adds Previous review threads have surfaced a number of open issues that require attention before merging:
Several threads were marked resolved by maintainers: the dual-scale asymmetry between swizzle and unswizzle will be addressed in a follow-up PR, and the Confidence Score: 2/5Not safe to merge — multiple open correctness and UB issues in swizzle.cu need to be resolved first. Several unresolved issues from previous review threads remain: output-size validation uses unpadded dimensions causing valid calls to be rejected for non-aligned shapes;
Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["nvte_unswizzle_scaling_factors(input, output, stream)"] --> B["unswizzle_scaling_factors()"]
B --> C{scaling_mode}
C -->|MXFP8_1D| D{has rowwise?}
C -->|NVFP4_1D| E{has rowwise?}
D -->|yes| F["m = output.scale_inv.shape[0]\nk = output.scale_inv.shape[1]\nrowwise = true"]
D -->|no, columnwise| G["m = output.columnwise.shape[1]\nk = output.columnwise.shape[0]\nrowwise = false"]
E -->|yes or no| H["rowwise = true\n(NVFP4 always uses row kernel)"]
F --> I["launch unswizzle_scaling_kernel\n(rowwise=true)"]
G --> J["launch unswizzle_scaling_kernel\n(rowwise=false)"]
H --> I
I --> K["unswizzle_row_scaling_kernel_impl\nLoad tiles to SLM\nregs_unshuffle\nWrite compact row-major"]
J --> L["unswizzle_col_scaling_kernel_impl\nLoad tiles to SLM\nregs_unshuffle_with_bit_shifts\nWrite compact col-major"]
M["nvte_multi_tensor_unswizzle_scaling_factors(inputs, outputs, n, stream)"]
M --> N["multi_tensor_unswizzle_scaling_factors()"]
N --> O{rowwise_unswizzle?}
N --> P{columnwise_unswizzle?}
O -->|yes| Q["launch_multi_tensor_unswizzle\nis_rowwise=true\nmulti_tensor_unswizzle_row_scaling_kernel"]
P -->|yes| R["launch_multi_tensor_unswizzle\nis_rowwise=false\nmulti_tensor_unswizzle_col_scaling_kernel"]
Reviews (13): Last reviewed commit: "Merge branch 'main' into feature/unswizz..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
vthumbe1503
added
the
community-contribution
Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
int-smart
force-pushed
the
feature/unswizzle
branch
from
85ea04b to
17dbb33
Compare
for more information, see https://pre-commit.ci
…ather than casting Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
for more information, see https://pre-commit.ci
|
@int-smart Please address the comments from Greptile and ideally also add the test case with the input not already padded to 128,128. |
|
@ptrendx Will look into these |
|
@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 |
- Updated unswizzling kernel implementations to remove original_M and original_K parameters, simplifying the function signatures. - Enhanced test suite to utilize new unswizzling data shapes, ensuring comprehensive coverage of aligned and padded cases. These changes improve the clarity and efficiency of the unswizzling process in the swizzle module. Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
for more information, see https://pre-commit.ci
|
@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:
|
|
@ptrendx Makes sense. I added that in the last commit. |
| if (has_rowwise_scale_inv) { | ||
| NVTE_CHECK(output->scale_inv.has_data(), | ||
| "Output tensor does not have row-wise scaling factors."); | ||
| } | ||
| if (has_columnwise_scale_inv) { | ||
| NVTE_CHECK(output->columnwise_scale_inv.has_data(), | ||
| "Output tensor does not have column-wise scaling factors."); | ||
| } |
There was a problem hiding this comment.
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).
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Yes, the changes look good. Please update the multitensor version accordingly.
Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
| std::vector<std::pair<size_t, size_t>> unswizzle_data_shapes = { | ||
| // Aligned: scale dims are already multiples of 128 and 4 | ||
| {128, 128}, | ||
| {128, 16896}, // K = 132 * 128, large K | ||
| {16896, 128}, // M = 132 * 128, large M | ||
| // M-padding only: M not a multiple of 128 (scale-M needs padding to 256) | ||
| {160, 128}, | ||
| // scale-K padding only: K/32 = 3, padded to 4 | ||
| {128, 96}, | ||
| // Both M and scale-K need padding | ||
| {160, 96}, | ||
| }; | ||
|
|
||
| std::vector<std::pair<bool, bool>> scaling_mode = { | ||
| {true, false}, |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
@int-smart Is there a reason for that difference between the tests?
There was a problem hiding this comment.
Just wanted to keep one test similar to swizzle tests which has aligned test cases. Moved to using unswizzled_data_shapes for roundtrip as well with the aligned cases as part of the unswizzled_data_shapes.
Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
…streamline processing. Need to check if rowwise and columnwise both can be true. If yes the if else needs to account for that Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
| const void* input = kernel_args.input_list[tensor_id]; | ||
| void* output = kernel_args.output_list[tensor_id]; | ||
| const int M = kernel_args.m_list[tensor_id]; | ||
| const int K = kernel_args.k_list[tensor_id]; | ||
|
|
||
| constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int); | ||
| constexpr int N_TILES_IN_TB = TB_DIM * N_TILE_PER_TD; | ||
|
|
||
| const int num_tiles_k = K / SF_TILE_DIM_K; | ||
| const int num_tiles_m = M / SF_TILE_DIM_M; | ||
| const int flat_offset = bid - kernel_args.block_range[tensor_id]; | ||
| const int grid_dim_x = DIVUP(num_tiles_k, N_TILES_IN_TB); | ||
| const int grid_dim_y = num_tiles_m; | ||
| const int bid_x = flat_offset / grid_dim_y; | ||
| const int bid_y = flat_offset % grid_dim_y; | ||
|
|
There was a problem hiding this comment.
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 = 8320k_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.
…put tensors for scaling mode and data validation. Updated checks for input and output tensor shapes to ensure proper handling of row-wise and column-wise scaling factors. Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
| #pragma unroll | ||
| for (int i = 0; i < kVectorSize; i++) | ||
| tmp[i % N_SF_PER_TD_PER_TILE * N_TILE_PER_TD + i / N_SF_PER_TD_PER_TILE] = ptr[i]; |
There was a problem hiding this comment.
In C++, % and * share the same precedence and are left-to-right associative, so i % N_SF_PER_TD_PER_TILE * N_TILE_PER_TD is evaluated as (i % N_SF_PER_TD_PER_TILE) * N_TILE_PER_TD. This happens to be the correct inverse formula, but the lack of explicit parentheses makes the intent unclear compared to regs_shuffle's i / N_TILE_PER_TD + i % N_TILE_PER_TD * N_SF_PER_TD_PER_TILE (which has the same ambiguity). Adding explicit parentheses improves readability:
| #pragma unroll | |
| for (int i = 0; i < kVectorSize; i++) | |
| tmp[i % N_SF_PER_TD_PER_TILE * N_TILE_PER_TD + i / N_SF_PER_TD_PER_TILE] = ptr[i]; | |
| tmp[(i % N_SF_PER_TD_PER_TILE) * N_TILE_PER_TD + i / N_SF_PER_TD_PER_TILE] = ptr[i]; |
| NVTE_CHECK(m % SF_TILE_DIM_M == 0, "Input should be padded in M/N dimension!"); | ||
| NVTE_CHECK(k % SF_TILE_DIM_K == 0, "Input should be padded in K dimension!"); | ||
| NVTE_CHECK(k > 0, "Input scale inverse should be 2D!"); | ||
| NVTE_CHECK(m * k == std::accumulate(output[i]->columnwise_scale_inv.shape.begin(), | ||
| output[i]->columnwise_scale_inv.shape.end(), 1, | ||
| std::multiplies<int>()), | ||
| "Input.columnwise_scale_inv size is not equal to " | ||
| "Output.columnwise_scale_inv size!"); |
There was a problem hiding this comment.
m and k are derived from output[i]->columnwise_scale_inv.shape (shape[1] and shape[0] respectively), so m * k equals output[i]->columnwise_scale_inv.numel() by construction. The std::accumulate on the output shape will always equal m * k, making this check a no-op that never catches a real mismatch. Furthermore, the error message says "Input.columnwise_scale_inv size is not equal to Output.columnwise_scale_inv size!" but the check never reads input[i]->columnwise_scale_inv.numel().
The missing check is against the input (swizzled) size — mirror the same validation added in the rowwise path (line 1380):
NVTE_CHECK(input[i]->columnwise_scale_inv.has_data(), "Input tensor ", i,
" does not have column-wise scaling factors.");
NVTE_CHECK(static_cast<size_t>(m) * k == input[i]->columnwise_scale_inv.numel(),
"Expected input tensor ", i, " to have ", static_cast<size_t>(m) * k,
" column-wise scaling factors, but got shape=",
input[i]->columnwise_scale_inv.shape, ".");… use input numel Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
| m * k, " row-wise scaling factors, but got shape=", output[i]->scale_inv.shape, | ||
| "."); |
There was a problem hiding this comment.
The NVTE_CHECK on line 1364 tests whether input[i]->scale_inv.numel() equals m * k (where m and k come from output[i]->scale_inv.shape). When the check fails the error message should tell the caller what the input actually contained, but it prints output[i]->scale_inv.shape instead.
Because m * k was already derived from the output shape, the message effectively says "got shape [Y, Z]" where Y*Z == m*k — the very value that was expected — giving the developer no information about what the mismatching input looked like.
| m * k, " row-wise scaling factors, but got shape=", output[i]->scale_inv.shape, | |
| "."); | |
| NVTE_CHECK(m * k == input[i]->scale_inv.numel(), "Expected input tensor ", i, " to have ", | |
| m * k, " row-wise scaling factors, but got shape=", input[i]->scale_inv.shape, | |
| "."); |
| NVTE_CHECK(m * k == std::accumulate(input[i]->columnwise_scale_inv.shape.begin(), | ||
| input[i]->columnwise_scale_inv.shape.end(), 1, | ||
| std::multiplies<int>()), | ||
| "Input.columnwise_scale_inv size is not equal to " | ||
| "Output.columnwise_scale_inv size!"); | ||
|
|
||
| int num_tiles_k = k / SF_TILE_DIM_K; | ||
| int vec_load_size_i = (num_tiles_k - 1) % 4 + 1; | ||
| vec_load_size = std::min(vec_load_size, vec_load_size_i); | ||
|
|
||
| const int pos = kernel_args.num_tensors; | ||
| kernel_args.input_list[pos] = const_cast<void*>(input[i]->columnwise_scale_inv.dptr); |
There was a problem hiding this comment.
has_data() validation for input columnwise scaling factors
The columnwise unswizzle path accesses input[i]->columnwise_scale_inv.shape (line 1435) and .dptr (line 1446) without first checking that the data actually exists. The validation loop at lines 1330–1332 only checks output[i]->columnwise_scale_inv.has_data() — not the input.
If input[i]->columnwise_scale_inv.has_data() is false, this block will silently compute with a null/zero numel() on line 1435 (causing the size check to fail spuriously) and then pass a null pointer to the kernel on line 1446, producing undefined behaviour.
By contrast, the analogous rowwise path at lines 1358–1359 explicitly guards against this with:
NVTE_CHECK(input[i]->scale_inv.has_data(), "Input tensor ", i,
" does not have row-wise scaling factors.");The same guard should be added before line 1435:
NVTE_CHECK(input[i]->columnwise_scale_inv.has_data(), "Input tensor ", i,
" does not have column-wise scaling factors.");| // Example for NVFP4 rowwise path: | ||
| NVTE_CHECK(static_cast<size_t>(m) * k == input->scale_inv.numel(), | ||
| "Expected input tensor to have ", static_cast<size_t>(m) * k, | ||
| " row-wise scaling factors, but got shape=", input->scale_inv.shape, "."); |
There was a problem hiding this comment.
The comment // Example for NVFP4 rowwise path: (with trailing whitespace on lines 1225–1226) appears to be an artefact of a reviewer's suggested code snippet that was copy-pasted directly into the source. It is not a meaningful code comment and will confuse future readers about why this block is labelled "example".
| // Example for NVFP4 rowwise path: | |
| NVTE_CHECK(static_cast<size_t>(m) * k == input->scale_inv.numel(), | |
| "Expected input tensor to have ", static_cast<size_t>(m) * k, | |
| " row-wise scaling factors, but got shape=", input->scale_inv.shape, "."); | |
| NVTE_CHECK(static_cast<size_t>(m) * k == input->scale_inv.numel(), | |
| "Expected input tensor to have ", static_cast<size_t>(m) * k, | |
| " row-wise scaling factors, but got shape=", input->scale_inv.shape, "."); |
| NVTE_CHECK(!has_rowwise_scale_inv || !has_columnwise_scale_inv, | ||
| "Output tensor has both row-wise and column-wise scaling factors"); |
There was a problem hiding this comment.
unswizzle_scaling_factors hard-rejects any tensor that has both rowwise and columnwise scaling factors (line 1155–1156). The corresponding swizzle_scaling_factors (line 599) has the same check, so calling the public round-trip pair:
nvte_swizzle_scaling_factors(input, swizzled, stream); // succeeds — handles both scales
nvte_unswizzle_scaling_factors(swizzled, output, stream); // FAILS — "Output tensor has both…"will throw at runtime for any dual-path MXFP8 tensor (common in training). The implementation already has independent rowwise and columnwise kernel paths (if (has_rowwise_scale_inv) / else if (has_columnwise_scale_inv) at lines 1183 and 1197), so the restriction is artificial. Lifting it requires changing else if (has_columnwise_scale_inv) to if (has_columnwise_scale_inv) and executing both paths sequentially when both flags are set — exactly mirroring swizzle_scaling_factors's design.
|
@ptrendx Let me know if this makes sense. All the major issues have been resolved |
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
|
/te-ci |
for more information, see https://pre-commit.ci
|
@int-smart Yes, it looks good. I resolved the merge conflict and launched the CI. |
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
|
/te-ci |
|
@int-smart Thank you for your contribution, merged :-)! |
3 participants
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
Changes
Please list the changes introduced in this PR:
transformer_engine/common/swizzle/swizzle.cuand declarations intransformer_engine/common/include/transformer_engine/swizzle.htests/cpp/operator/test_swizzle.cu, including standalone unswizzle and swizzle→unswizzle round-trip coverageChecklist: