IFU 2.17 upstream feature enablement

tests/cpp/operator/test_swizzle.cu

@@ -6,12 +6,14 @@
  * See LICENSE for license information.
  ************************************************************************/
 
+#include <algorithm>
 #include <cmath>
 #include <cstdint>
 #include <cstring>
 #include <memory>
 #include <iomanip>
 #include <iostream>
+#include <numeric>
 #include <random>
 #include <type_traits>
 
@@ -653,6 +655,263 @@ INSTANTIATE_TEST_SUITE_P(
   }
 );
 
+#endif  // !__HIP_PLATFORM_AMD__ (grouped unswizzle / roundtrip / variable suites)
+
+#ifdef __HIP_PLATFORM_AMD__
+// One tensor's scale block in the padded layout the grouped swizzle contracts
+// on. Columnwise swaps which logical dim is tiled by 128, and stores transposed.
+struct VariableScaleBlock {
+  size_t valid_128;
+  size_t valid_4;
+  size_t padded_128;
+  size_t padded_4;
+
+  size_t numel() const { return padded_128 * padded_4; }
+  // Offset of a valid element within the block, in its own storage order.
+  size_t offset(size_t i128, size_t i4, bool rowwise) const {
+    return rowwise ? i128 * padded_4 + i4 : i4 * padded_128 + i128;
+  }
+};
+
+static VariableScaleBlock variable_scale_block(size_t M, size_t K, bool rowwise) {
+  constexpr size_t BLOCK_SIZE = 32;
+  VariableScaleBlock block;
+  block.valid_128 = rowwise ? M : K;
+  block.valid_4 =
+      rowwise ? test::divide_round_up(K, BLOCK_SIZE) : test::divide_round_up(M, BLOCK_SIZE);
+  block.padded_128 = test::round_up_to_nearest_multiple(block.valid_128, 128);
+  block.padded_4 = test::round_up_to_nearest_multiple(block.valid_4, 4);
+  return block;
+}
+
+// Built here rather than gathered from test::Tensor, whose MXFP8 scales are
+// unpadded on ROCm and would carry a stride this entry point does not expect.
+void performTestGroupedSwizzleMXFP8VariableRocm(const std::vector<std::pair<size_t, size_t>>& shapes) {
+  using namespace transformer_engine;
+  using namespace test;
+
+  const size_t num_tensors = shapes.size();
+  std::vector<int64_t> first_dims(num_tensors), last_dims(num_tensors);
+  for (size_t i = 0; i < num_tensors; ++i) {
+    first_dims[i] = static_cast<int64_t>(shapes[i].first);
+    last_dims[i] = static_cast<int64_t>(shapes[i].second);
+  }
+  const bool same_first = std::all_of(first_dims.begin(), first_dims.end(),
+                                      [&](int64_t v) { return v == first_dims[0]; });
+  const bool same_last = std::all_of(last_dims.begin(), last_dims.end(),
+                                     [&](int64_t v) { return v == last_dims[0]; });
+
+  std::mt19937 gen(1234);
+  std::uniform_int_distribution<int> byte_dist(0, 255);
+
+  // Input padding stays zero so the reference, which permutes whatever it is
+  // given, agrees with the kernel, which zeroes past the valid extent.
+  auto build_side = [&](bool rowwise, std::vector<uint8_t>& input_host,
+                        std::vector<uint8_t>& ref_host, std::vector<size_t>& offsets) {
+    size_t total = 0;
+    offsets.resize(num_tensors);
+    for (size_t i = 0; i < num_tensors; ++i) {
+      offsets[i] = total;
+      total += variable_scale_block(shapes[i].first, shapes[i].second, rowwise).numel();
+    }
+    input_host.assign(total, 0);
+    ref_host.assign(total, 0);
+    for (size_t i = 0; i < num_tensors; ++i) {
+      const auto block = variable_scale_block(shapes[i].first, shapes[i].second, rowwise);
+      uint8_t* in = input_host.data() + offsets[i];
+      for (size_t a = 0; a < block.valid_128; ++a) {
+        for (size_t b = 0; b < block.valid_4; ++b) {
+          in[block.offset(a, b, rowwise)] = static_cast<uint8_t>(byte_dist(gen));
+        }
+      }
+      if (rowwise) {
+        compute_ref_swizzle<128, 4, true>(in, ref_host.data() + offsets[i], block.padded_128,
+                                          block.padded_4);
+      } else {
+        compute_ref_swizzle<128, 4, false>(in, ref_host.data() + offsets[i], block.padded_128,
+                                           block.padded_4);
+      }
+    }
+    return total;
+  };
+
+  std::vector<uint8_t> row_input, row_ref, col_input, col_ref;
+  std::vector<size_t> row_offsets, col_offsets;
+  const size_t row_total = build_side(true, row_input, row_ref, row_offsets);
+  const size_t col_total = build_side(false, col_input, col_ref, col_offsets);
+
+  auto upload = [](const std::vector<uint8_t>& host) {
+    CudaPtr<> dev = cuda_alloc(std::max<size_t>(host.size(), 1));
+    if (!host.empty()) {
+      NVTE_CHECK_CUDA(
+          cudaMemcpy(dev.get(), host.data(), host.size(), cudaMemcpyHostToDevice));
+    }
+    return dev;
+  };
+  CudaPtr<> row_in_dev = upload(row_input);
+  CudaPtr<> col_in_dev = upload(col_input);
+  // 0xCD rather than 0: a position the kernel fails to write must not pass by
+  // coincidentally matching the zeroed padding the reference expects.
+  CudaPtr<> row_out_dev = cuda_alloc(std::max<size_t>(row_total, 1));
+  CudaPtr<> col_out_dev = cuda_alloc(std::max<size_t>(col_total, 1));
+  NVTE_CHECK_CUDA(cudaMemset(row_out_dev.get(), 0xCD, row_total));
+  NVTE_CHECK_CUDA(cudaMemset(col_out_dev.get(), 0xCD, col_total));
+
+  CudaPtr<int64_t> first_dims_dev = cuda_alloc<int64_t>(num_tensors * sizeof(int64_t));
+  CudaPtr<int64_t> last_dims_dev = cuda_alloc<int64_t>(num_tensors * sizeof(int64_t));
+  NVTE_CHECK_CUDA(cudaMemcpy(first_dims_dev.get(), first_dims.data(),
+                             num_tensors * sizeof(int64_t), cudaMemcpyHostToDevice));
+  NVTE_CHECK_CUDA(cudaMemcpy(last_dims_dev.get(), last_dims.data(),
+                             num_tensors * sizeof(int64_t), cudaMemcpyHostToDevice));
+
+  size_t logical_data[2] = {static_cast<size_t>(first_dims[0]), static_cast<size_t>(last_dims[0])};
+  if (same_first && same_last) {
+    logical_data[0] = static_cast<size_t>(first_dims[0]) * num_tensors;
+  } else if (same_first) {
+    logical_data[1] = static_cast<size_t>(
+        std::accumulate(last_dims.begin(), last_dims.end(), int64_t{0}));
+  } else if (same_last) {
+    logical_data[0] = static_cast<size_t>(
+        std::accumulate(first_dims.begin(), first_dims.end(), int64_t{0}));
+  } else {
+    logical_data[0] = 1;
+    logical_data[1] = 0;
+    for (size_t i = 0; i < num_tensors; ++i) {
+      logical_data[1] += static_cast<size_t>(first_dims[i] * last_dims[i]);
+    }
+  }
+  const NVTEShape logical_shape = nvte_make_shape(logical_data, 2);
+
+  // Only scale_inv is read, but a grouped tensor is not allocated without data,
+  // and its CSR offsets are mandatory once any dimension varies.
+  size_t data_elems = 0;
+  std::vector<int64_t> data_offsets(num_tensors + 1, 0);
+  for (size_t i = 0; i < num_tensors; ++i) {
+    data_elems += static_cast<size_t>(first_dims[i] * last_dims[i]);
+    data_offsets[i + 1] = static_cast<int64_t>(data_elems);
+  }
+  CudaPtr<> data_dev = cuda_alloc(std::max<size_t>(data_elems, 1));
+  const size_t num_offsets = num_tensors + 1;
+  CudaPtr<int64_t> offsets_dev = cuda_alloc<int64_t>(num_offsets * sizeof(int64_t));
+  NVTE_CHECK_CUDA(cudaMemcpy(offsets_dev.get(), data_offsets.data(),
+                             num_offsets * sizeof(int64_t), cudaMemcpyHostToDevice));
+
+  auto make_grouped = [&](void* row_scales, size_t row_numel, void* col_scales, size_t col_numel,
+                          uint8_t swizzled) {
+    GroupedTensorHandle handle(
+        nvte_create_grouped_tensor(NVTE_MXFP8_1D_SCALING, num_tensors, logical_shape));
+    NVTEGroupedTensor h = handle.get();
+    NVTEShape dims_shape = nvte_make_shape(&num_tensors, 1);
+    NVTEShape data_shape = nvte_make_shape(&data_elems, 1);
+    NVTEBasicTensor data_t{data_dev.get(), kNVTEFloat8E4M3, data_shape};
+    nvte_set_grouped_tensor_param(h, kNVTEGroupedRowwiseData, &data_t, sizeof(data_t));
+    nvte_set_grouped_tensor_param(h, kNVTEGroupedColumnwiseData, &data_t, sizeof(data_t));
+    if (!same_first) {
+      NVTEBasicTensor t{first_dims_dev.get(), kNVTEInt64, dims_shape};
+      nvte_set_grouped_tensor_param(h, kNVTEGroupedFirstDims, &t, sizeof(t));
+    }
+    if (!same_last) {
+      NVTEBasicTensor t{last_dims_dev.get(), kNVTEInt64, dims_shape};
+      nvte_set_grouped_tensor_param(h, kNVTEGroupedLastDims, &t, sizeof(t));
+    }
+    if (!same_first || !same_last) {
+      NVTEShape off_shape = nvte_make_shape(&num_offsets, 1);
+      NVTEBasicTensor t{offsets_dev.get(), kNVTEInt64, off_shape};
+      nvte_set_grouped_tensor_param(h, kNVTEGroupedTensorOffsets, &t, sizeof(t));
+    }
+    NVTEShape row_shape = nvte_make_shape(&row_numel, 1);
+    NVTEBasicTensor row_t{row_scales, kNVTEFloat8E8M0, row_shape};
+    nvte_set_grouped_tensor_param(h, kNVTEGroupedRowwiseScaleInv, &row_t, sizeof(row_t));
+    NVTEShape col_shape = nvte_make_shape(&col_numel, 1);
+    NVTEBasicTensor col_t{col_scales, kNVTEFloat8E8M0, col_shape};
+    nvte_set_grouped_tensor_param(h, kNVTEGroupedColumnwiseScaleInv, &col_t, sizeof(col_t));
+    nvte_set_grouped_tensor_param(h, kNVTEGroupedWithGEMMSwizzledScales, &swizzled,
+                                  sizeof(swizzled));
+    return handle;
+  };
+
+  GroupedTensorHandle input =
+      make_grouped(row_in_dev.get(), row_total, col_in_dev.get(), col_total, 0);
+  GroupedTensorHandle output =
+      make_grouped(row_out_dev.get(), row_total, col_out_dev.get(), col_total, 1);
+
+  nvte_swizzle_grouped_scaling_factors(input.get(), output.get(), 0);
+  NVTE_CHECK_CUDA(cudaDeviceSynchronize());
+  NVTE_CHECK_CUDA(cudaGetLastError());
+
+  std::vector<uint8_t> row_out(row_total), col_out(col_total);
+  if (row_total > 0) {
+    NVTE_CHECK_CUDA(
+        cudaMemcpy(row_out.data(), row_out_dev.get(), row_total, cudaMemcpyDeviceToHost));
+  }
+  if (col_total > 0) {
+    NVTE_CHECK_CUDA(
+        cudaMemcpy(col_out.data(), col_out_dev.get(), col_total, cudaMemcpyDeviceToHost));
+  }
+
+  for (size_t i = 0; i < num_tensors; ++i) {
+    const auto row_block = variable_scale_block(shapes[i].first, shapes[i].second, true);
+    const auto col_block = variable_scale_block(shapes[i].first, shapes[i].second, false);
+    compareResults("grouped_swizzle_variable_rowwise_" + std::to_string(i),
+                   row_out.data() + row_offsets[i], row_ref.data() + row_offsets[i],
+                   row_block.numel());
+    compareResults("grouped_swizzle_variable_colwise_" + std::to_string(i),
+                   col_out.data() + col_offsets[i], col_ref.data() + col_offsets[i],
+                   col_block.numel());
+  }
+}
+
+class SwizzleGroupedVariableRocmTestSuite
+    : public ::testing::TestWithParam<std::vector<std::pair<size_t, size_t>>> {};
+
+TEST_P(SwizzleGroupedVariableRocmTestSuite, TestGroupedSwizzleMXFP8Variable) {
+  const auto shapes = GetParam();
+  performTestGroupedSwizzleMXFP8VariableRocm(shapes);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+  OperatorTest,
+  SwizzleGroupedVariableRocmTestSuite,
+  ::testing::Values(
+    // Case 1: num_tensors = 1 (n+3 = 4, even). Check simple alignment.
+    std::vector<std::pair<size_t, size_t>>{{1024, 1024}},
+
+    // Case 2: num_tensors = 2 (n+3 = 5, odd). Forces padding logic to trigger.
+    std::vector<std::pair<size_t, size_t>>{{128, 128}, {256, 256}},
+
+    // Case 3: Mixed small/irregular shapes.
+    std::vector<std::pair<size_t, size_t>>{{200, 160}, {33, 64}, {1, 32}},
+
+    // Case 4: Large workload to verify persistent grid
+    std::vector<std::pair<size_t, size_t>>(10, {4096, 4096}),
+
+    // Case 5: Variable M, Uniform K (Semi-variable)
+    std::vector<std::pair<size_t, size_t>>{{128, 256}, {512, 256}, {64, 256}},
+
+    // Case 6: Uniform M, Variable K (Semi-variable)
+    std::vector<std::pair<size_t, size_t>>{{512, 128}, {512, 1024}, {512, 32}},
+
+    // Case 7: Both dims varying, spanning padded scale-K of 4, 8 and 16 so that
+    // all three vectorized load widths are selected within one launch.
+    std::vector<std::pair<size_t, size_t>>{{256, 128}, {130, 256}, {512, 512}, {64, 96}},
+
+    // Case 8-10: a zero-token expert leading, middle and trailing; it must
+    // contribute no work without disturbing its neighbours' offsets.
+    std::vector<std::pair<size_t, size_t>>{{0, 256}, {128, 256}, {256, 256}},
+    std::vector<std::pair<size_t, size_t>>{{128, 256}, {0, 256}, {256, 256}},
+    std::vector<std::pair<size_t, size_t>>{{128, 256}, {256, 256}, {0, 256}},
+
+    // Case 11: a zero last dim, which empties the columnwise block instead.
+    std::vector<std::pair<size_t, size_t>>{{128, 0}, {128, 256}}
+  ),
+  [](const testing::TestParamInfo<SwizzleGroupedVariableRocmTestSuite::ParamType>& info) {
+    return "VariableShapes_" + std::to_string(info.index) + "_N" + std::to_string(info.param.size());
+  }
+);
+
+#endif  // __HIP_PLATFORM_AMD__ (ROCm variable-shape grouped swizzle)
+
+#ifndef __HIP_PLATFORM_AMD__
 class SwizzleGroupedTestSuite
     : public ::testing::TestWithParam<std::tuple<int, size_t, size_t>> {};
 

transformer_engine/common/include/transformer_engine/swizzle.h

@@ -1,4 +1,6 @@
 /*************************************************************************
+ * This file was modified for portability to AMDGPU
+ * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
  * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  *
  * See LICENSE for license information.
@@ -113,7 +115,9 @@ void nvte_swizzle_block_scaling_to_mxfp8_scaling_factors(const NVTETensor input,
  *  - scale_inv is stored in row-major per group.
  *  - scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale.
  *  - data is quantitized along K-dimension, i.e. 1D-scaling block lies along the K-dimension.
- *  - all tensors in the grouped tensor must have the same shape.
+ *  - tensors may differ in shape, in which case first_dims and/or last_dims must be set and both
+ *    sides use the per-tensor padded layout. A tensor with a zero extent contributes nothing.
+ *    The compact input layout is only accepted when all tensors have the same shape.
  */
 void nvte_swizzle_grouped_scaling_factors(const NVTEGroupedTensor input, NVTEGroupedTensor output,
                                           cudaStream_t stream);