diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index ebb0ab275d..787a83f8b5 100644 --- a/tests/cpp/operator/test_swizzle.cu +++ b/tests/cpp/operator/test_swizzle.cu @@ -6,12 +6,14 @@ * See LICENSE for license information. ************************************************************************/ +#include #include #include #include #include #include #include +#include #include #include @@ -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>& shapes) { + using namespace transformer_engine; + using namespace test; + + const size_t num_tensors = shapes.size(); + std::vector first_dims(num_tensors), last_dims(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + first_dims[i] = static_cast(shapes[i].first); + last_dims[i] = static_cast(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 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& input_host, + std::vector& ref_host, std::vector& 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(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 row_input, row_ref, col_input, col_ref; + std::vector 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& host) { + CudaPtr<> dev = cuda_alloc(std::max(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(row_total, 1)); + CudaPtr<> col_out_dev = cuda_alloc(std::max(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 first_dims_dev = cuda_alloc(num_tensors * sizeof(int64_t)); + CudaPtr last_dims_dev = cuda_alloc(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(first_dims[0]), static_cast(last_dims[0])}; + if (same_first && same_last) { + logical_data[0] = static_cast(first_dims[0]) * num_tensors; + } else if (same_first) { + logical_data[1] = static_cast( + std::accumulate(last_dims.begin(), last_dims.end(), int64_t{0})); + } else if (same_last) { + logical_data[0] = static_cast( + 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(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 data_offsets(num_tensors + 1, 0); + for (size_t i = 0; i < num_tensors; ++i) { + data_elems += static_cast(first_dims[i] * last_dims[i]); + data_offsets[i + 1] = static_cast(data_elems); + } + CudaPtr<> data_dev = cuda_alloc(std::max(data_elems, 1)); + const size_t num_offsets = num_tensors + 1; + CudaPtr offsets_dev = cuda_alloc(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 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>> {}; + +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>{{1024, 1024}}, + + // Case 2: num_tensors = 2 (n+3 = 5, odd). Forces padding logic to trigger. + std::vector>{{128, 128}, {256, 256}}, + + // Case 3: Mixed small/irregular shapes. + std::vector>{{200, 160}, {33, 64}, {1, 32}}, + + // Case 4: Large workload to verify persistent grid + std::vector>(10, {4096, 4096}), + + // Case 5: Variable M, Uniform K (Semi-variable) + std::vector>{{128, 256}, {512, 256}, {64, 256}}, + + // Case 6: Uniform M, Variable K (Semi-variable) + std::vector>{{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>{{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>{{0, 256}, {128, 256}, {256, 256}}, + std::vector>{{128, 256}, {0, 256}, {256, 256}}, + std::vector>{{128, 256}, {256, 256}, {0, 256}}, + + // Case 11: a zero last dim, which empties the columnwise block instead. + std::vector>{{128, 0}, {128, 256}} + ), + [](const testing::TestParamInfo& 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> {}; diff --git a/transformer_engine/common/include/transformer_engine/swizzle.h b/transformer_engine/common/include/transformer_engine/swizzle.h index 396093b543..b09cf55202 100644 --- a/transformer_engine/common/include/transformer_engine/swizzle.h +++ b/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); diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 4fd18cd91b..91d1c38246 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -2212,7 +2212,7 @@ void nvte_multi_tensor_unswizzle_scaling_factors(const NVTETensor* inputs, NVTET namespace transformer_engine { -#ifndef __HIP_PLATFORM_AMD__ // Disabled on ROCm +#ifndef __HIP_PLATFORM_AMD__ template __global__ void __launch_bounds__(TB_DIM* TB_DIM) grouped_swizzle_scaling_variable_shape_kernel(const void* input, void* output, @@ -2360,7 +2360,186 @@ int grouped_swizzle_variable_max_active_blocks_per_sm(int device_id) { std::call_once(flags[device_id], init); return cache[device_id]; } -#endif + +#else // __HIP_PLATFORM_AMD__ + +// Staging area for the widest vectorized load, chosen per tensor at runtime. +// Exactly the gfx942 LDS ceiling, so the block count shares its first word +// rather than being reserved alongside it. +constexpr int grouped_variable_shape_smem_size(int sf_tile_dim_m, int sf_tile_dim_k) { + return TB_DIM * 4 * sf_tile_dim_m * sf_tile_dim_k * static_cast(sizeof(int8_t)); +} + +// One tensor's scale-block extents; columnwise transposes the logical dims. +struct ScaleBlockDims { + size_t m; + size_t k; +}; + +__device__ __forceinline__ ScaleBlockDims variable_shape_dims(int i, const int64_t* m_array, + const int64_t* k_array, bool rowwise, + size_t common_m, size_t common_k) { + const size_t first_dim = m_array ? static_cast(m_array[i]) : common_m; + const size_t last_dim = k_array ? static_cast(k_array[i]) : common_k; + return rowwise ? ScaleBlockDims{first_dim, last_dim} : ScaleBlockDims{last_dim, first_dim}; +} + +// Per-tensor launch geometry. Walked twice, to measure and to resolve; the two +// passes must agree, so the arithmetic has a single definition. +struct VariableShapeTiling { + int grid_dim_x; + int grid_dim_y; + int vec_load_size; + size_t padded_m; + size_t padded_k; + + __device__ __forceinline__ int num_blocks() const { return grid_dim_x * grid_dim_y; } +}; + +template +__device__ __forceinline__ VariableShapeTiling variable_shape_tiling(ScaleBlockDims dims, + bool rowwise) { + VariableShapeTiling tiling; + tiling.padded_m = round_up_to_multiple(dims.m, 128); + tiling.padded_k = round_up_to_multiple(DIVUP(dims.k, static_cast(MXFP8_BLOCK_SIZE)), 4); + + const int num_tiles_m = static_cast(tiling.padded_m) / SF_TILE_DIM_M; + const int num_tiles_k = static_cast(tiling.padded_k) / SF_TILE_DIM_K; + + // A zero-token expert contributes no work, and returning keeps the zero out of + // the modulo below, which would leave vec_load_size at 0 and then divide by it. + if (num_tiles_m == 0 || num_tiles_k == 0) { + tiling.grid_dim_x = 0; + tiling.grid_dim_y = 0; + tiling.vec_load_size = 1; + return tiling; + } + + tiling.vec_load_size = rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1); + if (tiling.vec_load_size == 3) tiling.vec_load_size = 1; + + tiling.grid_dim_x = rowwise ? DIVUP(num_tiles_k, TB_DIM * tiling.vec_load_size) + : DIVUP(num_tiles_k, TB_DIM); + tiling.grid_dim_y = rowwise ? num_tiles_m : DIVUP(num_tiles_m, tiling.vec_load_size); + return tiling; +} + +template +__global__ void __launch_bounds__(TB_DIM* TB_DIM) + grouped_swizzle_scaling_variable_shape_kernel(const void* input, void* output, + const int64_t* m_array, const int64_t* k_array, + int num_tensors, bool rowwise, + size_t scale_elem_size, size_t common_m, + size_t common_k) { + // Shapes are device-side, so the grid cannot be sized on the host and the + // kernel measures its own block space. + extern __shared__ int s_total_blocks[]; + // Serial scan: HIP requires a 64-bit shuffle mask and full-wave participation, + // and num_tensors is the expert count. + if (threadIdx.x == 0 && threadIdx.y == 0) { + int total = 0; + for (int i = 0; i < num_tensors; ++i) { + const ScaleBlockDims dims = + variable_shape_dims(i, m_array, k_array, rowwise, common_m, common_k); + total += variable_shape_tiling(dims, rowwise).num_blocks(); + } + s_total_blocks[0] = total; + } + __syncthreads(); + // Into a register before the loop; the staging area starts at this same word. + const int total_blocks = s_total_blocks[0]; + + for (int linear_block_id = blockIdx.x; linear_block_id < total_blocks; + linear_block_id += gridDim.x) { + // The tile impl has no trailing barrier, so consecutive iterations would + // overlap in the staging area. The loop bound is block-uniform. + __syncthreads(); + + // Empty tensors are skipped naturally: the strict comparison never selects them. + ScaleBlockDims dims{}; + VariableShapeTiling tiling{}; + size_t scale_base_bytes = 0; + int block_base = 0; + for (int i = 0; i < num_tensors; ++i) { + dims = variable_shape_dims(i, m_array, k_array, rowwise, common_m, common_k); + tiling = variable_shape_tiling(dims, rowwise); + const int blocks_i = tiling.num_blocks(); + if (linear_block_id < block_base + blocks_i) break; + block_base += blocks_i; + scale_base_bytes += tiling.padded_m * tiling.padded_k * scale_elem_size; + } + + const int grid_dim_x = tiling.grid_dim_x; + const int grid_dim_y = tiling.grid_dim_y; + const int vec_load_size = tiling.vec_load_size; + const int local_block_id = linear_block_id - block_base; + const int block_x = local_block_id % grid_dim_x; + const int block_y = local_block_id / grid_dim_x; + + const uint8_t* input_base = reinterpret_cast(input) + scale_base_bytes; + uint8_t* output_base = reinterpret_cast(output) + scale_base_bytes; + + const int padded_m = static_cast(tiling.padded_m); + const int padded_k = static_cast(tiling.padded_k); + const int original_M = static_cast(dims.m); + const int original_K = static_cast(DIVUP(dims.k, static_cast(MXFP8_BLOCK_SIZE))); + const bool padding_m = (block_y == grid_dim_y - 1) && (original_M < padded_m); + const bool padding_k = (block_x == grid_dim_x - 1) && (original_K < padded_k); + + if (rowwise) { + if (vec_load_size == 4) { + dispatch_swizzle_row_scaling_kernel_impl( + input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y, + grid_dim_x, grid_dim_y, padding_k, padding_m); + } else if (vec_load_size == 2) { + dispatch_swizzle_row_scaling_kernel_impl( + input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y, + grid_dim_x, grid_dim_y, padding_k, padding_m); + } else { + dispatch_swizzle_row_scaling_kernel_impl( + input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y, + grid_dim_x, grid_dim_y, padding_k, padding_m); + } + } else { + if (vec_load_size == 4) { + dispatch_swizzle_col_scaling_kernel_impl( + input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y, + grid_dim_x, grid_dim_y, padding_k, padding_m); + } else if (vec_load_size == 2) { + dispatch_swizzle_col_scaling_kernel_impl( + input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y, + grid_dim_x, grid_dim_y, padding_k, padding_m); + } else { + dispatch_swizzle_col_scaling_kernel_impl( + input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y, + grid_dim_x, grid_dim_y, padding_k, padding_m); + } + } + } +} + +template +int grouped_swizzle_variable_max_active_blocks_per_sm(int device_id) { + static std::vector cache(cuda::num_devices(), -1); + static std::vector flags(cuda::num_devices()); + NVTE_CHECK(0 <= device_id && device_id < cuda::num_devices(), "invalid CUDA device ID"); + + auto init = [&]() { + constexpr int dynamic_smem_size = + grouped_variable_shape_smem_size(SF_TILE_DIM_M, SF_TILE_DIM_K); + int max_active_blocks_per_sm; + NVTE_CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks_per_sm, + grouped_swizzle_scaling_variable_shape_kernel, + TB_DIM * TB_DIM, dynamic_smem_size)); + NVTE_CHECK(max_active_blocks_per_sm > 0, "Occupancy query returned 0 blocks per SM."); + cache[device_id] = max_active_blocks_per_sm; + }; + std::call_once(flags[device_id], init); + return cache[device_id]; +} + +#endif // __HIP_PLATFORM_AMD__ void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* output, cudaStream_t stream) { @@ -2507,16 +2686,19 @@ void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* launch_grouped_swizzle(false); } } else { -#ifndef __HIP_PLATFORM_AMD__ // Variable shape implementation using Device-Side Block Scheduler size_t num_tensors = input->num_tensors; constexpr int SF_TILE_DIM_M = 128; constexpr int SF_TILE_DIM_K = 4; const dim3 block_size(TB_DIM, TB_DIM); +#ifndef __HIP_PLATFORM_AMD__ const int max_slm_size = TB_DIM * 4 * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); const int metadata_shmem = sizeof(int); // s_total_blocks const int dynamic_smem_size = max_slm_size + metadata_shmem; +#else + const int dynamic_smem_size = grouped_variable_shape_smem_size(SF_TILE_DIM_M, SF_TILE_DIM_K); +#endif size_t common_m = input->all_same_first_dim() ? input->get_common_first_dim() : 0; size_t common_k = input->all_same_last_dim() ? input->get_common_last_dim() : 0; @@ -2553,9 +2735,6 @@ void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* if (has_columnwise_scale_inv) { launch_grouped_swizzle_variable(false); } -#else - NVTE_ERROR("Variable-shape grouped scale swizzling is not supported on ROCm."); -#endif } }