diff --git a/docs/envvars.rst b/docs/envvars.rst index 09e2b46278..5ff430eabe 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 4bd1800cea..69ef1cef99 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 // =================================================================== @@ -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 ? (K + BLOCK_SIZE - 1) / BLOCK_SIZE : (M + BLOCK_SIZE - 1) / 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 ? (K + BLOCK_SIZE - 1) / BLOCK_SIZE : (M + BLOCK_SIZE - 1) / 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) // =================================================================== @@ -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 a442215ae1..7025a4b055 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,156 @@ 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. +} + +void performTestRejectUnalignedMXFP8Scales(bool rowwise, bool offset_input) { + using namespace test; + constexpr size_t kDataSize = 128 * 128; + constexpr size_t kScaleSize = 128 * 4; + auto input_data = cuda_alloc(kDataSize); + auto output_data = cuda_alloc(kDataSize); + auto input_storage = cuda_alloc(kScaleSize + 4); + auto output_storage = cuda_alloc(kScaleSize + 4); + uint8_t* input_scales = input_storage.get() + (offset_input ? 4 : 0); + uint8_t* output_scales = output_storage.get() + (offset_input ? 0 : 4); + TensorWrapper input(NVTE_MXFP8_1D_SCALING); + TensorWrapper output(NVTE_MXFP8_1D_SCALING); + const std::vector data_shape{128, 128}; + const std::vector scale_shape = rowwise ? std::vector{128, 4} + : std::vector{4, 128}; + 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 kScaleSize = 128 * 4; + const std::vector data_shape{128, 128}; + const std::vector scale_shape{128, 4}; + 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(2); + output_tensors.reserve(2); + for (int i = 0; i < 2; i++) { + input_data.push_back(cuda_alloc(128 * 128)); + output_data.push_back(cuda_alloc(128 * 128)); + input_storage.push_back(cuda_alloc(kScaleSize + 4)); + output_storage.push_back(cuda_alloc(kScaleSize + 4)); + uint8_t* input_scales = input_storage.back().get() + (i == 0 && offset_input ? 4 : 0); + uint8_t* output_scales = output_storage.back().get() + (i == 0 && !offset_input ? 4 : 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 +447,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 +459,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 +1137,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 +1182,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 +1190,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 f85d88fdce..1c54e5f898 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 19fb8e848a..af8fe32fe1 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -223,17 +223,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 +383,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, @@ -867,102 +699,25 @@ __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. +// 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. 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 kRowwiseMaxSharedMemoryBytes = 32 * 1024; +// Wide rows use the full-width kernel for tiny launches or once the aggregate grid is large enough +// to occupy the GPU. Narrow rows use it regardless of grid size. +constexpr int kRowwiseNarrowSharedMemoryBytes = 17 * 1024; +constexpr int kRowwiseWideSharedMemoryBytes = 30 * 1024; +constexpr int kRowwiseMaxTinyTiles = 16; +constexpr int kRowwiseMinSaturatingTiles = 128; +// Row-wise: the largest grid before several row 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"); } - // 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, @@ -983,10 +738,10 @@ inline uint32_t division_magic(int divisor) { return divisor <= 1 ? 0u : static_cast(((1ull << 32) + divisor - 1) / 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, @@ -1084,18 +839,18 @@ __device__ __forceinline__ void swizzle_columnwise_chunk(const uint8_t* __restri } __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) { + 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[32 * 8]; swizzle_columnwise_chunk<32>(input, output, padded_m, padded_k, original_m, original_k, blockIdx.x, blockIdx.y * 32, columnwise_smem); @@ -1104,33 +859,33 @@ __global__ void __launch_bounds__(kColumnwiseThreadsPerBlock) // 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 / 128 - 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,7 +893,7 @@ __global__ void __launch_bounds__(kRowwiseThreadsPerBlock) template __global__ void __launch_bounds__(kColumnwiseThreadsPerBlock) - multi_tensor_swizzle_columnwise_per_tile_kernel(const PerTileMultiSwizzleArgs args) { + multi_tensor_swizzle_columnwise_chunk_kernel(const MultiTensorSwizzleArgs args) { __shared__ int4 columnwise_multi_smem[kChunkRows * 8]; const int t = tensor_of_block(args, blockIdx.x); const int block_in_tensor = blockIdx.x - args.first_block[t]; @@ -1155,116 +910,114 @@ 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 = 128 * ((padded_k >> 2) | 1) * 4; + if (smem_bytes <= kRowwiseNarrowSharedMemoryBytes) return true; + return smem_bytes <= kRowwiseWideSharedMemoryBytes && + (total_row_tiles <= kRowwiseMaxTinyTiles || total_row_tiles >= kRowwiseMinSaturatingTiles); +} + +// 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); - } + total_row_tiles += tensors[i].padded_m / 128; + max_smem_row_words = max(max_smem_row_words, (tensors[i].padded_k >> 2) | 1); } - 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; + int row_tiles_per_block = 1; + const int max_row_tiles_per_block = kRowwiseMaxSharedMemoryBytes / (128 * max_smem_row_words * 4); + 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.tiles_per_block = tiles_per_block; + args.row_tiles_per_block = row_tiles_per_block; 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); + const SwizzleTensorDescriptor& tensor = tensors[i]; + const int blocks = (tensor.padded_m / 128 + row_tiles_per_block - 1) / row_tiles_per_block; 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; + multi_tensor_swizzle_rowwise_full_width_kernel<<< + grid, kRowwiseThreadsPerBlock, row_tiles_per_block * 128 * max_smem_row_words * 4, 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++) { + args.tensors[i] = tensors[i]; + total_chunks += + static_cast(tensors[i].padded_m / 128) * ((tensors[i].padded_k + 31) / 32); + } + const int chunk_rows = total_chunks > kColumnwiseMaxBlocksWith32RowChunks ? 128 : 32; + for (int i = 0; i < num_tensors; i++) { + const SwizzleTensorDescriptor& tensor = tensors[i]; + const int blocks = (tensor.padded_m / 128) * ((tensor.padded_k + chunk_rows - 1) / 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; + if (chunk_rows == 128) { + multi_tensor_swizzle_columnwise_chunk_kernel<128> <<>>(args); } else { - multi_tensor_swizzle_columnwise_per_tile_kernel<32> + multi_tensor_swizzle_columnwise_chunk_kernel<32> <<>>(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>>>( +// 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 / 128; + const int smem = 128 * ((tensor.padded_k >> 2) | 1) * 4; + swizzle_rowwise_full_width_kernel<<>>( tensor.input, tensor.output, tensor.padded_k, division_magic(tensor.padded_k >> 2), 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; +// 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 / 128, (tensor.padded_k + 31) / 32); - swizzle_columnwise_per_tile_kernel<<>>( + 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 +1028,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 +1093,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 +1120,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 +1176,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 +1236,26 @@ 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) * ((k + 31) / 32); + 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 +1268,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 +1391,49 @@ 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 % 128 == 0 && padded_k % 4 == 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 +1473,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]) / 128; + } + } + 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 +1504,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 +1547,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 +1576,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); } }