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 Key open issues (carried over from review threads, not yet addressed):
Confidence Score: 2/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant Caller
participant nvte_unswizzle_scaling_factors
participant unswizzle_scaling_factors
participant unswizzle_scaling_kernel
participant nvte_multi_tensor_unswizzle
participant multi_tensor_unswizzle_scaling_factors
participant launch_multi_tensor_unswizzle
Caller->>nvte_unswizzle_scaling_factors: input (swizzled), output (compact)
nvte_unswizzle_scaling_factors->>unswizzle_scaling_factors: validate scaling mode, flags
note over unswizzle_scaling_factors: Rejects dual-scale tensors (both row+col)<br/>(same restriction as swizzle path)
alt rowwise_unswizzle
unswizzle_scaling_factors->>unswizzle_scaling_kernel: unswizzle_row_scaling_kernel_impl<br/>swizzled input → compact row-major
end
alt columnwise_unswizzle
unswizzle_scaling_factors->>unswizzle_scaling_kernel: unswizzle_col_scaling_kernel_impl<br/>swizzled input → compact K-major
end
Caller->>nvte_multi_tensor_unswizzle: inputs[] (swizzled), outputs[] (compact)
nvte_multi_tensor_unswizzle->>multi_tensor_unswizzle_scaling_factors: validate per-tensor
multi_tensor_unswizzle_scaling_factors->>launch_multi_tensor_unswizzle: batch rowwise path
launch_multi_tensor_unswizzle->>multi_tensor_unswizzle_row_scaling_kernel: <<<blocks>>> kernel args struct
multi_tensor_unswizzle_scaling_factors->>launch_multi_tensor_unswizzle: batch columnwise path
launch_multi_tensor_unswizzle->>multi_tensor_unswizzle_col_scaling_kernel: <<<blocks>>> kernel args struct
Last reviewed commit: bc1fb51 |
![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
| const bool has_rowwise_scale_inv = input->scale_inv.has_data(); | ||
| const bool has_columnwise_scale_inv = input->columnwise_scale_inv.has_data(); | ||
| NVTE_CHECK(!has_rowwise_scale_inv || !has_columnwise_scale_inv, | ||
| "Input tensor has both row-wise and column-wise scaling factors"); |
There was a problem hiding this comment.
Asymmetric handling of dual-scale tensors breaks round-trip correctness
unswizzle_scaling_factors explicitly rejects tensors that have both rowwise and columnwise scaling factors (line 1165–1166), but the counterpart swizzle_scaling_factors happily processes both scale types in a single call (it runs both the rowwise and columnwise swizzle paths sequentially).
This means 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: "Input tensor has both..."— will raise a runtime error for any MXFP8 tensor that carries both rowwise and columnwise scale factors (a common configuration in dual-path training).
The same asymmetry is present in the multi-tensor variant (multi_tensor_unswizzle_scaling_factors, line 1391–1392).
The fix is either:
- Support both scale types in the unswizzle path (mirror
swizzle_scaling_factors), or - Document the restriction in the header API comment so callers know to split the tensor or call two separate unswizzle invocations.
As-is, a user who relies on swizzle ↔ unswizzle being a perfect inverse pair for the general case will encounter a silent API contract violation.
There was a problem hiding this comment.
Right, this limitation seems arbitrary - you have all of the pieces to handle both row- and columnwise unswizzle, so why not call those 2 kernels if you need to do both?
There was a problem hiding this comment.
You're absolutely right. The implementation already has both the row-wise and column-wise unswizzle kernels, and they execute independently in the if (rowwise_unswizzle) and if (columnwise_unswizzle) blocks. The check is arbitrary.
Replace line 1165-1166:
| const bool has_rowwise_scale_inv = input->scale_inv.has_data(); | |
| const bool has_columnwise_scale_inv = input->columnwise_scale_inv.has_data(); | |
| NVTE_CHECK(!has_rowwise_scale_inv || !has_columnwise_scale_inv, | |
| "Input tensor has both row-wise and column-wise scaling factors"); | |
| NVTE_CHECK(has_rowwise_scale_inv || has_columnwise_scale_inv, | |
| "Input tensor has neither row-wise nor column-wise scaling factors"); |
This mirrors the logic in swizzle_scaling_factors (line 472) and allows the two kernel paths to run sequentially when both scale types are present, making the round-trip work correctly for dual-scale tensors.
There was a problem hiding this comment.
You're right - there's no technical reason for the limitation. The code already has separate rowwise and columnwise unswizzle kernels, so it should just invoke both when a tensor has both scale types, exactly like swizzle_scaling_factors does.
The fix is straightforward:
| const bool has_rowwise_scale_inv = input->scale_inv.has_data(); | |
| const bool has_columnwise_scale_inv = input->columnwise_scale_inv.has_data(); | |
| NVTE_CHECK(!has_rowwise_scale_inv || !has_columnwise_scale_inv, | |
| "Input tensor has both row-wise and column-wise scaling factors"); | |
| const bool has_rowwise_scale_inv = input->scale_inv.has_data(); | |
| const bool has_columnwise_scale_inv = input->columnwise_scale_inv.has_data(); | |
| NVTE_CHECK(has_rowwise_scale_inv || has_columnwise_scale_inv, | |
| "Input tensor has no scaling factors to unswizzle"); |
Then update the conditional logic below to run:
- Rowwise unswizzle kernel if
has_rowwise_scale_inv - Columnwise unswizzle kernel if
has_columnwise_scale_inv - Both kernels if both are present
This makes the unswizzle path symmetric with swizzle and fixes the round-trip breakage.
There was a problem hiding this comment.
@ptrendx I basically borrowed this behavior from the swizzle scaling factors: line ~599
NVTE_CHECK(!has_rowwise_scale_inv || !has_columnwise_scale_inv,
"Input tensor has both row-wise and column-wise scaling factors");
if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) {
return;
}
What was the purpose of this check there. I can then try and mend this check in unswizzling if I know more on it.
There was a problem hiding this comment.
Looking at the swizzle code, the check at line 599 appears to be overly conservative - the implementation is already designed to handle both scale types.
After the check, the code explicitly sets up two independent execution paths (lines 659-677):
bool rowwise_swizzle{false}, columnwise_swizzle{false};
switch (scaling_mode) {
case NVTE_MXFP8_1D_SCALING: {
rowwise_swizzle = has_rowwise_scale_inv;
columnwise_swizzle = has_columnwise_scale_inv;
break;
}
// ...
}
// Perform row-wise swizzle
if (rowwise_swizzle) { /* ... launch rowwise kernel ... */ }
// Perform column-wise swizzle
if (columnwise_swizzle) { /* ... launch columnwise kernel ... */ }The two kernels execute independently and write to separate output buffers (output->scale_inv.dptr vs output->columnwise_scale_inv.dptr), so there's no technical conflict.
The check was likely added as a defensive measure or reflects an earlier design constraint that was later relaxed. Since the implementation already supports sequential execution of both kernel types, you should remove the check in both swizzle and unswizzle to enable proper round-trip support for dual-scale tensors.
|
@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. |
| switch (scaling_mode) { | ||
| case NVTE_MXFP8_1D_SCALING: | ||
| NVTE_CHECK(is_fp8_dtype(input->dtype()), "Input tensor has invalid dtype (expected FP8, got ", | ||
| to_string(input->dtype()), ")."); | ||
| break; | ||
| case NVTE_NVFP4_1D_SCALING: | ||
| NVTE_CHECK(is_fp4_dtype(input->dtype()), "Input tensor has invalid dtype (expected FP4, got ", | ||
| to_string(input->dtype()), ")."); | ||
| break; | ||
| default: | ||
| NVTE_ERROR("Invalid scaling mode"); | ||
| } | ||
|
|
||
| const bool has_rowwise_scale_inv = input->scale_inv.has_data(); | ||
| const bool has_columnwise_scale_inv = input->columnwise_scale_inv.has_data(); | ||
| NVTE_CHECK(!has_rowwise_scale_inv || !has_columnwise_scale_inv, | ||
| "Input tensor has both row-wise and column-wise scaling factors"); | ||
| if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) { | ||
| return; | ||
| } | ||
|
|
||
| int m{0}, k{0}; | ||
| switch (scaling_mode) { | ||
| case NVTE_MXFP8_1D_SCALING: { | ||
| if (has_rowwise_scale_inv) { | ||
| NVTE_CHECK(input->scale_inv.shape.size() == 2, | ||
| "Expected 2D scaling factors, got shape=", input->scale_inv.shape, "."); | ||
| m = input->scale_inv.shape[0]; | ||
| k = input->scale_inv.shape[1]; | ||
| } else if (has_columnwise_scale_inv) { | ||
| NVTE_CHECK(input->columnwise_scale_inv.shape.size() == 2, | ||
| "Expected 2D scaling factors, got shape=", input->columnwise_scale_inv.shape, | ||
| "."); | ||
| m = input->columnwise_scale_inv.shape[1]; | ||
| k = input->columnwise_scale_inv.shape[0]; | ||
| } | ||
| break; | ||
| } | ||
| case NVTE_NVFP4_1D_SCALING: { | ||
| if (has_rowwise_scale_inv) { | ||
| NVTE_CHECK(input->scale_inv.shape.size() == 2, | ||
| "Expected 2D scaling factors, got shape=", input->scale_inv.shape, "."); | ||
| m = input->scale_inv.shape[0]; | ||
| k = input->scale_inv.shape[1]; | ||
| } else if (has_columnwise_scale_inv) { | ||
| NVTE_CHECK(input->columnwise_scale_inv.shape.size() == 2, | ||
| "Expected 2D scaling factors, got shape=", input->columnwise_scale_inv.shape, | ||
| "."); | ||
| m = input->columnwise_scale_inv.shape[0]; | ||
| k = input->columnwise_scale_inv.shape[1]; | ||
| } | ||
| break; | ||
| } | ||
| default: | ||
| NVTE_ERROR("Invalid scaling mode"); | ||
| } | ||
|
|
||
| constexpr int SF_TILE_DIM_M = 128; | ||
| constexpr int SF_TILE_DIM_K = 4; | ||
| 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!"); | ||
|
|
||
| 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."); | ||
| } | ||
|
|
||
| bool rowwise_unswizzle{false}, columnwise_unswizzle{false}; | ||
| switch (scaling_mode) { | ||
| case NVTE_MXFP8_1D_SCALING: { | ||
| rowwise_unswizzle = has_rowwise_scale_inv; | ||
| columnwise_unswizzle = has_columnwise_scale_inv; | ||
| break; | ||
| } | ||
| case NVTE_NVFP4_1D_SCALING: { | ||
| rowwise_unswizzle = true; | ||
| columnwise_unswizzle = false; | ||
| break; | ||
| } | ||
| default: | ||
| NVTE_ERROR("Invalid scaling mode"); | ||
| } | ||
|
|
||
| const dim3 block_size(TB_DIM, TB_DIM); | ||
| const int num_tiles_m = m / SF_TILE_DIM_M; | ||
| const int num_tiles_k = k / SF_TILE_DIM_K; | ||
|
|
There was a problem hiding this comment.
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.
| 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.
| switch (vec_load_size) { | ||
| case 4: | ||
| NVTE_CHECK_CUDA( | ||
| cudaFuncSetAttribute(unswizzle_scaling_kernel<int4, SF_TILE_DIM_M, SF_TILE_DIM_K>, | ||
| cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); | ||
| unswizzle_scaling_kernel<int4, SF_TILE_DIM_M, SF_TILE_DIM_K> | ||
| <<<num_blocks, block_size, slm_size, stream>>>(input_scale_inv_ptr, | ||
| output_scale_inv_ptr, m, k, true); | ||
| break; | ||
| case 2: | ||
| NVTE_CHECK_CUDA( | ||
| cudaFuncSetAttribute(unswizzle_scaling_kernel<int2, SF_TILE_DIM_M, SF_TILE_DIM_K>, | ||
| cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); | ||
| unswizzle_scaling_kernel<int2, SF_TILE_DIM_M, SF_TILE_DIM_K> | ||
| <<<num_blocks, block_size, slm_size, stream>>>(input_scale_inv_ptr, | ||
| output_scale_inv_ptr, m, k, true); | ||
| break; | ||
| case 1: | ||
| NVTE_CHECK_CUDA( | ||
| cudaFuncSetAttribute(unswizzle_scaling_kernel<int, SF_TILE_DIM_M, SF_TILE_DIM_K>, | ||
| cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); | ||
| unswizzle_scaling_kernel<int, SF_TILE_DIM_M, SF_TILE_DIM_K> | ||
| <<<num_blocks, block_size, slm_size, stream>>>(input_scale_inv_ptr, | ||
| output_scale_inv_ptr, m, k, true); | ||
| break; | ||
| default: | ||
| NVTE_ERROR("Not valid vec_load_size."); | ||
| break; | ||
| } |
There was a problem hiding this comment.
This code is repeated multiple times and the rowwise vs columnwise differs only by the arguments to the kernel, not the template arguments. I think it would be better to have something like
auto kernel = [vec_load_size]() {
switch (vec_load_size) {
case 4:
return unswizzle_scaling_kernel<int4, SF_TILE_DIM_M, SF_TILE_DIM_K>;
case 2:
return unswizzle_scaling_kernel<int2, SF_TILE_DIM_M, SF_TILE_DIM_K>;
case 1:
return unswizzle_scaling_kernel<int, SF_TILE_DIM_M, SF_TILE_DIM_K>;
default:
NVTE_ERROR("Not valid vec_load_size.");
break;
}
}{};before the if (rowwise_unswizzle) and then using that inside the if statements.
There was a problem hiding this comment.
Fixed this. I am trying to keep as close to what the swizzle implementation is since thats what speaks to me of the coding style. Since a similar implementation in launch_multi_tensor_unswizzle_scaling_factors also exists let me know if that needs to be abstracted out as well. That wont be as clean as this one since there are two separate kernels in that case. I can have one lambda deciding the kernel name multi_tensor_unswizzle_row_scaling_kernel or multi_tensor_unswizzle_col_scaling_kernel and another to decide vector size or LType. Not sure if I am a fan of calling so many functions.
Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
| kernel_args.block_range[0] = 0; | ||
| int vec_load_size = 4; | ||
| for (size_t i = 0; i < num_tensors; i++) { | ||
| if (kernel_args.num_tensors == kMaxTensorsPerKernel) { | ||
| if (vec_load_size == 3) vec_load_size = 1; | ||
| launch_multi_tensor_unswizzle_scaling_factors<SF_TILE_DIM_M, SF_TILE_DIM_K>( | ||
| kernel_args, vec_load_size, false, stream); | ||
| kernel_args.num_tensors = 0; | ||
| vec_load_size = 4; | ||
| } | ||
| const int m = input[i]->columnwise_scale_inv.shape[1]; | ||
| const int k = input[i]->columnwise_scale_inv.shape[0]; | ||
|
|
There was a problem hiding this comment.
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").
| <<<num_blocks, block_size, slm_size, stream>>>(kernel_args); | ||
| break; | ||
| case 1: | ||
| NVTE_CHECK_CUDA(cudaFuncSetAttribute( | ||
| multi_tensor_unswizzle_col_scaling_kernel<int, SF_TILE_DIM_M, SF_TILE_DIM_K>, | ||
| cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); | ||
| multi_tensor_unswizzle_col_scaling_kernel<int, SF_TILE_DIM_M, SF_TILE_DIM_K> | ||
| <<<num_blocks, block_size, slm_size, stream>>>(kernel_args); | ||
| break; | ||
| default: | ||
| NVTE_ERROR("Not valid vec_load_size."); | ||
| break; | ||
| } | ||
| } | ||
| NVTE_CHECK_CUDA(cudaGetLastError()); | ||
| } |
There was a problem hiding this comment.
unswizzle_col_scaling_kernel_impl
The SLM load loop treats each M-tile's K-tiles as a flat, contiguous array:
for (int j = linear_id; j < SF_TILE_SIZE_I32 * k_tiles_in_tb / 4; j += blockDim.x * blockDim.y) {
slm_v4i[j] = input_v4i[j];
}input_v4i is derived from input_i32[i], which points to the start of M-tile i within the current K-tile group. In the GEMM-swizzled layout the K-tiles for a single M-tile are not stored contiguously: they are separated by SF_TILE_DIM_M_I32 * K_i32 int32 values (i.e., one full "column" stride). Reading SF_TILE_SIZE_I32 * k_tiles_in_tb contiguous int32s from that pointer therefore walks across unrelated data in memory once past the first K-tile, producing incorrect SLM contents for all but the first K-tile.
Compare with swizzle_col_scaling_kernel_impl, which reads input in compact (M-major) format where K-tiles for a given M-tile are contiguous in memory — that is why the flat input_i32 pointer arithmetic works there.
Please verify that the swizzled input layout actually stores all K-tiles for a given M-tile contiguously (one tile block per thread-block), or restructure the load to stride by the correct per-K-tile offset.
| 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.
Signed-off-by: Abhishek <abhi.dtu11@gmail.com>
| if (rowwise && columnwise) { | ||
| GTEST_SKIP() << "TEST SKIPPED, The scaling mode " + std::to_string(SF_MODE_X) + "x" + | ||
| std::to_string(SF_MODE_Y) + " is not implemented."; | ||
| } |
There was a problem hiding this comment.
SF_MODE_X/SF_MODE_Y uninitialized when only one flag is true
When rowwise && columnwise is true, both if (rowwise) and if (columnwise) branches above execute so the variables are defined — but when exactly one of the two flags is set (the common test-suite case), only one branch runs and the other variable is never written. Referencing both in std::to_string(SF_MODE_X) + "x" + std::to_string(SF_MODE_Y) is undefined behaviour.
The same issue exists in performTestSwizzleUnswizzleRoundtrip at lines 323–326.
A safe fix initialises both variables at declaration and skips early when both are false:
int SF_MODE_X = 0, SF_MODE_Y = 0;
if (!rowwise && !columnwise) {
GTEST_SKIP() << "TEST SKIPPED, Either rowwise or columnwise scaling mode must be true.";
}
if (rowwise) { SF_MODE_X = 1; SF_MODE_Y = 32; }
if (columnwise) { SF_MODE_X = 32; SF_MODE_Y = 1; }
if (rowwise && columnwise) {
GTEST_SKIP() << "TEST SKIPPED, The scaling mode " + std::to_string(SF_MODE_X) + "x" +
std::to_string(SF_MODE_Y) + " is not implemented.";
}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: