From 4442942d7695fda7a7951b04f236b6f5d0befa7c Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 27 Oct 2026 16:55:41 -0700 Subject: [PATCH] Optimize scale swizzle kernels Use optimized rowwise full-width and columnwise chunk kernels for MXFP8 and NVFP4 scales, including multi-tensor swizzling. Keep the 16-byte scale-buffer alignment contract and cover its validation in C++ tests. Choose the rowwise full-width kernel after the fallback grid exceeds one SM wave, based on paired GB200 and VR200 measurements. Signed-off-by: Varun Thumbe --- docs/envvars.rst | 6 - tests/cpp/operator/test_multi_swizzle.cu | 224 +++- tests/cpp/operator/test_swizzle.cu | 196 +++- .../include/transformer_engine/swizzle.h | 3 + transformer_engine/common/swizzle/swizzle.cu | 1044 +++++------------ 5 files changed, 676 insertions(+), 797 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 09e2b462780..5ff430eabea 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -351,12 +351,6 @@ Kernel Configuration :Default: ``0`` :Description: Disable NVRTC (CUDA Runtime Compilation) support. When set to ``1``, runtime kernel compilation is disabled. Existing transpose operations select their static fallback automatically. NVFP4 4over6 quantization, fused softmax, and normalization paths require their corresponding ``NVTE_BUILD_LEGACY_STATIC_NVFP4``, ``NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX``, or ``NVTE_BUILD_LEGACY_STATIC_NORM`` CMake option to have been enabled when the library was built; otherwise no static fallback is available. -.. envvar:: NVTE_MXFP8_PER_TILE_SWIZZLE - - :Type: ``int`` (0 or 1) - :Default: ``0`` - :Description: Swizzle MXFP8 scaling factors for GEMM with kernels that give each 128-row tile of scaling factors (row-wise) or each 128-column by 32-row block (column-wise) its own thread block, instead of batching 32 tiles in one block. This lowers the latency of swizzling the scaling factors of tensors with few tiles, such as activations. - .. envvar:: NVTE_USE_CUTLASS_GROUPED_GEMM :Type: ``int`` (0 or 1) diff --git a/tests/cpp/operator/test_multi_swizzle.cu b/tests/cpp/operator/test_multi_swizzle.cu index 4bd1800cea2..c0ac1501391 100644 --- a/tests/cpp/operator/test_multi_swizzle.cu +++ b/tests/cpp/operator/test_multi_swizzle.cu @@ -51,18 +51,25 @@ void compute_ref_swizzle(const uint8_t *h_input, uint8_t *h_output, } } -static void zero_scale_inv_padding(uint8_t *buf, - size_t padded_rows, size_t padded_cols, - size_t orig_rows, size_t orig_cols) { +static void set_scale_inv_padding(uint8_t *buf, + size_t padded_rows, size_t padded_cols, + size_t orig_rows, size_t orig_cols, + uint8_t value) { for (size_t r = 0; r < padded_rows; ++r) { for (size_t c = 0; c < padded_cols; ++c) { if (r >= orig_rows || c >= orig_cols) { - buf[r * padded_cols + c] = 0; + buf[r * padded_cols + c] = value; } } } } +static void zero_scale_inv_padding(uint8_t *buf, + size_t padded_rows, size_t padded_cols, + size_t orig_rows, size_t orig_cols) { + set_scale_inv_padding(buf, padded_rows, padded_cols, orig_rows, orig_cols, 0); +} + // =================================================================== // Multi-tensor swizzle test // =================================================================== @@ -93,12 +100,12 @@ void performTestMultiTensorSwizzle(const int num_tensors, const size_t M, const const NVTEShape rs = input->rowwise_scale_inv_shape(); zero_scale_inv_padding(input->rowwise_cpu_scale_inv_ptr(), rs.data[0], rs.data[1], - M, (K + BLOCK_SIZE - 1) / BLOCK_SIZE); + M, divide_round_up(K, BLOCK_SIZE)); } else { const NVTEShape cs = input->columnwise_scale_inv_shape(); zero_scale_inv_padding(input->columnwise_cpu_scale_inv_ptr(), cs.data[0], cs.data[1], - (M + BLOCK_SIZE - 1) / BLOCK_SIZE, K); + divide_round_up(M, BLOCK_SIZE), K); } input->from_cpu(); @@ -141,6 +148,159 @@ void performTestMultiTensorSwizzle(const int num_tensors, const size_t M, const } } +void performTestMultiTensorSwizzleNVFP4(const int num_tensors, const size_t M, const size_t K, + bool rowwise) { + using namespace test; + constexpr size_t BLOCK_SIZE = 16; + std::vector> inputs; + std::vector> outputs; + std::vector input_handles; + std::vector output_handles; + + for (int i = 0; i < num_tensors; ++i) { + auto input = std::make_unique("input_" + std::to_string(i), + std::vector{M, K}, DType::kFloat4E2M1, rowwise, + !rowwise, NVTE_NVFP4_1D_SCALING); + auto output = std::make_unique("output_" + std::to_string(i), + std::vector{M, K}, DType::kFloat4E2M1, rowwise, + !rowwise, NVTE_NVFP4_1D_SCALING); + output->set_with_gemm_swizzled_scales(true); + if (rowwise) { + input->fill_uniform_rowwise_scale_inv(); + } else { + input->fill_uniform_columnwise_scale_inv(); + } + const NVTEShape scale_shape = + rowwise ? input->rowwise_scale_inv_shape() : input->columnwise_scale_inv_shape(); + uint8_t *input_scales = rowwise ? input->rowwise_cpu_scale_inv_ptr() + : input->columnwise_cpu_scale_inv_ptr(); + const size_t original_rows = rowwise ? M : K; + const size_t original_cols = + rowwise ? divide_round_up(K, BLOCK_SIZE) : divide_round_up(M, BLOCK_SIZE); + set_scale_inv_padding(input_scales, scale_shape.data[0], scale_shape.data[1], original_rows, + original_cols, 0xFF); + input->from_cpu(); + input_handles.push_back(input->data()); + output_handles.push_back(output->data()); + inputs.emplace_back(std::move(input)); + outputs.emplace_back(std::move(output)); + } + + nvte_multi_tensor_swizzle_scaling_factors(input_handles.data(), output_handles.data(), + num_tensors, 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + for (int i = 0; i < num_tensors; ++i) { + outputs[i]->to_cpu(); + const NVTEShape scale_shape = + rowwise ? inputs[i]->rowwise_scale_inv_shape() : inputs[i]->columnwise_scale_inv_shape(); + const uint8_t *input_scales = rowwise ? inputs[i]->rowwise_cpu_scale_inv_ptr() + : inputs[i]->columnwise_cpu_scale_inv_ptr(); + const uint8_t *actual = rowwise ? outputs[i]->rowwise_cpu_scale_inv_ptr() + : outputs[i]->columnwise_cpu_scale_inv_ptr(); + std::vector reference_input( + input_scales, input_scales + scale_shape.data[0] * scale_shape.data[1]); + const size_t original_rows = rowwise ? M : K; + const size_t original_cols = + rowwise ? divide_round_up(K, BLOCK_SIZE) : divide_round_up(M, BLOCK_SIZE); + zero_scale_inv_padding(reference_input.data(), scale_shape.data[0], scale_shape.data[1], + original_rows, original_cols); + std::vector reference(scale_shape.data[0] * scale_shape.data[1]); + compute_ref_swizzle<128, 4, true>(reference_input.data(), reference.data(), scale_shape.data[0], + scale_shape.data[1]); + compareResults("multi_nvfp4_swizzle", actual, reference.data(), reference.size()); + } +} + +TEST(MultiTensorSwizzleTest, MXFP8ColumnwiseBatchIgnoresUnusedRowwiseBuffer) { + using namespace test; + constexpr int num_tensors = 2; + const std::vector shape{256, 128}; + std::vector> inputs; + std::vector> outputs; + std::vector input_handles; + std::vector output_handles; + + for (int i = 0; i < num_tensors; ++i) { + const bool has_rowwise_scales = i == 0; + auto input = std::make_unique("input_" + std::to_string(i), shape, + DType::kFloat8E4M3, has_rowwise_scales, true, + NVTE_MXFP8_1D_SCALING); + auto output = std::make_unique("output_" + std::to_string(i), shape, + DType::kFloat8E4M3, has_rowwise_scales, true, + NVTE_MXFP8_1D_SCALING); + if (has_rowwise_scales) input->fill_uniform_rowwise_scale_inv(); + input->fill_uniform_columnwise_scale_inv(); + output->set_with_gemm_swizzled_scales(true); + input->from_cpu(); + input_handles.push_back(input->data()); + output_handles.push_back(output->data()); + inputs.emplace_back(std::move(input)); + outputs.emplace_back(std::move(output)); + } + + nvte_multi_tensor_swizzle_scaling_factors(input_handles.data(), output_handles.data(), + num_tensors, 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + for (int i = 0; i < num_tensors; ++i) { + outputs[i]->to_cpu(); + const NVTEShape scale_shape = inputs[i]->columnwise_scale_inv_shape(); + std::vector reference(scale_shape.data[0] * scale_shape.data[1]); + compute_ref_swizzle<128, 4, false>( + inputs[i]->columnwise_cpu_scale_inv_ptr(), reference.data(), scale_shape.data[1], + scale_shape.data[0]); + compareResults("multi_mixed_buffers_columnwise_" + std::to_string(i), + outputs[i]->columnwise_cpu_scale_inv_ptr(), reference.data(), + reference.size()); + } +} + +TEST(MultiTensorSwizzleTest, NVFP4FullWidthChunkLaunchSplitting) { + for (const bool rowwise : {false, true}) { + performTestMultiTensorSwizzleNVFP4(4, 160, 96, rowwise); + performTestMultiTensorSwizzleNVFP4(4, 128, 128, rowwise); + performTestMultiTensorSwizzleNVFP4(65, 128, 128, rowwise); + } +} + +TEST(MultiTensorSwizzleTest, RowwiseMixedSharedMemoryFallback) { + using namespace test; + // The first tensor fits the full-width shared-memory limit, but the second does not. + const std::vector> shapes{{128, 7168}, {128, 8064}}; + std::vector> inputs; + std::vector> outputs; + std::vector input_handles; + std::vector output_handles; + for (size_t i = 0; i < shapes.size(); ++i) { + const auto [M, K] = shapes[i]; + auto input = std::make_unique("input_" + std::to_string(i), + std::vector{M, K}, DType::kFloat8E4M3, true, + false, NVTE_MXFP8_1D_SCALING); + auto output = std::make_unique("output_" + std::to_string(i), + std::vector{M, K}, DType::kFloat8E4M3, true, + false, NVTE_MXFP8_1D_SCALING); + output->set_with_gemm_swizzled_scales(true); + input->fill_uniform_rowwise_scale_inv(); + input->from_cpu(); + input_handles.push_back(input->data()); + output_handles.push_back(output->data()); + inputs.emplace_back(std::move(input)); + outputs.emplace_back(std::move(output)); + } + + nvte_multi_tensor_swizzle_scaling_factors( + input_handles.data(), output_handles.data(), static_cast(inputs.size()), 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + for (size_t i = 0; i < inputs.size(); ++i) { + outputs[i]->to_cpu(); + const NVTEShape scale_shape = inputs[i]->rowwise_scale_inv_shape(); + std::vector reference(scale_shape.data[0] * scale_shape.data[1]); + compute_ref_swizzle<128, 4, true>(inputs[i]->rowwise_cpu_scale_inv_ptr(), + reference.data(), scale_shape.data[0], scale_shape.data[1]); + compareResults("multi_rowwise_fallback", outputs[i]->rowwise_cpu_scale_inv_ptr(), + reference.data(), reference.size()); + } +} + // =================================================================== // Multi-tensor unswizzle test (uses single-tensor swizzle to prepare) // =================================================================== @@ -172,12 +332,12 @@ void performTestMultiTensorUnswizzle(const int num_tensors, const size_t M, cons const NVTEShape rs = orig->rowwise_scale_inv_shape(); zero_scale_inv_padding(orig->rowwise_cpu_scale_inv_ptr(), rs.data[0], rs.data[1], - M, (K + BLOCK_SIZE - 1) / BLOCK_SIZE); + M, divide_round_up(K, BLOCK_SIZE)); } else { const NVTEShape cs = orig->columnwise_scale_inv_shape(); zero_scale_inv_padding(orig->columnwise_cpu_scale_inv_ptr(), cs.data[0], cs.data[1], - (M + BLOCK_SIZE - 1) / BLOCK_SIZE, K); + divide_round_up(M, BLOCK_SIZE), K); } orig->from_cpu(); @@ -249,12 +409,12 @@ void performTestMultiTensorRoundtrip(const int num_tensors, const size_t M, cons const NVTEShape rs = orig->rowwise_scale_inv_shape(); zero_scale_inv_padding(orig->rowwise_cpu_scale_inv_ptr(), rs.data[0], rs.data[1], - M, (K + BLOCK_SIZE - 1) / BLOCK_SIZE); + M, divide_round_up(K, BLOCK_SIZE)); } else { const NVTEShape cs = orig->columnwise_scale_inv_shape(); zero_scale_inv_padding(orig->columnwise_cpu_scale_inv_ptr(), cs.data[0], cs.data[1], - (M + BLOCK_SIZE - 1) / BLOCK_SIZE, K); + divide_round_up(M, BLOCK_SIZE), K); } orig->from_cpu(); @@ -300,22 +460,11 @@ void performTestMultiTensorRoundtrip(const int num_tensors, const size_t M, cons // Test suites // =================================================================== -// Sets NVTE_MXFP8_PER_TILE_SWIZZLE while it lives, so that MXFP8 swizzles use the per-tile -// kernels (enable) or the default ones. -class PerTileSwizzle { - public: - explicit PerTileSwizzle(bool enable) { - setenv("NVTE_MXFP8_PER_TILE_SWIZZLE", enable ? "1" : "0", 1); - } - ~PerTileSwizzle() { unsetenv("NVTE_MXFP8_PER_TILE_SWIZZLE"); } -}; - class MultiTensorSwizzleTestSuite - : public ::testing::TestWithParam, bool>> {}; + : public ::testing::TestWithParam> {}; TEST_P(MultiTensorSwizzleTestSuite, TestMultiTensorSwizzle) { - const auto& [num_tensors, M, K, rowwise] = std::get<0>(GetParam()); - const PerTileSwizzle per_tile_swizzle(std::get<1>(GetParam())); + const auto& [num_tensors, M, K, rowwise] = GetParam(); performTestMultiTensorSwizzle(num_tensors, M, K, rowwise); } @@ -343,17 +492,7 @@ TEST_P(MultiTensorRoundtripTestSuite, TestMultiTensorRoundtrip) { namespace { -// Shapes that exercise the narrow_k kernel (rowwise) / narrow_m kernel (colwise): -// Narrow-K fires when ALL tensors have scale num_tiles_k < TB_DIM (32), -// i.e. padded ceil(K/32) < 128. -// Narrow-M fires analogously for colwise when padded K < 4096 -// (since colwise m = K padded to 128, num_tiles_m = m / 128 < 32). -// -// Shapes that bypass narrow and use the regular multi_tensor kernel: -// K >= 4096 makes num_tiles_k >= 32 (rowwise) and num_tiles_m >= 32 (colwise). - std::vector> multi_tensor_test_cases = { - // --- Narrow path cases (K small → narrow_k for row, narrow_m for col) --- // M and K both aligned to 128 {3, 256, 256, true}, {3, 256, 256, false}, @@ -365,7 +504,7 @@ std::vector> multi_tensor_test_cases = { {3, 192, 256, false}, {2, 64, 256, true}, {2, 64, 256, false}, - // Larger narrow K (num_tiles_k = 8, shared mem = 128 KB) + // Larger K {2, 128, 1024, true}, {2, 128, 1024, false}, // K not divisible by 128 @@ -380,13 +519,16 @@ std::vector> multi_tensor_test_cases = { {4, 32, 64, true}, {4, 32, 64, false}, - // --- Non-narrow path cases (K >= 4096 → regular multi_tensor kernel) --- {3, 256, 4096, true}, {3, 256, 4096, false}, {2, 128, 8192, true}, {2, 128, 8192, false}, // Rowwise num_tiles_k = 34 selects vec_load_size = 2 (int2 kernel) {2, 128, 4352, true}, + // Wide row-wise scales fall back for an underfilled grid and use the full-width kernel once + // the aggregate grid is large enough. + {2, 2048, 7168, true}, // 32 aggregate row tiles: regular fallback. + {16, 2048, 7168, true}, // 256 aggregate row tiles: rowwise full-width kernel. // Colwise num_tiles_k = 512 / 32 / 4 = 4 selects vec_load_size = 4 (int4 kernel) {2, 512, 4096, false}, // MXFP8 swizzle: more tensors than one launch takes, several row-wise tiles per CTA, @@ -404,14 +546,12 @@ std::vector> multi_tensor_test_cases = { INSTANTIATE_TEST_SUITE_P( OperatorTest, MultiTensorSwizzleTestSuite, - ::testing::Combine(::testing::ValuesIn(multi_tensor_test_cases), ::testing::Values(false, true)), + ::testing::ValuesIn(multi_tensor_test_cases), [](const testing::TestParamInfo& info) { - const auto& test_case = std::get<0>(info.param); - return "n" + std::to_string(std::get<0>(test_case)) + - "_M" + std::to_string(std::get<1>(test_case)) + - "_K" + std::to_string(std::get<2>(test_case)) + - (std::get<3>(test_case) ? "_row" : "_col") + - (std::get<1>(info.param) ? "_per_tile" : ""); + return "n" + std::to_string(std::get<0>(info.param)) + + "_M" + std::to_string(std::get<1>(info.param)) + + "_K" + std::to_string(std::get<2>(info.param)) + + (std::get<3>(info.param) ? "_row" : "_col"); }); INSTANTIATE_TEST_SUITE_P( diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index a442215ae15..4ebb23fdc51 100644 --- a/tests/cpp/operator/test_swizzle.cu +++ b/tests/cpp/operator/test_swizzle.cu @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -204,18 +205,167 @@ void performTestUnswizzle1D(const size_t M, const size_t K, bool rowwise, bool c // The buffer is stored in leading-dim-major order (row-major for rowwise, // column-major for colwise). `padded_rows x padded_cols` is the full // (padded) shape; `orig_rows` / `orig_cols` are the unpadded extents. -static void zero_scale_inv_padding(uint8_t *buf, - size_t padded_rows, size_t padded_cols, - size_t orig_rows, size_t orig_cols) { +static void set_scale_inv_padding(uint8_t *buf, + size_t padded_rows, size_t padded_cols, + size_t orig_rows, size_t orig_cols, + uint8_t value) { for (size_t r = 0; r < padded_rows; ++r) { for (size_t c = 0; c < padded_cols; ++c) { if (r >= orig_rows || c >= orig_cols) { - buf[r * padded_cols + c] = 0; + buf[r * padded_cols + c] = value; } } } } +static void zero_scale_inv_padding(uint8_t *buf, + size_t padded_rows, size_t padded_cols, + size_t orig_rows, size_t orig_cols) { + set_scale_inv_padding(buf, padded_rows, padded_cols, orig_rows, orig_cols, 0); +} + +void performTestSwizzleNVFP4(const size_t M, const size_t K, bool rowwise) { + using namespace test; + constexpr size_t BLOCK_SIZE = 16; + + const std::vector shape{M, K}; + Tensor input("input", shape, DType::kFloat4E2M1, rowwise, !rowwise, NVTE_NVFP4_1D_SCALING); + Tensor output("output", shape, DType::kFloat4E2M1, rowwise, !rowwise, NVTE_NVFP4_1D_SCALING); + output.set_with_gemm_swizzled_scales(true); + if (rowwise) { + input.fill_uniform_rowwise_scale_inv(); + } else { + input.fill_uniform_columnwise_scale_inv(); + } + + const NVTEShape scale_shape = + rowwise ? input.rowwise_scale_inv_shape() : input.columnwise_scale_inv_shape(); + uint8_t *input_scales = rowwise ? input.rowwise_cpu_scale_inv_ptr() + : input.columnwise_cpu_scale_inv_ptr(); + const size_t original_rows = rowwise ? M : K; + const size_t original_cols = + rowwise ? divide_round_up(K, BLOCK_SIZE) : divide_round_up(M, BLOCK_SIZE); + set_scale_inv_padding(input_scales, scale_shape.data[0], scale_shape.data[1], original_rows, + original_cols, 0xFF); + std::vector reference_input( + input_scales, input_scales + scale_shape.data[0] * scale_shape.data[1]); + zero_scale_inv_padding(reference_input.data(), scale_shape.data[0], scale_shape.data[1], + original_rows, original_cols); + input.from_cpu(); + + std::vector reference(scale_shape.data[0] * scale_shape.data[1]); + compute_ref_swizzle<128, 4, true>(reference_input.data(), reference.data(), scale_shape.data[0], + scale_shape.data[1]); + nvte_swizzle_scaling_factors(input.data(), output.data(), 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + output.to_cpu(); + + const uint8_t *actual = rowwise ? output.rowwise_cpu_scale_inv_ptr() + : output.columnwise_cpu_scale_inv_ptr(); + compareResults("nvfp4_swizzle", actual, reference.data(), reference.size()); +} + +TEST(SwizzleTest, NVFP4FullWidthChunkAndFallback) { + performTestSwizzleNVFP4(128, 128, true); + performTestSwizzleNVFP4(128, 128, false); + performTestSwizzleNVFP4(160, 96, true); // Padding in both scale dimensions. + performTestSwizzleNVFP4(160, 96, false); // Transposed padding in both scale dimensions. + performTestSwizzleNVFP4(128, 3584, true); // 28.5 KiB, 1 tile: rowwise full-width kernel. + performTestSwizzleNVFP4(2048, 3584, true); // 28.5 KiB, 16 tiles: rowwise full-width kernel. + performTestSwizzleNVFP4(2176, 3584, true); // 28.5 KiB, 17 tiles: regular fallback. + performTestSwizzleNVFP4(128, 4032, true); // 31.5 KiB: regular fallback. + performTestSwizzleNVFP4(128, 4096, true); // 32.5 KiB: regular fallback. + performTestSwizzleNVFP4(3584, 128, false); // 28.5 KiB, 1 tile: transposed full-width kernel. + performTestSwizzleNVFP4(3584, 2048, false); // 28.5 KiB, 16 tiles: transposed full-width kernel. + performTestSwizzleNVFP4(3584, 2176, false); // 28.5 KiB, 17 tiles: regular fallback. + performTestSwizzleNVFP4(4032, 128, false); // 31.5 KiB: regular fallback. + performTestSwizzleNVFP4(4096, 128, false); // 32.5 KiB: regular fallback. +} + +constexpr size_t kAlignmentTestM = 128; +constexpr size_t kAlignmentTestK = 128; +constexpr size_t kAlignmentTestBlockSize = 32; +constexpr size_t kAlignmentOffsetBytes = 4; + +void performTestRejectUnalignedMXFP8Scales(bool rowwise, bool offset_input) { + using namespace test; + constexpr size_t kDataSize = kAlignmentTestM * kAlignmentTestK; + constexpr size_t kScaleSize = kDataSize / kAlignmentTestBlockSize; + auto input_data = cuda_alloc(kDataSize); + auto output_data = cuda_alloc(kDataSize); + auto input_storage = cuda_alloc(kScaleSize + kAlignmentOffsetBytes); + auto output_storage = cuda_alloc(kScaleSize + kAlignmentOffsetBytes); + uint8_t* input_scales = input_storage.get() + (offset_input ? kAlignmentOffsetBytes : 0); + uint8_t* output_scales = output_storage.get() + (offset_input ? 0 : kAlignmentOffsetBytes); + TensorWrapper input(NVTE_MXFP8_1D_SCALING); + TensorWrapper output(NVTE_MXFP8_1D_SCALING); + const std::vector data_shape{kAlignmentTestM, kAlignmentTestK}; + const std::vector scale_shape = + rowwise ? std::vector{kAlignmentTestM, kAlignmentTestK / kAlignmentTestBlockSize} + : std::vector{kAlignmentTestM / kAlignmentTestBlockSize, kAlignmentTestK}; + if (rowwise) { + input.set_rowwise_data(input_data.get(), DType::kFloat8E4M3, data_shape); + input.set_rowwise_scale_inv(input_scales, DType::kFloat8E8M0, scale_shape); + output.set_rowwise_data(output_data.get(), DType::kFloat8E4M3, data_shape); + output.set_rowwise_scale_inv(output_scales, DType::kFloat8E8M0, scale_shape); + } else { + input.set_columnwise_data(input_data.get(), DType::kFloat8E4M3, data_shape); + input.set_columnwise_scale_inv(input_scales, DType::kFloat8E8M0, scale_shape); + output.set_columnwise_data(output_data.get(), DType::kFloat8E4M3, data_shape); + output.set_columnwise_scale_inv(output_scales, DType::kFloat8E8M0, scale_shape); + } + output.set_with_gemm_swizzled_scales(true); + EXPECT_THROW(nvte_swizzle_scaling_factors(input.data(), output.data(), 0), std::runtime_error); +} + +TEST(SwizzleTest, RejectUnalignedMXFP8ScaleBuffers) { + for (bool rowwise : {true, false}) { + performTestRejectUnalignedMXFP8Scales(rowwise, true); + performTestRejectUnalignedMXFP8Scales(rowwise, false); + } +} + +TEST(SwizzleTest, RejectUnalignedMultiMXFP8ScaleBuffers) { + using namespace test; + constexpr size_t kNumTensors = 2; + constexpr size_t kDataSize = kAlignmentTestM * kAlignmentTestK; + constexpr size_t kScaleSize = kDataSize / kAlignmentTestBlockSize; + const std::vector data_shape{kAlignmentTestM, kAlignmentTestK}; + const std::vector scale_shape{kAlignmentTestM, + kAlignmentTestK / kAlignmentTestBlockSize}; + for (bool offset_input : {true, false}) { + std::vector> input_data, output_data, input_storage, output_storage; + std::vector input_tensors, output_tensors; + std::vector inputs, outputs; + input_tensors.reserve(kNumTensors); + output_tensors.reserve(kNumTensors); + for (size_t i = 0; i < kNumTensors; i++) { + input_data.push_back(cuda_alloc(kDataSize)); + output_data.push_back(cuda_alloc(kDataSize)); + input_storage.push_back(cuda_alloc(kScaleSize + kAlignmentOffsetBytes)); + output_storage.push_back(cuda_alloc(kScaleSize + kAlignmentOffsetBytes)); + uint8_t* input_scales = + input_storage.back().get() + (i == 0 && offset_input ? kAlignmentOffsetBytes : 0); + uint8_t* output_scales = + output_storage.back().get() + (i == 0 && !offset_input ? kAlignmentOffsetBytes : 0); + input_tensors.emplace_back(NVTE_MXFP8_1D_SCALING); + output_tensors.emplace_back(NVTE_MXFP8_1D_SCALING); + input_tensors.back().set_rowwise_data(input_data.back().get(), DType::kFloat8E4M3, + data_shape); + input_tensors.back().set_rowwise_scale_inv(input_scales, DType::kFloat8E8M0, scale_shape); + output_tensors.back().set_rowwise_data(output_data.back().get(), DType::kFloat8E4M3, + data_shape); + output_tensors.back().set_rowwise_scale_inv(output_scales, DType::kFloat8E8M0, scale_shape); + output_tensors.back().set_with_gemm_swizzled_scales(true); + inputs.push_back(input_tensors.back().data()); + outputs.push_back(output_tensors.back().data()); + } + EXPECT_THROW( + nvte_multi_tensor_swizzle_scaling_factors(inputs.data(), outputs.data(), inputs.size(), 0), + std::runtime_error); + } +} + void performTestGroupedSwizzleMXFP8(const int num_tensors, const size_t M, const size_t K) { using namespace transformer_engine; using namespace test; @@ -308,17 +458,9 @@ void performTestGroupedSwizzleMXFP8(const int num_tensors, const size_t M, const num_tensors * col_numel); } -// Sets NVTE_MXFP8_PER_TILE_SWIZZLE while it lives, so that MXFP8 swizzles use the per-tile -// kernels (enable) or the default ones. -class PerTileSwizzle { - public: - explicit PerTileSwizzle(bool enable) { - setenv("NVTE_MXFP8_PER_TILE_SWIZZLE", enable ? "1" : "0", 1); - } - ~PerTileSwizzle() { unsetenv("NVTE_MXFP8_PER_TILE_SWIZZLE"); } -}; - -class SwizzleTestSuite : public ::testing::TestWithParam, std::pair, bool, bool>> {}; +class SwizzleTestSuite + : public ::testing::TestWithParam< + std::tuple, std::pair, bool>> {}; TEST_P(SwizzleTestSuite, TestSwizzle) { @@ -328,7 +470,6 @@ TEST_P(SwizzleTestSuite, TestSwizzle) { const auto num_tiles = std::get<0>(GetParam()); const auto scaling_mode = std::get<1>(GetParam()); const auto transa = std::get<2>(GetParam()); - const PerTileSwizzle per_tile_swizzle(std::get<3>(GetParam())); performTestSwizzle1D(num_tiles.first, num_tiles.second, scaling_mode.first, scaling_mode.second, @@ -1007,14 +1148,15 @@ std::vector> num_tiles = { {65, 257}, {65, 258}, {65, 259}, - // Additional narrow-path coverage: narrow_k (row) when num_tiles_K < 32, - // narrow_m (col) when num_tiles_M < 32. - {1, 4}, // narrow_k with 4 K-tiles - {1, 8}, // narrow_k with 8 K-tiles - {4, 1}, // narrow_m with 4 M-tiles - {8, 1}, // narrow_m with 8 M-tiles - {31, 1}, // narrow_m at boundary (31 < TB_DIM=32) - {1, 31}, // narrow_k at boundary (31 < TB_DIM=32) + {1, 4}, + {1, 8}, + {4, 1}, + {8, 1}, + {31, 1}, + {1, 31}, + {2, 63}, // Above 30 KiB: regular fallback. + {2, 64}, // Above 32 KiB: regular fallback. + {16, 56}, // 28.5 KiB: rowwise full-width kernel. {1600, 1}, // MXFP8: more row-wise tiles than one CTA per tile allows }; @@ -1051,8 +1193,7 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Combine( ::testing::ValuesIn(num_tiles), ::testing::ValuesIn(scaling_mode), - ::testing::ValuesIn(transa), - ::testing::Values(false, true) + ::testing::ValuesIn(transa) ), [](const testing::TestParamInfo& info) { std::string name = "ntiles" + @@ -1060,8 +1201,7 @@ INSTANTIATE_TEST_SUITE_P( std::to_string(std::get<0>(info.param).second) + "smode" + std::to_string(std::get<1>(info.param).first) + "X"+ std::to_string(std::get<1>(info.param).second) + "trans" + - std::to_string(std::get<2>(info.param)) + - (std::get<3>(info.param) ? "PerTile" : ""); + std::to_string(std::get<2>(info.param)); return name; }); diff --git a/transformer_engine/common/include/transformer_engine/swizzle.h b/transformer_engine/common/include/transformer_engine/swizzle.h index f85d88fdce8..1c54e5f8981 100644 --- a/transformer_engine/common/include/transformer_engine/swizzle.h +++ b/transformer_engine/common/include/transformer_engine/swizzle.h @@ -26,6 +26,7 @@ extern "C" { * Requirements: * - scale_inv is stored in row-major. * - scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale. + * - Input and output scale_inv buffer pointers used by this operation must be 16-byte aligned. * - data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. */ void nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, cudaStream_t stream); @@ -40,6 +41,7 @@ void nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, cud * Requirements: * - scale_inv is stored in row-major. * - scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale. + * - Input and output scale_inv buffer pointers used by this operation must be 16-byte aligned. * - data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. */ void nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETensor* outputs, @@ -50,6 +52,7 @@ void nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETen * * Use this variant when the data and scale_inv tensors intentionally have * different shapes, e.g. when scale_invs have been transposed for attention. + * The 16-byte alignment requirement for scale_inv buffer pointers used by this operation still applies. */ void nvte_multi_tensor_swizzle_scaling_factors_unchecked(const NVTETensor* inputs, NVTETensor* outputs, diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 19fb8e848ad..7089a981dc2 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -33,6 +33,38 @@ constexpr __device__ __host__ int N_SF_PER_TD_PER_TILE = 4; constexpr __device__ __host__ int NEW_SF_TILE_DIM_K_I32 = NEW_SF_TILE_DIM_K / 4; constexpr __device__ __host__ int NEW_SF_TILE_DIM_M_I32 = 32; +// Shared layout for the optimized rowwise full-width and columnwise chunk kernels. Their local K +// dimension indexes blocks along the quantized axis; local M is the other axis. An M tile spans +// 128 rows in rowwise scales or 128 columns in columnwise scales. Each scale is one byte; a +// uint32_t word packs four scales. One output block holds 128 M positions by four K scale +// positions, or 512 bytes. +constexpr int kSwizzleTileDimM = 128; +constexpr int kScalesPerWord = 4; +constexpr int kBytesPerWord = sizeof(uint32_t); +// One byte per M position makes a 128-position span eight 16-byte vectors. +constexpr int kVectorsPerMSpan = kSwizzleTileDimM / sizeof(int4); +constexpr int kMPositionsPerSwizzleGroup = kSwizzleTileDimM / kScalesPerWord; +constexpr int kSwizzleOutputBlockBytes = kSwizzleTileDimM * kScalesPerWord; + +// Both optimized multi-tensor variants store at most this many descriptors per launch. +constexpr int kMaxTensorsPerLaunch = 64; + +// Rowwise full-width kernels and their dispatch thresholds. Narrow rows always use this kernel; +// wider rows use it for tiny grids or grids large enough to occupy the GPU. +constexpr int kRowwiseThreadsPerBlock = 128; +constexpr int kRowwiseMaxSharedMemoryBytes = 32 * 1024; +constexpr int kRowwiseNarrowSharedMemoryBytes = 17 * 1024; +constexpr int kRowwiseWideSharedMemoryBytes = 30 * 1024; +constexpr int kRowwiseMaxTinyTiles = 16; +// Largest rowwise grid before several row tiles share a block. +constexpr int kRowwiseMaxBlocks = 1536; + +// Columnwise chunk kernels and the dispatch threshold for switching from 32 to 128 scale rows. +constexpr int kColumnwiseSmallChunkRows = 32; +constexpr int kColumnwiseLargeChunkRows = 128; +constexpr int kColumnwiseThreadsPerBlock = 256; +constexpr int kColumnwiseMaxBlocksWith32RowChunks = 1024; + template __device__ inline void regs_shuffle_with_bit_shifts(LType* regs_vec) { // inp, 4-byte chunks [0,1,2,3, 4,5,6,7, 8,9,10,11, 12,13,14,15] @@ -223,17 +255,6 @@ __device__ __forceinline__ void dispatch_swizzle_col_scaling_kernel_impl( } } -template -__global__ void __launch_bounds__(TB_DIM* TB_DIM) - swizzle_col_scaling_kernel(const void* input, void* output, const int M, const int K, - const int original_M, const int original_K) { - const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M); - const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K); - dispatch_swizzle_col_scaling_kernel_impl( - input, output, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y, - padding_k, padding_m); -} - template __device__ inline void regs_shuffle(LType* regs_vec) { constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int); @@ -394,163 +415,6 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) padding_k, padding_m); } -// Narrow-K specialization for row scaling swizzle. -// When K is small (num_tiles_k < TB_DIM), the standard kernel wastes threadIdx.x -// because there aren't enough K-tiles to distribute across threads. -// This kernel repurposes the thread dimensions: threadIdx.x iterates rows within -// an M-tile, threadIdx.y indexes M-tiles within the block, processing TB_DIM -// M-tiles per block with full thread utilization. -template -__device__ void swizzle_row_scaling_narrow_k_kernel_impl(const void* input, void* output, - const int M, const int K, - const int original_M, const int original_K, - const int bid, const int grid_dim) { - constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4; - const int K_i32 = K / 4; - const int num_tiles_m = M / SF_TILE_DIM_M; - - const int m_tile = bid * blockDim.y + threadIdx.y; - const bool active = (m_tile < num_tiles_m); - - extern __shared__ int4 slm_v4i[]; - const int slm_tile_v4i = K_i32 * (SF_TILE_SIZE_I32 / 4); - - if (active) { - const bool padding_m = (m_tile == num_tiles_m - 1) && (original_M < M); - const bool padding_k = (original_K < K); - - int4* my_slm = slm_v4i + threadIdx.y * slm_tile_v4i; - - for (int k = 0; k < K_i32; k++) { - const int input_base = m_tile * SF_TILE_DIM_M * K_i32 + k; - const int* input_i32 = reinterpret_cast(input) + input_base; - - int regs[N_SF_PER_TD_PER_TILE]; -#pragma unroll - for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) { - const int row = i * TB_DIM + threadIdx.x; - regs[i] = __ldg(input_i32 + row * K_i32); - if (padding_m || padding_k) { - for (int j = 0; j < 4; j++) { - const int byte_row = m_tile * SF_TILE_DIM_M + row; - const int byte_col = k * 4 + j; - if (byte_row >= original_M || byte_col >= original_K) { - reinterpret_cast(®s[i])[j] = 0; - } - } - } - } - - my_slm[k * (SF_TILE_SIZE_I32 / 4) + threadIdx.x] = *reinterpret_cast(regs); - } - } - - __syncthreads(); - - if (active) { - int4* my_slm = slm_v4i + threadIdx.y * slm_tile_v4i; - int4* out_v4i = - reinterpret_cast(reinterpret_cast(output) + m_tile * SF_TILE_DIM_M * K_i32); - - for (int i = threadIdx.x; i < slm_tile_v4i; i += blockDim.x) { - out_v4i[i] = my_slm[i]; - } - } -} - -template -__global__ void __launch_bounds__(TB_DIM* TB_DIM) - swizzle_row_scaling_narrow_k_kernel(const void* input, void* output, const int M, const int K, - const int original_M, const int original_K) { - swizzle_row_scaling_narrow_k_kernel_impl( - input, output, M, K, original_M, original_K, blockIdx.x, gridDim.x); -} - -// Narrow-M variant of the column scaling swizzle kernel, for when num_tiles_m < TB_DIM. -// Analogous to the narrow-K row kernel: when the M dimension is small, the normal -// col kernel underutilizes threads in the load phase because threadIdx.x covers M -// positions with vectorized loads, leaving many threads idle. This kernel repurposes -// thread dimensions: threadIdx.y indexes K-tiles within the block, threadIdx.x covers -// one int32 column of an M-tile, and M-tiles are iterated serially. -template -__device__ void swizzle_col_scaling_narrow_m_kernel_impl(const void* input, void* output, - const int M, const int K, - const int original_M, const int original_K, - const int bid, const int grid_dim) { - constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4; - constexpr int SF_TILE_DIM_M_I32 = SF_TILE_DIM_M / 4; - constexpr int SF_TILE_DIM_K_I32 = SF_TILE_DIM_K; - - const int M_i32 = M / 4; - const int K_i32 = K; - const int num_tiles_m = M / SF_TILE_DIM_M; - const int num_tiles_k = K / SF_TILE_DIM_K; - - const int k_tile = bid * blockDim.y + threadIdx.y; - const bool active = (k_tile < num_tiles_k); - const int remaining = num_tiles_k - bid * static_cast(blockDim.y); - const int k_tiles_in_block = remaining <= 0 ? 0 : (remaining < TB_DIM ? remaining : TB_DIM); - - extern __shared__ int slm_narrow_m[]; - - if (active) { - const bool padding_k = (k_tile == num_tiles_k - 1) && (original_K < K); - const int32_t* input_i32 = reinterpret_cast(input); - - for (int m_tile = 0; m_tile < num_tiles_m; m_tile++) { - const bool padding_m = (m_tile == num_tiles_m - 1) && (original_M < M); - - int regs[N_SF_PER_TD_PER_TILE]; -#pragma unroll - for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) { - const int k_row = k_tile * SF_TILE_DIM_K_I32 + i; - const int m_col = m_tile * SF_TILE_DIM_M_I32 + threadIdx.x; - regs[i] = __ldg(input_i32 + k_row * M_i32 + m_col); - if (padding_m || padding_k) { - for (int j = 0; j < 4; j++) { - if (m_col * 4 + j >= original_M || k_row >= original_K) { - reinterpret_cast(®s[i])[j] = 0; - } - } - } - } - - regs_shuffle_with_bit_shifts(regs); - - int tM = threadIdx.x * N_SF_PER_TD_PER_TILE; - int* slm_tile = - slm_narrow_m + m_tile * TB_DIM * SF_TILE_SIZE_I32 + threadIdx.y * SF_TILE_SIZE_I32; -#pragma unroll - for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) { - slm_tile[(tM % SF_TILE_DIM_M) / NEW_SF_TILE_DIM_M_I32 + - ((tM + i) % NEW_SF_TILE_DIM_M_I32) * NEW_SF_TILE_DIM_K_I32] = regs[i]; - } - } - } - - __syncthreads(); - - const int linear_id = threadIdx.y * blockDim.x + threadIdx.x; - for (int m_tile = 0; m_tile < num_tiles_m; m_tile++) { - int4* out_v4i = reinterpret_cast(reinterpret_cast(output) + - m_tile * SF_TILE_DIM_M_I32 * K_i32 + - bid * TB_DIM * SF_TILE_SIZE_I32); - int4* slm_v4i = reinterpret_cast(slm_narrow_m + m_tile * TB_DIM * SF_TILE_SIZE_I32); - const int n_v4i = k_tiles_in_block * SF_TILE_SIZE_I32 / 4; - for (int j = linear_id; j < n_v4i; j += blockDim.x * blockDim.y) { - out_v4i[j] = slm_v4i[j]; - } - } -} - -template -__global__ void __launch_bounds__(TB_DIM* TB_DIM) - swizzle_col_scaling_narrow_m_kernel(const void* input, void* output, const int M, const int K, - const int original_M, const int original_K) { - swizzle_col_scaling_narrow_m_kernel_impl( - input, output, M, K, original_M, original_K, blockIdx.x, gridDim.x); -} - template __device__ void unswizzle_row_scaling_kernel_impl(const void* input, void* output, const int M, const int K, const int bid_x, const int bid_y, @@ -712,10 +576,6 @@ struct MultiSwizzleArgs { int num_tensors; }; -constexpr size_t round_up_to_multiple(size_t value, size_t multiple) { - return DIVUP(value, multiple) * multiple; -} - template __global__ void __launch_bounds__(TB_DIM* TB_DIM) grouped_swizzle_row_scaling_uniform_shape_kernel(const void* input, void* output, const int M, @@ -867,108 +727,16 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) padding_m); } -template -__global__ void __launch_bounds__(TB_DIM* TB_DIM) - multi_tensor_swizzle_col_scaling_kernel(MultiSwizzleArgs kernel_args) { - // Find tensor corresponding to block - const int bid = blockIdx.x; - int tensor_id = 0; - while (kernel_args.block_range[tensor_id + 1] <= bid) { - ++tensor_id; - } - // Get args corresponding to block - 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]; - const int original_M = kernel_args.original_m_list[tensor_id]; - const int original_K = kernel_args.original_k_list[tensor_id]; - - constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int); - - // Get block index in grid. Emulate 2D grid. - const int num_tiles_k = K / SF_TILE_DIM_K; - const int num_tiles_m = M / SF_TILE_DIM_M; - const int grid_dim_x = DIVUP(num_tiles_k, TB_DIM); - const int grid_dim_y = DIVUP(num_tiles_m, N_TILE_PER_TD); - const int bid_x = (bid - kernel_args.block_range[tensor_id]) / grid_dim_y; - const int bid_y = (bid - kernel_args.block_range[tensor_id]) % grid_dim_y; - - const bool padding_m = (bid_y == grid_dim_y - 1) && (original_M < M); - const bool padding_k = (bid_x == grid_dim_x - 1) && (original_K < K); - dispatch_swizzle_col_scaling_kernel_impl( - input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y, padding_k, - padding_m); -} - -template -__global__ void __launch_bounds__(TB_DIM* TB_DIM) - multi_tensor_swizzle_row_scaling_narrow_k_kernel(MultiSwizzleArgs kernel_args) { - const int bid = blockIdx.x; - int tensor_id = 0; - while (kernel_args.block_range[tensor_id + 1] <= bid) { - ++tensor_id; - } - 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]; - const int original_M = kernel_args.original_m_list[tensor_id]; - const int original_K = kernel_args.original_k_list[tensor_id]; - const int flat_bid = bid - kernel_args.block_range[tensor_id]; - const int num_tiles_m = M / SF_TILE_DIM_M; - const int grid_dim = DIVUP(num_tiles_m, TB_DIM); - - swizzle_row_scaling_narrow_k_kernel_impl( - input, output, M, K, original_M, original_K, flat_bid, grid_dim); -} - -template -__global__ void __launch_bounds__(TB_DIM* TB_DIM) - multi_tensor_swizzle_col_scaling_narrow_m_kernel(MultiSwizzleArgs kernel_args) { - const int bid = blockIdx.x; - int tensor_id = 0; - while (kernel_args.block_range[tensor_id + 1] <= bid) { - ++tensor_id; - } - 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]; - const int original_M = kernel_args.original_m_list[tensor_id]; - const int original_K = kernel_args.original_k_list[tensor_id]; - const int flat_bid = bid - kernel_args.block_range[tensor_id]; - const int num_tiles_k = K / SF_TILE_DIM_K; - const int grid_dim = DIVUP(num_tiles_k, TB_DIM); - - swizzle_col_scaling_narrow_m_kernel_impl( - input, output, M, K, original_M, original_K, flat_bid, grid_dim); -} - -// MXFP8 scale swizzle with one thread block per scale tile, used when -// NVTE_MXFP8_PER_TILE_SWIZZLE=1. The narrow-K / narrow-M kernels above put 32 tiles in one -// 1024-thread block, so e.g. the scales of a 4096-row activation are swizzled by a single block. -// These kernels give every 128-row tile (row-wise) or 128-column x 32-row block (column-wise) its -// own block, which moves a contiguous chunk with coalesced 16-byte loads and stores through shared -// memory. Multi-tensor launches put several tiles in one block only when the grid would otherwise -// be very large. -constexpr int kRowwiseThreadsPerBlock = 128; -constexpr int kColumnwiseThreadsPerBlock = 256; -constexpr int kMaxTensorsPerLaunch = 64; -// Row-wise: the largest grid before several tiles share a block. -constexpr int kRowwiseMaxBlocks = 1536; -// Column-wise: the largest grid with 32-row chunks before switching to 128-row chunks. -constexpr int kColumnwiseMaxBlocksWith32RowChunks = 1024; - -// Whether MXFP8 swizzles use the per-tile kernels (NVTE_MXFP8_PER_TILE_SWIZZLE=1). -bool use_per_tile_swizzle() { return getenv("NVTE_MXFP8_PER_TILE_SWIZZLE"); } - +// Optimized scale swizzle for narrow scale widths. "Rowwise full width" means that a block stages +// the entire padded K dimension for one or more 128-row tiles. "Columnwise chunk" means that a +// block stages only 32 (or 128 for large grids) K rows for one 128-column M tile. Both paths use +// coalesced 16-byte loads and stores through shared memory. // The 4 scales packed in `word` (the first one in column first_column), with those in columns // >= num_columns set to zero. __device__ __forceinline__ uint32_t zero_padding_columns(uint32_t word, int first_column, int num_columns) { const int valid_columns = num_columns - first_column; - if (valid_columns >= 4) return word; + if (valid_columns >= kScalesPerWord) return word; if (valid_columns <= 0) return 0u; return word & ((1u << (8 * valid_columns)) - 1u); } @@ -980,31 +748,33 @@ __device__ __forceinline__ uint32_t divide_by_magic(uint32_t numerator, uint32_t } inline uint32_t division_magic(int divisor) { - return divisor <= 1 ? 0u : static_cast(((1ull << 32) + divisor - 1) / divisor); + return divisor <= 1 + ? 0u + : static_cast(DIVUP(uint64_t{1} << 32, static_cast(divisor))); } -// Row-wise: tiles [first_tile, first_tile + num_tiles) of one tensor, a contiguous chunk of the +// Row-wise: tiles [first_tile, first_tile + num_tiles) of one tensor, a consecutive chunk of the // input and of the output. Shared-memory rows are padded to an odd number of words, so the column // gathers are free of bank conflicts. Requires num_tiles * 128 * padded_k / 4 < 2^16 (guaranteed -// by the 48 KB shared-memory limit). +// by the row-wise shared-memory limit = 30 KB). __device__ __forceinline__ void swizzle_rowwise_tiles(const uint8_t* __restrict__ input, uint8_t* __restrict__ output, int padded_k, uint32_t words_per_row_magic, int original_m, int original_k, int first_tile, int num_tiles, uint32_t* smem) { - const int words_per_row = padded_k >> 2; + const int words_per_row = padded_k / kScalesPerWord; const int smem_row_words = words_per_row | 1; - const int num_vectors = - num_tiles * 8 * padded_k; // 16-byte vectors: num_tiles * 128 * padded_k / 16 - const size_t chunk_offset = static_cast(first_tile) * 128 * padded_k; + const int num_vectors = num_tiles * kVectorsPerMSpan * padded_k; + const size_t chunk_offset = static_cast(first_tile) * kSwizzleTileDimM * padded_k; const int4* input_vectors = reinterpret_cast(input + chunk_offset); for (int i = threadIdx.x; i < num_vectors; i += kRowwiseThreadsPerBlock) { const int4 vector = __ldg(input_vectors + i); - const uint32_t words[4] = {static_cast(vector.x), static_cast(vector.y), - static_cast(vector.z), static_cast(vector.w)}; + const uint32_t words[kScalesPerWord] = { + static_cast(vector.x), static_cast(vector.y), + static_cast(vector.z), static_cast(vector.w)}; #pragma unroll - for (int w = 0; w < 4; w++) { - const uint32_t word_index = 4 * i + w; + for (int w = 0; w < kScalesPerWord; w++) { + const uint32_t word_index = kScalesPerWord * i + w; const uint32_t row = divide_by_magic(word_index, words_per_row_magic); smem[row * smem_row_words + (word_index - row * words_per_row)] = words[w]; } @@ -1013,22 +783,22 @@ __device__ __forceinline__ void swizzle_rowwise_tiles(const uint8_t* __restrict_ // An output tile holds padded_k / 4 blocks of 512 bytes, i.e. 32 vectors each; vector r of a // block packs rows r, r + 32, r + 64 and r + 96 of its 4 scale columns. - const int vectors_per_tile = 8 * padded_k; + const int vectors_per_tile = kVectorsPerMSpan * padded_k; for (int tile = 0; tile < num_tiles; tile++) { - int4* output_vectors = - reinterpret_cast(output + chunk_offset + static_cast(tile) * 128 * padded_k); - const int tile_first_row = tile * 128; - const int global_first_row = first_tile * 128 + tile_first_row; + int4* output_vectors = reinterpret_cast( + output + chunk_offset + static_cast(tile) * kSwizzleTileDimM * padded_k); + const int tile_first_row = tile * kSwizzleTileDimM; + const int global_first_row = first_tile * kSwizzleTileDimM + tile_first_row; for (int v = threadIdx.x; v < vectors_per_tile; v += kRowwiseThreadsPerBlock) { - const int k_block = v >> 5; - const int row_in_group = v & 31; - uint32_t words[4]; + const int k_block = v / kMPositionsPerSwizzleGroup; + const int row_in_group = v % kMPositionsPerSwizzleGroup; + uint32_t words[kScalesPerWord]; #pragma unroll - for (int group = 0; group < 4; group++) { - const int row = row_in_group + 32 * group; + for (int group = 0; group < kScalesPerWord; group++) { + const int row = row_in_group + kMPositionsPerSwizzleGroup * group; uint32_t word = smem[(tile_first_row + row) * smem_row_words + k_block]; if (global_first_row + row >= original_m) word = 0u; - words[group] = zero_padding_columns(word, k_block * 4, original_k); + words[group] = zero_padding_columns(word, k_block * kScalesPerWord, original_k); } output_vectors[v] = make_int4(static_cast(words[0]), static_cast(words[1]), static_cast(words[2]), static_cast(words[3])); @@ -1045,92 +815,97 @@ __device__ __forceinline__ void swizzle_columnwise_chunk(const uint8_t* __restri int padded_k, int original_m, int original_k, int m_tile, int first_k, int4* smem) { - const int num_rows = min(kChunkRows, padded_k - first_k); // multiple of 4 - const int num_vectors = num_rows * 8; + const int num_rows = min(kChunkRows, padded_k - first_k); // multiple of kScalesPerWord + const int num_vectors = num_rows * kVectorsPerMSpan; for (int i = threadIdx.x; i < num_vectors; i += kColumnwiseThreadsPerBlock) { - const int row = i >> 3; - const int vector_in_row = i & 7; - smem[i] = __ldg(reinterpret_cast( - input + static_cast(first_k + row) * padded_m + m_tile * 128) + - vector_in_row); + const int row = i / kVectorsPerMSpan; + const int vector_in_row = i % kVectorsPerMSpan; + smem[i] = + __ldg(reinterpret_cast(input + static_cast(first_k + row) * padded_m + + m_tile * kSwizzleTileDimM) + + vector_in_row); } __syncthreads(); const uint8_t* smem_bytes = reinterpret_cast(smem); - const size_t first_output_block = static_cast(m_tile) * (padded_k >> 2) + (first_k >> 2); + const size_t first_output_block = + static_cast(m_tile) * (padded_k / kScalesPerWord) + (first_k / kScalesPerWord); for (int v = threadIdx.x; v < num_vectors; v += kColumnwiseThreadsPerBlock) { - const int k_block = v >> 5; - const int column_in_group = v & 31; - uint32_t words[4]; + const int k_block = v / kMPositionsPerSwizzleGroup; + const int column_in_group = v % kMPositionsPerSwizzleGroup; + uint32_t words[kScalesPerWord]; #pragma unroll - for (int group = 0; group < 4; group++) { - const int column = column_in_group + 32 * group; + for (int group = 0; group < kScalesPerWord; group++) { + const int column = column_in_group + kMPositionsPerSwizzleGroup * group; uint32_t word = 0u; - if (m_tile * 128 + column < original_m) { + if (m_tile * kSwizzleTileDimM + column < original_m) { #pragma unroll - for (int j = 0; j < 4; j++) { - const int row = k_block * 4 + j; + for (int j = 0; j < kScalesPerWord; j++) { + const int row = k_block * kScalesPerWord + j; if (first_k + row < original_k) { - word |= static_cast(smem_bytes[row * 128 + column]) << (8 * j); + word |= static_cast(smem_bytes[row * kSwizzleTileDimM + column]) << (8 * j); } } } words[group] = word; } - reinterpret_cast(output + (first_output_block + k_block) * 512)[column_in_group] = + reinterpret_cast(output + (first_output_block + k_block) * + kSwizzleOutputBlockBytes)[column_in_group] = make_int4(static_cast(words[0]), static_cast(words[1]), static_cast(words[2]), static_cast(words[3])); } } __global__ void __launch_bounds__(kRowwiseThreadsPerBlock) - swizzle_rowwise_per_tile_kernel(const uint8_t* __restrict__ input, uint8_t* __restrict__ output, - int padded_k, uint32_t words_per_row_magic, int original_m, - int original_k) { + swizzle_rowwise_full_width_kernel(const uint8_t* __restrict__ input, + uint8_t* __restrict__ output, int padded_k, + uint32_t words_per_row_magic, int original_m, + int original_k) { extern __shared__ uint32_t rowwise_smem[]; swizzle_rowwise_tiles(input, output, padded_k, words_per_row_magic, original_m, original_k, blockIdx.x, 1, rowwise_smem); } __global__ void __launch_bounds__(kColumnwiseThreadsPerBlock) - swizzle_columnwise_per_tile_kernel(const uint8_t* __restrict__ input, - uint8_t* __restrict__ output, int padded_m, int padded_k, - int original_m, int original_k) { - __shared__ int4 columnwise_smem[32 * 8]; - swizzle_columnwise_chunk<32>(input, output, padded_m, padded_k, original_m, original_k, - blockIdx.x, blockIdx.y * 32, columnwise_smem); + swizzle_columnwise_chunk_kernel(const uint8_t* __restrict__ input, uint8_t* __restrict__ output, + int padded_m, int padded_k, int original_m, int original_k) { + __shared__ int4 columnwise_smem[kColumnwiseSmallChunkRows * kVectorsPerMSpan]; + swizzle_columnwise_chunk( + input, output, padded_m, padded_k, original_m, original_k, blockIdx.x, + blockIdx.y * kColumnwiseSmallChunkRows, columnwise_smem); } // The scales of one tensor: padded_m (a multiple of 128) by padded_k (a multiple of 4) in // row-wise terms; for column-wise scales, m runs along the columns of the data. Scales at // m >= original_m or k >= original_k are padding and written as zeros. -struct PerTileSwizzleTensor { +struct SwizzleTensorDescriptor { const uint8_t* input; uint8_t* output; int padded_m, padded_k, original_m, original_k; }; -struct PerTileMultiSwizzleArgs { - PerTileSwizzleTensor tensors[kMaxTensorsPerLaunch]; +struct MultiTensorSwizzleArgs { + SwizzleTensorDescriptor tensors[kMaxTensorsPerLaunch]; int first_block[kMaxTensorsPerLaunch + 1]; uint32_t words_per_row_magic[kMaxTensorsPerLaunch]; // row-wise only - int tiles_per_block; // row-wise only + int row_tiles_per_block; // row-wise only int num_tensors; }; -__device__ __forceinline__ int tensor_of_block(const PerTileMultiSwizzleArgs& args, int block) { +__device__ __forceinline__ int tensor_of_block(const MultiTensorSwizzleArgs& args, int block) { int tensor = 0; while (tensor + 1 < args.num_tensors && args.first_block[tensor + 1] <= block) tensor++; return tensor; } __global__ void __launch_bounds__(kRowwiseThreadsPerBlock) - multi_tensor_swizzle_rowwise_per_tile_kernel(const PerTileMultiSwizzleArgs args) { + multi_tensor_swizzle_rowwise_full_width_kernel(const MultiTensorSwizzleArgs args) { extern __shared__ uint32_t rowwise_multi_smem[]; const int t = tensor_of_block(args, blockIdx.x); - const int tiles_per_block = args.tiles_per_block; - const int first_tile = (blockIdx.x - args.first_block[t]) * tiles_per_block; - const int num_tiles = min(tiles_per_block, args.tensors[t].padded_m / 128 - first_tile); + const int row_tiles_per_block = args.row_tiles_per_block; + const int first_tile = (blockIdx.x - args.first_block[t]) * row_tiles_per_block; + const int num_tiles = + min(row_tiles_per_block, args.tensors[t].padded_m / kSwizzleTileDimM - first_tile); swizzle_rowwise_tiles(args.tensors[t].input, args.tensors[t].output, args.tensors[t].padded_k, args.words_per_row_magic[t], args.tensors[t].original_m, args.tensors[t].original_k, first_tile, num_tiles, rowwise_multi_smem); @@ -1138,11 +913,11 @@ __global__ void __launch_bounds__(kRowwiseThreadsPerBlock) template __global__ void __launch_bounds__(kColumnwiseThreadsPerBlock) - multi_tensor_swizzle_columnwise_per_tile_kernel(const PerTileMultiSwizzleArgs args) { - __shared__ int4 columnwise_multi_smem[kChunkRows * 8]; + multi_tensor_swizzle_columnwise_chunk_kernel(const MultiTensorSwizzleArgs args) { + __shared__ int4 columnwise_multi_smem[kChunkRows * kVectorsPerMSpan]; const int t = tensor_of_block(args, blockIdx.x); const int block_in_tensor = blockIdx.x - args.first_block[t]; - const int num_m_tiles = args.tensors[t].padded_m / 128; + const int num_m_tiles = args.tensors[t].padded_m / kSwizzleTileDimM; const int m_tile = block_in_tensor % num_m_tiles; const int first_k = (block_in_tensor / num_m_tiles) * kChunkRows; swizzle_columnwise_chunk(args.tensors[t].input, args.tensors[t].output, @@ -1155,116 +930,119 @@ inline bool is_16_byte_aligned(const void* ptr) { return (reinterpret_cast(ptr) & 15u) == 0; } -inline bool per_tile_swizzle_supported(const PerTileSwizzleTensor& tensor, bool rowwise) { - if (tensor.padded_m <= 0 || tensor.padded_k <= 0 || tensor.padded_m % 128 != 0 || - tensor.padded_k % 4 != 0) { - return false; +inline void check_scale_inv_alignment(const Tensor* tensor, const char* name, bool rowwise) { + const auto& scales = rowwise ? tensor->scale_inv : tensor->columnwise_scale_inv; + if (scales.has_data()) { + NVTE_CHECK(is_16_byte_aligned(scales.dptr), name, rowwise ? " row-wise" : " column-wise", + " scale_inv buffer must be 16-byte aligned."); } - if (!is_16_byte_aligned(tensor.input) || !is_16_byte_aligned(tensor.output)) return false; - if (rowwise && 128 * ((tensor.padded_k >> 2) | 1) * 4 > 48 * 1024) return false; - return true; } -// One kernel for up to kMaxTensorsPerLaunch tensors, all row-wise or all column-wise. Row-wise: -// the tiles per block are the smallest power of two that keeps the grid within kRowwiseMaxBlocks -// (shared memory permitting). Column-wise: 128-row chunks once 32-row chunks would exceed -// kColumnwiseMaxBlocksWith32RowChunks blocks. Returns false, without launching anything, if any -// tensor is not supported. -inline bool launch_multi_tensor_per_tile_swizzle(const PerTileSwizzleTensor* tensors, - int num_tensors, bool rowwise, - cudaStream_t stream) { - if (num_tensors <= 0 || num_tensors > kMaxTensorsPerLaunch) return false; - PerTileMultiSwizzleArgs args; +inline bool should_use_rowwise_full_width_swizzle(int padded_k, int64_t total_row_tiles) { + const int smem_bytes = kSwizzleTileDimM * ((padded_k / kScalesPerWord) | 1) * kBytesPerWord; + if (smem_bytes <= kRowwiseNarrowSharedMemoryBytes) return true; + // The fallback launches one block per row tile at these widths. Switch after its grid + // exceeds one wave on the current GPU. + return smem_bytes <= kRowwiseWideSharedMemoryBytes && + (total_row_tiles <= kRowwiseMaxTinyTiles || total_row_tiles > cuda::sm_count()); +} + +// Launches the rowwise full-width kernel, where each block stages the entire padded K dimension, +// for up to kMaxTensorsPerLaunch tensors. The row tiles per block are the smallest power of two +// that keeps the grid within kRowwiseMaxBlocks, subject to the shared-memory limit. +inline void launch_multi_tensor_rowwise_full_width_swizzle(const SwizzleTensorDescriptor* tensors, + int num_tensors, cudaStream_t stream) { + MultiTensorSwizzleArgs args{}; args.num_tensors = num_tensors; args.first_block[0] = 0; - int64_t total_tiles = 0; + int64_t total_row_tiles = 0; int max_smem_row_words = 1; for (int i = 0; i < num_tensors; i++) { - if (!per_tile_swizzle_supported(tensors[i], rowwise)) return false; args.tensors[i] = tensors[i]; - args.words_per_row_magic[i] = division_magic(tensors[i].padded_k >> 2); - if (rowwise) { - total_tiles += tensors[i].padded_m / 128; - max_smem_row_words = max(max_smem_row_words, (tensors[i].padded_k >> 2) | 1); - } else { - total_tiles += - static_cast(tensors[i].padded_m / 128) * ((tensors[i].padded_k + 31) / 32); - } + args.words_per_row_magic[i] = division_magic(tensors[i].padded_k / kScalesPerWord); + total_row_tiles += tensors[i].padded_m / kSwizzleTileDimM; + max_smem_row_words = max(max_smem_row_words, (tensors[i].padded_k / kScalesPerWord) | 1); + } + int row_tiles_per_block = 1; + const int max_row_tiles_per_block = + kRowwiseMaxSharedMemoryBytes / (kSwizzleTileDimM * max_smem_row_words * kBytesPerWord); + while (total_row_tiles > static_cast(kRowwiseMaxBlocks) * row_tiles_per_block && + row_tiles_per_block * 2 <= max_row_tiles_per_block) { + row_tiles_per_block *= 2; + } + args.row_tiles_per_block = row_tiles_per_block; + for (int i = 0; i < num_tensors; i++) { + const SwizzleTensorDescriptor& tensor = tensors[i]; + const int blocks = DIVUP(tensor.padded_m / kSwizzleTileDimM, row_tiles_per_block); + args.first_block[i + 1] = args.first_block[i] + blocks; } - int tiles_per_block = 1, chunk_rows = 32; - if (rowwise) { - const int max_tiles_per_block = (48 * 1024) / (128 * max_smem_row_words * 4); - while (total_tiles > static_cast(kRowwiseMaxBlocks) * tiles_per_block && - tiles_per_block * 2 <= max_tiles_per_block) { - tiles_per_block *= 2; - } - } else { - if (total_tiles > kColumnwiseMaxBlocksWith32RowChunks) chunk_rows = 128; + for (int i = num_tensors + 1; i <= kMaxTensorsPerLaunch; i++) { + args.first_block[i] = args.first_block[num_tensors]; } - args.tiles_per_block = tiles_per_block; + const int grid = args.first_block[num_tensors]; + if (grid == 0) return; + multi_tensor_swizzle_rowwise_full_width_kernel<<< + grid, kRowwiseThreadsPerBlock, + row_tiles_per_block * kSwizzleTileDimM * max_smem_row_words * kBytesPerWord, stream>>>(args); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +// Launches the columnwise chunk kernel, where each block stages part of K for one 128-column M +// tile, for up to kMaxTensorsPerLaunch tensors. It uses 128-row K chunks once 32-row chunks would +// exceed kColumnwiseMaxBlocksWith32RowChunks blocks. +inline void launch_multi_tensor_columnwise_chunk_swizzle(const SwizzleTensorDescriptor* tensors, + int num_tensors, cudaStream_t stream) { + MultiTensorSwizzleArgs args{}; + args.num_tensors = num_tensors; + args.first_block[0] = 0; + int64_t total_chunks = 0; for (int i = 0; i < num_tensors; i++) { - const PerTileSwizzleTensor& tensor = tensors[i]; - const int blocks = - rowwise ? (tensor.padded_m / 128 + tiles_per_block - 1) / tiles_per_block - : (tensor.padded_m / 128) * ((tensor.padded_k + chunk_rows - 1) / chunk_rows); + args.tensors[i] = tensors[i]; + total_chunks += static_cast(tensors[i].padded_m / kSwizzleTileDimM) * + DIVUP(tensors[i].padded_k, kColumnwiseSmallChunkRows); + } + const int chunk_rows = total_chunks > kColumnwiseMaxBlocksWith32RowChunks + ? kColumnwiseLargeChunkRows + : kColumnwiseSmallChunkRows; + for (int i = 0; i < num_tensors; i++) { + const SwizzleTensorDescriptor& tensor = tensors[i]; + const int blocks = (tensor.padded_m / kSwizzleTileDimM) * DIVUP(tensor.padded_k, chunk_rows); args.first_block[i + 1] = args.first_block[i] + blocks; } for (int i = num_tensors + 1; i <= kMaxTensorsPerLaunch; i++) { args.first_block[i] = args.first_block[num_tensors]; } const int grid = args.first_block[num_tensors]; - if (grid == 0) return true; - if (rowwise) { - multi_tensor_swizzle_rowwise_per_tile_kernel<<< - grid, kRowwiseThreadsPerBlock, tiles_per_block * 128 * max_smem_row_words * 4, stream>>>( - args); - } else if (chunk_rows == 128) { - multi_tensor_swizzle_columnwise_per_tile_kernel<128> + if (grid == 0) return; + if (chunk_rows == kColumnwiseLargeChunkRows) { + multi_tensor_swizzle_columnwise_chunk_kernel <<>>(args); } else { - multi_tensor_swizzle_columnwise_per_tile_kernel<32> + multi_tensor_swizzle_columnwise_chunk_kernel <<>>(args); } NVTE_CHECK_CUDA(cudaGetLastError()); - return true; } -inline bool launch_rowwise_per_tile_swizzle(const PerTileSwizzleTensor& tensor, - cudaStream_t stream) { - if (!per_tile_swizzle_supported(tensor, true)) return false; - swizzle_rowwise_per_tile_kernel<<> 2) | 1) * 4, stream>>>( - tensor.input, tensor.output, tensor.padded_k, division_magic(tensor.padded_k >> 2), - tensor.original_m, tensor.original_k); +// Each block stages the full padded K dimension for one 128-row tile. +inline void launch_rowwise_full_width_swizzle(const SwizzleTensorDescriptor& tensor, + cudaStream_t stream) { + const int grid = tensor.padded_m / kSwizzleTileDimM; + const int smem = kSwizzleTileDimM * ((tensor.padded_k / kScalesPerWord) | 1) * kBytesPerWord; + swizzle_rowwise_full_width_kernel<<>>( + tensor.input, tensor.output, tensor.padded_k, + division_magic(tensor.padded_k / kScalesPerWord), tensor.original_m, tensor.original_k); NVTE_CHECK_CUDA(cudaGetLastError()); - return true; } -inline bool launch_columnwise_per_tile_swizzle(const PerTileSwizzleTensor& tensor, - cudaStream_t stream) { - if (!per_tile_swizzle_supported(tensor, false)) return false; - dim3 grid(tensor.padded_m / 128, (tensor.padded_k + 31) / 32); - swizzle_columnwise_per_tile_kernel<<>>( +// Each block stages a 32-row K chunk for one 128-column M tile. +inline void launch_columnwise_chunk_swizzle(const SwizzleTensorDescriptor& tensor, + cudaStream_t stream) { + dim3 grid(tensor.padded_m / kSwizzleTileDimM, DIVUP(tensor.padded_k, kColumnwiseSmallChunkRows)); + swizzle_columnwise_chunk_kernel<<>>( tensor.input, tensor.output, tensor.padded_m, tensor.padded_k, tensor.original_m, tensor.original_k); NVTE_CHECK_CUDA(cudaGetLastError()); - return true; -} - -// A single tensor uses the one-tile-per-block kernels unless its grid would be large; several -// tensors go through one multi-tensor launch. -inline bool launch_per_tile_swizzle(const PerTileSwizzleTensor* tensors, int num_tensors, - bool rowwise, cudaStream_t stream) { - const bool grid_too_large = - num_tensors == 1 && (rowwise ? tensors[0].padded_m / 128 > kRowwiseMaxBlocks - : static_cast(tensors[0].padded_m / 128) * - ((tensors[0].padded_k + 31) / 32) > - kColumnwiseMaxBlocksWith32RowChunks); - if (num_tensors == 1 && !grid_too_large) { - return rowwise ? launch_rowwise_per_tile_swizzle(tensors[0], stream) - : launch_columnwise_per_tile_swizzle(tensors[0], stream); - } - return launch_multi_tensor_per_tile_swizzle(tensors, num_tensors, rowwise, stream); } } // namespace @@ -1275,7 +1053,6 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s NVTE_CHECK(scaling_mode == NVTE_MXFP8_1D_SCALING || scaling_mode == NVTE_NVFP4_1D_SCALING, "Input tensor has invalid scaling mode (", to_string(input->scaling_mode), ")."); - // Check tensors CheckInputTensor(*input, "scaling_factor_input"); CheckInputTensor(*output, "scaling_factor_output"); NVTE_CHECK(!input->with_gemm_swizzled_scales, @@ -1341,9 +1118,10 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s NVTE_ERROR("Invalid scaling mode"); } - // Check dims constexpr int SF_TILE_DIM_M = 128; constexpr int SF_TILE_DIM_K = 4; + + // Check dims 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!"); @@ -1367,6 +1145,9 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s return; } + check_scale_inv_alignment(input, "Input", has_rowwise_scale_inv); + check_scale_inv_alignment(output, "Output", has_rowwise_scale_inv); + // Choose swizzle implementation bool rowwise_swizzle{false}, columnwise_swizzle{false}; switch (scaling_mode) { @@ -1420,31 +1201,24 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s NVTE_ERROR("Invalid scaling mode"); } - const int narrow_k_slm_size = - TB_DIM * num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t)); - const PerTileSwizzleTensor per_tile_tensor{static_cast(input_scale_inv_ptr), - static_cast(output_scale_inv_ptr), - m, - k, - original_M, - original_K}; - if (scaling_mode == NVTE_MXFP8_1D_SCALING && use_per_tile_swizzle() && - per_tile_swizzle_supported(per_tile_tensor, true)) { - launch_per_tile_swizzle(&per_tile_tensor, 1, true, stream); - } else if (num_tiles_k < TB_DIM && static_cast(narrow_k_slm_size) <= - cuda::max_shared_memory_per_block_optin()) { - // Narrow-K: batch TB_DIM M-tiles per block, fully utilizing all threads. - dim3 num_blocks_narrow(DIVUP(num_tiles_m, TB_DIM)); - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_row_scaling_narrow_k_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, narrow_k_slm_size)); - swizzle_row_scaling_narrow_k_kernel - <<>>( - input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K); + int vec_load_size = (num_tiles_k - 1) % 4 + 1; + /* there is no int3 and misaligned if using int4/int2 */ + if (vec_load_size == 3) vec_load_size = 1; + // Shared by the row-wise dispatch heuristic and optimized kernels. + const SwizzleTensorDescriptor descriptor{static_cast(input_scale_inv_ptr), + static_cast(output_scale_inv_ptr), + m, + k, + original_M, + original_K}; + if (should_use_rowwise_full_width_swizzle(descriptor.padded_k, + descriptor.padded_m / SF_TILE_DIM_M)) { + if (m / SF_TILE_DIM_M > kRowwiseMaxBlocks) { + launch_multi_tensor_rowwise_full_width_swizzle(&descriptor, 1, stream); + } else { + launch_rowwise_full_width_swizzle(descriptor, stream); + } } else { - int vec_load_size = (num_tiles_k - 1) % 4 + 1; - /* there is no int3 and misaligned if using int4/int2 */ - if (vec_load_size == 3) vec_load_size = 1; int n_tiles_in_tb = TB_DIM * vec_load_size; dim3 num_blocks(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m); int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); @@ -1487,77 +1261,27 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s const int original_M = input->flat_last_dim(); const int original_K = input->flat_first_dim() / MXFP8_BLOCK_SIZE; - const int narrow_m_slm_size = - TB_DIM * num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t)); - const PerTileSwizzleTensor per_tile_tensor{ + const SwizzleTensorDescriptor descriptor{ static_cast(input->columnwise_scale_inv.dptr), static_cast(output->columnwise_scale_inv.dptr), m, k, original_M, original_K}; - if (use_per_tile_swizzle() && per_tile_swizzle_supported(per_tile_tensor, false)) { - launch_per_tile_swizzle(&per_tile_tensor, 1, false, stream); - } else if (num_tiles_m < TB_DIM && static_cast(narrow_m_slm_size) <= - cuda::max_shared_memory_per_block_optin()) { - // Narrow-M: batch TB_DIM K-tiles per block, fully utilizing all threads. - dim3 num_blocks_narrow(DIVUP(num_tiles_k, TB_DIM)); - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_col_scaling_narrow_m_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, narrow_m_slm_size)); - swizzle_col_scaling_narrow_m_kernel - <<>>( - input->columnwise_scale_inv.dptr, output->columnwise_scale_inv.dptr, m, k, original_M, - original_K); + const int64_t grid_size = + static_cast(m / SF_TILE_DIM_M) * DIVUP(k, kColumnwiseSmallChunkRows); + if (grid_size > kColumnwiseMaxBlocksWith32RowChunks) { + launch_multi_tensor_columnwise_chunk_swizzle(&descriptor, 1, stream); } else { - int vec_load_size = (num_tiles_m - 1) % 4 + 1; - if (vec_load_size == 3) vec_load_size = 1; /* no int3 and misaligned if using int4/int2 */ - int n_tiles_in_tb = TB_DIM * vec_load_size; - dim3 num_blocks(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size)); - int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); - - switch (vec_load_size) { - case 4: - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_col_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - swizzle_col_scaling_kernel - <<>>(input->columnwise_scale_inv.dptr, - output->columnwise_scale_inv.dptr, m, - k, original_M, original_K); - break; - case 2: - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_col_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - swizzle_col_scaling_kernel - <<>>(input->columnwise_scale_inv.dptr, - output->columnwise_scale_inv.dptr, m, - k, original_M, original_K); - break; - case 1: - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_col_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - swizzle_col_scaling_kernel - <<>>(input->columnwise_scale_inv.dptr, - output->columnwise_scale_inv.dptr, m, - k, original_M, original_K); - break; - default: - NVTE_ERROR("Not valid vec_load_size."); - break; - } + launch_columnwise_chunk_swizzle(descriptor, stream); } NVTE_CHECK_CUDA(cudaGetLastError()); } } template -void launch_multi_tensor_swizzle_scaling_factors(MultiSwizzleArgs& kernel_args, - const int vec_load_size, const bool is_rowwise, - const bool use_narrow_k, const bool use_narrow_m, - cudaStream_t stream) { +void launch_multi_tensor_row_scaling_swizzle(MultiSwizzleArgs& kernel_args, const int vec_load_size, + cudaStream_t stream) { // cudaFuncSetAttribute is a host-synchronous driver call; cache the max shared memory // setting per kernel variant so we only pay the cost when slm_size actually increases. auto set_smem_if_needed = [](auto kernel_fn, int slm, int& cached) { @@ -1570,125 +1294,40 @@ void launch_multi_tensor_swizzle_scaling_factors(MultiSwizzleArgs& kernel_args, dim3 block_size(TB_DIM, TB_DIM); - if (is_rowwise && use_narrow_k) { - // Narrow-K path: each block handles TB_DIM M-tiles with full thread utilization. - // slm_size depends on num_tiles_k, which can vary per tensor — use the max. - int max_num_tiles_k = 0; - for (size_t j = 0; j < kernel_args.num_tensors; j++) { - const int num_tiles_m = kernel_args.m_list[j] / SF_TILE_DIM_M; - const int num_tiles_k = kernel_args.k_list[j] / SF_TILE_DIM_K; - max_num_tiles_k = std::max(max_num_tiles_k, num_tiles_k); - kernel_args.block_range[j + 1] = kernel_args.block_range[j] + DIVUP(num_tiles_m, TB_DIM); - } - int slm_size = TB_DIM * max_num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); - const int num_blocks = kernel_args.block_range[kernel_args.num_tensors]; - - static int cached_narrow_k = -1; - set_smem_if_needed( - multi_tensor_swizzle_row_scaling_narrow_k_kernel, slm_size, - cached_narrow_k); - multi_tensor_swizzle_row_scaling_narrow_k_kernel - <<>>(kernel_args); - } else if (!is_rowwise && use_narrow_m) { - // Narrow-M path: each block handles TB_DIM K-tiles with full thread utilization. - // slm_size depends on num_tiles_m, which can vary per tensor — use the max. - int max_num_tiles_m = 0; - for (size_t j = 0; j < kernel_args.num_tensors; j++) { - const int num_tiles_m = kernel_args.m_list[j] / SF_TILE_DIM_M; - const int num_tiles_k = kernel_args.k_list[j] / SF_TILE_DIM_K; - max_num_tiles_m = std::max(max_num_tiles_m, num_tiles_m); - kernel_args.block_range[j + 1] = kernel_args.block_range[j] + DIVUP(num_tiles_k, TB_DIM); - } - int slm_size = TB_DIM * max_num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); - const int num_blocks = kernel_args.block_range[kernel_args.num_tensors]; - - static int cached_narrow_m = -1; - set_smem_if_needed( - multi_tensor_swizzle_col_scaling_narrow_m_kernel, slm_size, - cached_narrow_m); - multi_tensor_swizzle_col_scaling_narrow_m_kernel - <<>>(kernel_args); - } else { - int n_tiles_in_tb = TB_DIM * vec_load_size; - int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); - /* Calculate number of CUDA blocks needed for each tensor. - * We have to do it here because we have to iterate over all tensors in this batch to - * get the minimum vec_load_size. - */ - for (size_t j = 0; j < kernel_args.num_tensors; j++) { - const int m = kernel_args.m_list[j]; - const int k = kernel_args.k_list[j]; - int num_tiles_m = m / SF_TILE_DIM_M; - int num_tiles_k = k / SF_TILE_DIM_K; - if (is_rowwise) { - kernel_args.block_range[j + 1] = - kernel_args.block_range[j] + DIVUP(num_tiles_k, n_tiles_in_tb) * num_tiles_m; - } else { - kernel_args.block_range[j + 1] = - kernel_args.block_range[j] + - DIVUP(num_tiles_k, TB_DIM) * DIVUP(num_tiles_m, vec_load_size); - } - } - const int num_blocks = kernel_args.block_range[kernel_args.num_tensors]; - - static int cached_row_int4 = -1, cached_row_int2 = -1, cached_row_int1 = -1; - static int cached_col_int4 = -1, cached_col_int2 = -1, cached_col_int1 = -1; - - if (is_rowwise) { - switch (vec_load_size) { - case 4: - set_smem_if_needed( - multi_tensor_swizzle_row_scaling_kernel, slm_size, - cached_row_int4); - multi_tensor_swizzle_row_scaling_kernel - <<>>(kernel_args); - break; - case 2: - set_smem_if_needed( - multi_tensor_swizzle_row_scaling_kernel, slm_size, - cached_row_int2); - multi_tensor_swizzle_row_scaling_kernel - <<>>(kernel_args); - break; - case 1: - set_smem_if_needed( - multi_tensor_swizzle_row_scaling_kernel, slm_size, - cached_row_int1); - multi_tensor_swizzle_row_scaling_kernel - <<>>(kernel_args); - break; - default: - NVTE_ERROR("Not valid vec_load_size."); - break; - } - } else { - switch (vec_load_size) { - case 4: - set_smem_if_needed( - multi_tensor_swizzle_col_scaling_kernel, slm_size, - cached_col_int4); - multi_tensor_swizzle_col_scaling_kernel - <<>>(kernel_args); - break; - case 2: - set_smem_if_needed( - multi_tensor_swizzle_col_scaling_kernel, slm_size, - cached_col_int2); - multi_tensor_swizzle_col_scaling_kernel - <<>>(kernel_args); - break; - case 1: - set_smem_if_needed( - multi_tensor_swizzle_col_scaling_kernel, slm_size, - cached_col_int1); - multi_tensor_swizzle_col_scaling_kernel - <<>>(kernel_args); - break; - default: - NVTE_ERROR("Not valid vec_load_size."); - break; - } - } + int n_tiles_in_tb = TB_DIM * vec_load_size; + int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); + for (size_t j = 0; j < kernel_args.num_tensors; j++) { + const int num_tiles_m = kernel_args.m_list[j] / SF_TILE_DIM_M; + const int num_tiles_k = kernel_args.k_list[j] / SF_TILE_DIM_K; + kernel_args.block_range[j + 1] = + kernel_args.block_range[j] + DIVUP(num_tiles_k, n_tiles_in_tb) * num_tiles_m; + } + const int num_blocks = kernel_args.block_range[kernel_args.num_tensors]; + + static int cached_row_int4 = -1, cached_row_int2 = -1, cached_row_int1 = -1; + switch (vec_load_size) { + case 4: + set_smem_if_needed( + multi_tensor_swizzle_row_scaling_kernel, slm_size, + cached_row_int4); + multi_tensor_swizzle_row_scaling_kernel + <<>>(kernel_args); + break; + case 2: + set_smem_if_needed( + multi_tensor_swizzle_row_scaling_kernel, slm_size, + cached_row_int2); + multi_tensor_swizzle_row_scaling_kernel + <<>>(kernel_args); + break; + case 1: + set_smem_if_needed(multi_tensor_swizzle_row_scaling_kernel, + slm_size, cached_row_int1); + multi_tensor_swizzle_row_scaling_kernel + <<>>(kernel_args); + break; + default: + NVTE_ERROR("Not valid vec_load_size."); } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -1778,36 +1417,50 @@ void launch_multi_tensor_unswizzle_scaling_factors(MultiSwizzleArgs& kernel_args namespace { -// MXFP8 multi-tensor swizzle with the per-tile kernels. Returns false without launching anything -// if any tensor is not supported (the caller then uses the generic kernels, which also report -// shape errors). -bool multi_tensor_per_tile_swizzle(const std::vector& input, std::vector& output, - bool rowwise, cudaStream_t stream) { +// Multi-tensor swizzle using rowwise full-width or columnwise chunk kernels with aligned int4 +// (16-byte) loads and stores. +void int4_vectorized_multi_tensor_swizzle(const std::vector& input, + std::vector& output, bool rowwise, + cudaStream_t stream) { const size_t num_tensors = input.size(); - std::vector tensors(num_tensors); + std::vector tensors(num_tensors); for (size_t i = 0; i < num_tensors; i++) { - if (!is_mxfp8_scaling(input[i]->scaling_mode)) return false; - const auto& in_scales = rowwise ? input[i]->scale_inv : input[i]->columnwise_scale_inv; - const auto& out_scales = rowwise ? output[i]->scale_inv : output[i]->columnwise_scale_inv; - if (in_scales.shape.size() != 2) return false; - const int padded_m = static_cast(rowwise ? in_scales.shape[0] : in_scales.shape[1]); - const int padded_k = static_cast(rowwise ? in_scales.shape[1] : in_scales.shape[0]); - if (padded_m <= 0 || padded_k <= 0 || padded_m % 128 != 0 || padded_k % 4 != 0) return false; - if (out_scales.numel() != static_cast(padded_m) * padded_k) return false; + const bool nvfp4 = is_nvfp4_scaling(input[i]->scaling_mode); + const bool use_rowwise_buffer = input[i]->scale_inv.has_data(); + const bool selected_rowwise = rowwise && use_rowwise_buffer; + const auto& in_scales = selected_rowwise ? input[i]->scale_inv : input[i]->columnwise_scale_inv; + const auto& out_scales = + selected_rowwise ? output[i]->scale_inv : output[i]->columnwise_scale_inv; + NVTE_CHECK(in_scales.shape.size() == 2, "Expected two-dimensional scaling factors."); + const int padded_m = + static_cast((rowwise || nvfp4) ? in_scales.shape[0] : in_scales.shape[1]); + const int padded_k = + static_cast((rowwise || nvfp4) ? in_scales.shape[1] : in_scales.shape[0]); + NVTE_CHECK(padded_m > 0 && padded_k > 0 && padded_m % kSwizzleTileDimM == 0 && + padded_k % kScalesPerWord == 0, + "Invalid scaling-factor shape (", padded_m, ", ", padded_k, ")."); + NVTE_CHECK(in_scales.numel() == static_cast(padded_m) * padded_k && + out_scales.numel() == static_cast(padded_m) * padded_k, + "Input and output scaling-factor shapes must match."); const auto [first_dim, last_dim] = input[i]->flat_2d_dims(); + const int block_size = nvfp4 ? NVFP4_BLOCK_SIZE : MXFP8_BLOCK_SIZE; + const int original_m = static_cast(selected_rowwise ? first_dim : last_dim); + const int original_k = static_cast((selected_rowwise ? last_dim : first_dim) / block_size); tensors[i] = {static_cast(in_scales.dptr), static_cast(out_scales.dptr), padded_m, padded_k, - static_cast(rowwise ? first_dim : last_dim), - static_cast((rowwise ? last_dim : first_dim) / MXFP8_BLOCK_SIZE)}; - if (!per_tile_swizzle_supported(tensors[i], rowwise)) return false; + original_m, + original_k}; } for (size_t first = 0; first < num_tensors; first += kMaxTensorsPerLaunch) { const int count = static_cast(std::min(kMaxTensorsPerLaunch, num_tensors - first)); - launch_per_tile_swizzle(tensors.data() + first, count, rowwise, stream); + if (rowwise) { + launch_multi_tensor_rowwise_full_width_swizzle(tensors.data() + first, count, stream); + } else { + launch_multi_tensor_columnwise_chunk_swizzle(tensors.data() + first, count, stream); + } } - return true; } } // namespace @@ -1847,11 +1500,30 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, const bool rowwise_swizzle = all_has_data || all_nvfp4; const bool columnwise_swizzle = all_has_columnwise_data && !all_nvfp4; - if (!all_nvfp4 && use_per_tile_swizzle() && - multi_tensor_per_tile_swizzle(input, output, rowwise_swizzle, stream)) { + int max_rowwise_padded_k = 0; + int64_t total_rowwise_tiles = 0; + for (size_t i = 0; i < num_tensors; i++) { + const bool use_rowwise_buffer = input[i]->scale_inv.has_data(); + const bool selected_rowwise = rowwise_swizzle && use_rowwise_buffer; + const auto& in_scales = selected_rowwise ? input[i]->scale_inv : input[i]->columnwise_scale_inv; + check_scale_inv_alignment(input[i], "Input", selected_rowwise); + check_scale_inv_alignment(output[i], "Output", selected_rowwise); + if (rowwise_swizzle) { + NVTE_CHECK(in_scales.shape.size() == 2, "Expected two-dimensional scaling factors."); + max_rowwise_padded_k = std::max(max_rowwise_padded_k, static_cast(in_scales.shape[1])); + total_rowwise_tiles += static_cast(in_scales.shape[0]) / kSwizzleTileDimM; + } + } + if (rowwise_swizzle) { + if (should_use_rowwise_full_width_swizzle(max_rowwise_padded_k, total_rowwise_tiles)) { + int4_vectorized_multi_tensor_swizzle(input, output, true, stream); + return; + } + } else if (columnwise_swizzle) { + int4_vectorized_multi_tensor_swizzle(input, output, false, stream); return; } - + // Fallback to the generic rowwise swizzle kernel. constexpr int SF_TILE_DIM_M = 128; constexpr int SF_TILE_DIM_K = 4; if (rowwise_swizzle) { @@ -1859,18 +1531,16 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, kernel_args.num_tensors = 0; kernel_args.block_range[0] = 0; int vec_load_size = 4; - bool all_narrow_k = true; for (size_t i = 0; i < num_tensors; i++) { //Launch kernel if argument struct is full if (kernel_args.num_tensors == kMaxTensorsPerKernel) { // There is no int3 and misaligned if using int4/int2. if (vec_load_size == 3) vec_load_size = 1; - launch_multi_tensor_swizzle_scaling_factors( - kernel_args, vec_load_size, true, all_narrow_k, false, stream); + launch_multi_tensor_row_scaling_swizzle( + kernel_args, vec_load_size, stream); // Reset the argument struct and vec_load_size kernel_args.num_tensors = 0; vec_load_size = 4; - all_narrow_k = true; } int m, k; @@ -1904,11 +1574,6 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, } int num_tiles_k = k / SF_TILE_DIM_K; - const int narrow_k_slm = - TB_DIM * num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t)); - all_narrow_k = - all_narrow_k && (num_tiles_k < TB_DIM) && - (static_cast(narrow_k_slm) <= cuda::max_shared_memory_per_block_optin()); int vec_load_size_i = (num_tiles_k - 1) % 4 + 1; // We use the minimum vec_load_size across all tensors. // TODO(zhongbo): fix vec_load_size for NVFP4 @@ -1938,69 +1603,8 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, // Launch the remaining tensors // There is no int3 and misaligned if using int4/int2. if (vec_load_size == 3) vec_load_size = 1; - launch_multi_tensor_swizzle_scaling_factors( - kernel_args, vec_load_size, true, all_narrow_k, false, stream); - } - - if (columnwise_swizzle) { - // NVFP4 shouldn't end up here because it only needs rowwise swizzle - NVTE_CHECK(!all_nvfp4, "NVFP4 shouldn't end up here because it only needs rowwise swizzle"); - - MultiSwizzleArgs kernel_args; - kernel_args.num_tensors = 0; - kernel_args.block_range[0] = 0; - int vec_load_size = 4; - bool all_narrow_m = true; - for (size_t i = 0; i < num_tensors; i++) { - //Launch kernel if argument struct is full - if (kernel_args.num_tensors == kMaxTensorsPerKernel) { - // There is no int3 and misaligned if using int4/int2. - if (vec_load_size == 3) vec_load_size = 1; - launch_multi_tensor_swizzle_scaling_factors( - kernel_args, vec_load_size, false, false, all_narrow_m, stream); - // Reset the argument struct and vec_load_size - kernel_args.num_tensors = 0; - vec_load_size = 4; - all_narrow_m = true; - } - const int m = input[i]->columnwise_scale_inv.shape[1]; - const int k = input[i]->columnwise_scale_inv.shape[0]; - - 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()), - "Input.columnwise_scale_inv size is not equal to " - "Output.columnwise_scale_inv size!"); - - int num_tiles_m = m / SF_TILE_DIM_M; - int num_tiles_k = k / SF_TILE_DIM_K; - const int narrow_m_slm = - TB_DIM * num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t)); - all_narrow_m = - all_narrow_m && (num_tiles_m < TB_DIM) && - (static_cast(narrow_m_slm) <= cuda::max_shared_memory_per_block_optin()); - int vec_load_size_i = (num_tiles_k - 1) % 4 + 1; - // We use the minimum vec_load_size across all tensors. - 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(input[i]->columnwise_scale_inv.dptr); - kernel_args.output_list[pos] = output[i]->columnwise_scale_inv.dptr; - kernel_args.m_list[pos] = m; - kernel_args.k_list[pos] = k; - const auto [first_dim, last_dim] = input[i]->flat_2d_dims(); - kernel_args.original_m_list[pos] = last_dim; - kernel_args.original_k_list[pos] = first_dim / MXFP8_BLOCK_SIZE; - kernel_args.num_tensors++; - } - // Launch the remaining tensors - // There is no int3 and misaligned if using int4/int2. - if (vec_load_size == 3) vec_load_size = 1; - launch_multi_tensor_swizzle_scaling_factors( - kernel_args, vec_load_size, false, false, all_narrow_m, stream); + launch_multi_tensor_row_scaling_swizzle(kernel_args, + vec_load_size, stream); } } @@ -2406,8 +2010,8 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) size_t m = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k); size_t k = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m); - size_t padded_m = round_up_to_multiple(m, 128); - size_t padded_k = round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); + size_t padded_m = DIVUP_TO_MULTIPLE(m, 128); + size_t padded_k = DIVUP_TO_MULTIPLE(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); int num_tiles_m = padded_m / SF_TILE_DIM_M; int num_tiles_k = padded_k / SF_TILE_DIM_K; @@ -2446,8 +2050,8 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) M = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k); K = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m); - size_t padded_m = round_up_to_multiple(M, 128); - size_t padded_k = round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4); + size_t padded_m = DIVUP_TO_MULTIPLE(M, 128); + size_t padded_k = DIVUP_TO_MULTIPLE(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4); int num_tiles_m = padded_m / SF_TILE_DIM_M; int num_tiles_k = padded_k / SF_TILE_DIM_K; @@ -2475,9 +2079,9 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) const uint8_t* input_base = reinterpret_cast(input) + current_scale_base; uint8_t* output_base = reinterpret_cast(output) + current_scale_base; - const int padded_m = static_cast(round_up_to_multiple(M, 128)); + const int padded_m = static_cast(DIVUP_TO_MULTIPLE(M, 128)); const int padded_k = - static_cast(round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4)); + static_cast(DIVUP_TO_MULTIPLE(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4)); const int original_M = static_cast(M); const int original_K = static_cast(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE))); const bool padding_m = (block_y == grid_dim_y - 1) && (original_M < padded_m); @@ -2577,9 +2181,8 @@ void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* auto launch_grouped_swizzle = [&](bool rowwise) { const size_t m = rowwise ? first_dim : last_dim; const size_t k = rowwise ? last_dim : first_dim; - const size_t padded_m = round_up_to_multiple(m, 128); - const size_t padded_k = - round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); + const size_t padded_m = DIVUP_TO_MULTIPLE(m, 128); + const size_t padded_k = DIVUP_TO_MULTIPLE(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); // Per-tensor scale-element counts: // - "padded" layout: each tensor occupies padded_m * padded_k elements // (total buffer = num_tensors * padded_m * padded_k). @@ -2599,8 +2202,8 @@ void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* const size_t compact_scale_elems = rowwise ? m * padded_k : DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)) * padded_m; const size_t compact_total_scale_elems = - rowwise ? round_up_to_multiple(input->num_tensors * m, 128) * padded_k - : round_up_to_multiple( + rowwise ? DIVUP_TO_MULTIPLE(input->num_tensors * m, 128) * padded_k + : DIVUP_TO_MULTIPLE( input->num_tensors * DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4) * padded_m; @@ -2765,9 +2368,8 @@ void unswizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor auto launch_grouped_unswizzle = [&](bool rowwise) { const size_t m = rowwise ? first_dim : last_dim; const size_t k = rowwise ? last_dim : first_dim; - const size_t padded_m = round_up_to_multiple(m, 128); - const size_t padded_k = - round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); + const size_t padded_m = DIVUP_TO_MULTIPLE(m, 128); + const size_t padded_k = DIVUP_TO_MULTIPLE(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); const size_t scale_elems = padded_m * padded_k; const size_t scale_elem_size = rowwise ? typeToSize(output->scale_inv.dtype)