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.

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

Author: vthumbe1503

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)

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<std::unique_ptr<Tensor>> inputs;
+  std::vector<std::unique_ptr<Tensor>> outputs;
+  std::vector<NVTETensor> input_handles;
+  std::vector<NVTETensor> output_handles;
+
+  for (int i = 0; i < num_tensors; ++i) {
+    auto input = std::make_unique<Tensor>("input_" + std::to_string(i),
+                                          std::vector<size_t>{M, K}, DType::kFloat4E2M1, rowwise,
+                                          !rowwise, NVTE_NVFP4_1D_SCALING);
+    auto output = std::make_unique<Tensor>("output_" + std::to_string(i),
+                                           std::vector<size_t>{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<uint8_t>()
+                                    : input->columnwise_cpu_scale_inv_ptr<uint8_t>();
+    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<uint8_t>()
+                                          : inputs[i]->columnwise_cpu_scale_inv_ptr<uint8_t>();
+    const uint8_t *actual = rowwise ? outputs[i]->rowwise_cpu_scale_inv_ptr<uint8_t>()
+                                    : outputs[i]->columnwise_cpu_scale_inv_ptr<uint8_t>();
+    std::vector<uint8_t> 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<uint8_t> 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<size_t> shape{256, 128};
+  std::vector<std::unique_ptr<Tensor>> inputs;
+  std::vector<std::unique_ptr<Tensor>> outputs;
+  std::vector<NVTETensor> input_handles;
+  std::vector<NVTETensor> output_handles;
+
+  for (int i = 0; i < num_tensors; ++i) {
+    const bool has_rowwise_scales = i == 0;
+    auto input = std::make_unique<Tensor>("input_" + std::to_string(i), shape,
+                                          DType::kFloat8E4M3, has_rowwise_scales, true,
+                                          NVTE_MXFP8_1D_SCALING);
+    auto output = std::make_unique<Tensor>("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<uint8_t> reference(scale_shape.data[0] * scale_shape.data[1]);
+    compute_ref_swizzle<128, 4, false>(
+        inputs[i]->columnwise_cpu_scale_inv_ptr<uint8_t>(), 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<uint8_t>(), 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<std::pair<size_t, size_t>> shapes{{128, 7168}, {128, 8064}};
+  std::vector<std::unique_ptr<Tensor>> inputs;
+  std::vector<std::unique_ptr<Tensor>> outputs;
+  std::vector<NVTETensor> input_handles;
+  std::vector<NVTETensor> output_handles;
+  for (size_t i = 0; i < shapes.size(); ++i) {
+    const auto [M, K] = shapes[i];
+    auto input = std::make_unique<Tensor>("input_" + std::to_string(i),
+                                          std::vector<size_t>{M, K}, DType::kFloat8E4M3, true,
+                                          false, NVTE_MXFP8_1D_SCALING);
+    auto output = std::make_unique<Tensor>("output_" + std::to_string(i),
+                                           std::vector<size_t>{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<int>(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<uint8_t> reference(scale_shape.data[0] * scale_shape.data[1]);
+    compute_ref_swizzle<128, 4, true>(inputs[i]->rowwise_cpu_scale_inv_ptr<uint8_t>(),
+                                      reference.data(), scale_shape.data[0], scale_shape.data[1]);
+    compareResults("multi_rowwise_fallback", outputs[i]->rowwise_cpu_scale_inv_ptr<uint8_t>(),
+                   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<std::tuple<std::tuple<int, size_t, size_t, bool>, bool>> {};
+    : public ::testing::TestWithParam<std::tuple<int, size_t, size_t, bool>> {};
 
 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<std::tuple<int, size_t, size_t, bool>> 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<std::tuple<int, size_t, size_t, bool>> 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<std::tuple<int, size_t, size_t, bool>> 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<std::tuple<int, size_t, size_t, bool>> 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<MultiTensorSwizzleTestSuite::ParamType>& 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(