From a1ee272c912a2fb658b72538dca98eb54605f3db Mon Sep 17 00:00:00 2001 From: root Date: Tue, 4 Aug 2026 22:30:41 +0000 Subject: [PATCH 1/3] Enable variable-shape grouped scale swizzling on ROCm The grouped MXFP8 scale swizzle refused any group whose members differ in shape, which is the MoE case where each expert receives a different token count. Two things blocked it, both in the device-side block scheduler; the tile kernels it dispatches to were already portable. The workload measurement used a wave32 butterfly over a 32-bit shuffle mask. HIP requires a 64-bit mask and full-wave participation, and the reduction ran under a half-populated wave. Replace it with a serial scan by one thread: the kernel already performs the same O(num_tensors) scan on every persistent-grid iteration, so measuring it once costs strictly less than work already being done, and it leaves no assumption about how many lanes execute in lockstep. The dynamic shared memory request was the tile staging area plus one int for the block count. The staging area alone is exactly 64 KiB, which is also the per-workgroup LDS ceiling on gfx942-class parts, so the extra word made the launch unsatisfiable there. Alias the count into the first word of the staging buffer instead; every thread copies it into a register before the barrier that hands the buffer over. Two defects in the path being enabled, both reachable on either platform: The tile implementation leaves no trailing barrier, so consecutive iterations of the persistent loop overlapped in the staging area. Add the barrier at the top of the loop body rather than at the end of the implementation, which is shared with the non-persistent kernels where it would be pure cost. An expert routed no tokens yields zero tiles, and the load-width expression then evaluated to zero and divided by it. Return an empty tiling instead. The per-tensor geometry was also duplicated verbatim between the counting and resolution passes; drift between the two copies would have selected the wrong tensor's scales rather than failing. Give it one definition. The test builds its scale blocks directly rather than gathering them from test::Tensor, whose MXFP8 scale allocation is unpadded on ROCm and so carries a per-tensor stride this entry point does not expect. That also makes a zero-token expert expressible, and lets one test body serve both platforms. Co-Authored-By: Claude Opus 5 (1M context) --- tests/cpp/operator/test_swizzle.cu | 308 +++++++++++++----- .../include/transformer_engine/swizzle.h | 6 +- transformer_engine/common/swizzle/swizzle.cu | 200 +++++++----- 3 files changed, 341 insertions(+), 173 deletions(-) diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index ebb0ab275..7ef7576fb 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 @@ -517,104 +519,219 @@ void performTestGroupedSwizzleUnswizzleRoundtrip(const int num_tensors, const si num_tensors * col_numel); } -void performTestGroupedSwizzleMXFP8Variable(const std::vector>& shapes) { - using namespace transformer_engine; - using namespace test; - - int num_tensors = shapes.size(); - std::vector> input_tensors; - std::vector> output_tensors; - std::vector input_ptrs; - std::vector output_ptrs; - input_tensors.reserve(num_tensors); - output_tensors.reserve(num_tensors); - input_ptrs.reserve(num_tensors); - output_ptrs.reserve(num_tensors); +#endif // !__HIP_PLATFORM_AMD__ (grouped unswizzle / roundtrip suites) + +// Geometry of one tensor's MXFP8 scale block in the per-tensor padded layout the +// grouped swizzle contracts on. `dim_128` is the extent the swizzle tiles by 128 +// and `dim_4` the one it tiles by 4; columnwise scales swap which logical +// dimension supplies each, and store 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; - for (int i = 0; i < num_tensors; ++i) { - const std::vector shape{shapes[i].first, shapes[i].second}; - auto input = std::make_unique("input_" + std::to_string(i), shape, - DType::kFloat8E4M3, true, true, - NVTE_MXFP8_1D_SCALING); - auto output = std::make_unique("output_" + std::to_string(i), shape, - DType::kFloat8E4M3, true, true, - NVTE_MXFP8_1D_SCALING); - fillUniform(input.get()); - fillUniform(output.get()); + 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; +} - // Zero padding - input->to_cpu(); - const NVTEShape rs = input->rowwise_scale_inv_shape(); - zero_scale_inv_padding(input->rowwise_cpu_scale_inv_ptr(), - rs.data[0], rs.data[1], - shapes[i].first, (shapes[i].second + BLOCK_SIZE - 1) / BLOCK_SIZE); - const NVTEShape cs = input->columnwise_scale_inv_shape(); - zero_scale_inv_padding(input->columnwise_cpu_scale_inv_ptr(), - cs.data[0], cs.data[1], - (shapes[i].first + BLOCK_SIZE - 1) / BLOCK_SIZE, shapes[i].second); - input->from_cpu(); +// Variable-shape grouped swizzle, against the CPU reference. +// +// The scale buffers are built here rather than gathered from test::Tensor. This +// entry point contracts on the per-tensor *padded* layout, which is not what +// test::Tensor allocates on both platforms -- ROCm quantizers emit unpadded +// scales, so scale_tensor_alignment_* is 1 there (test_common.h) and a gathered +// buffer would carry a per-tensor stride the kernel does not expect. Building +// the layout under test directly also makes a zero-token expert expressible, +// which is the case variable shapes exist for. +void performTestGroupedSwizzleMXFP8Variable(const std::vector>& shapes) { + using namespace transformer_engine; + using namespace test; - input_ptrs.push_back(input.get()); - output_ptrs.push_back(output.get()); - input_tensors.emplace_back(std::move(input)); - output_tensors.emplace_back(std::move(output)); + 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); + + // Host-side input blocks and their expected swizzled form. Padding stays zero + // on the input so the reference, which permutes whatever it is given, agrees + // with the kernel, which zeroes anything 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); - GroupedBuffers grouped_input = build_grouped_tensor(input_ptrs, NVTE_MXFP8_1D_SCALING); - GroupedBuffers grouped_output = build_grouped_tensor(output_ptrs, NVTE_MXFP8_1D_SCALING); - - const uint8_t input_swizzled = 0; - nvte_set_grouped_tensor_param(grouped_input.get_handle(), - kNVTEGroupedWithGEMMSwizzledScales, - &input_swizzled, sizeof(input_swizzled)); - const uint8_t output_swizzled = 1; - nvte_set_grouped_tensor_param(grouped_output.get_handle(), - kNVTEGroupedWithGEMMSwizzledScales, - &output_swizzled, sizeof(output_swizzled)); - - nvte_swizzle_grouped_scaling_factors(grouped_input.get_handle(), - grouped_output.get_handle(), - 0); - - cudaDeviceSynchronize(); + // The swizzle reads only scale_inv, but a grouped tensor is not considered + // allocated without data, so back it with a buffer of the right element count. + // The matching 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()); - // Verification - size_t row_offset = 0; - size_t col_offset = 0; - for (int i = 0; i < num_tensors; ++i) { - const NVTEShape row_shape = input_tensors[i]->rowwise_scale_inv_shape(); - const NVTEShape col_shape = input_tensors[i]->columnwise_scale_inv_shape(); - const size_t row_numel = row_shape.data[0] * row_shape.data[1]; - const size_t col_numel = col_shape.data[0] * col_shape.data[1]; - - std::vector output_row_host(row_numel); - std::vector output_col_host(col_numel); - NVTE_CHECK_CUDA(cudaMemcpy(output_row_host.data(), - static_cast(grouped_output.scale_inv.get()) + row_offset, - row_numel, cudaMemcpyDeviceToHost)); - NVTE_CHECK_CUDA(cudaMemcpy(output_col_host.data(), - static_cast(grouped_output.columnwise_scale_inv.get()) + col_offset, - col_numel, cudaMemcpyDeviceToHost)); - - std::vector ref_row(row_numel); - std::vector ref_col(col_numel); - compute_ref_swizzle<128, 4, true>(input_tensors[i]->rowwise_cpu_scale_inv_ptr(), - ref_row.data(), - row_shape.data[0], row_shape.data[1]); - compute_ref_swizzle<128, 4, false>( - input_tensors[i]->columnwise_cpu_scale_inv_ptr(), - ref_col.data(), - col_shape.data[1], col_shape.data[0]); + 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), - output_row_host.data(), ref_row.data(), row_numel); + row_out.data() + row_offsets[i], row_ref.data() + row_offsets[i], + row_block.numel()); compareResults("grouped_swizzle_variable_colwise_" + std::to_string(i), - output_col_host.data(), ref_col.data(), col_numel); - - row_offset += row_numel; - col_offset += col_numel; + col_out.data() + col_offsets[i], col_ref.data() + col_offsets[i], + col_block.numel()); } } @@ -646,13 +763,28 @@ INSTANTIATE_TEST_SUITE_P( std::vector>{{128, 256}, {512, 256}, {64, 256}}, // Case 6: Uniform M, Variable K (Semi-variable) - std::vector>{{512, 128}, {512, 1024}, {512, 32}} + 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: an expert routed no tokens, in leading, middle and trailing + // position. Its scale block is empty and must contribute no work without + // disturbing the offsets of its neighbours. + 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()); } ); +#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 396093b54..b09cf5520 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 4fd18cd91..acfb649f8 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -2212,7 +2212,73 @@ void nvte_multi_tensor_unswizzle_scaling_factors(const NVTETensor* inputs, NVTET namespace transformer_engine { -#ifndef __HIP_PLATFORM_AMD__ // Disabled on ROCm +// Dynamic shared memory the variable-shape kernel reserves. The vectorized load +// width is a per-tensor runtime value here, so every block must be able to stage +// the widest specialization. This comes to exactly 64 KiB, which is also the +// per-workgroup LDS ceiling on gfx942-class parts, so the measured block count +// shares the first word of this buffer 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)); +} + +// Extents of one tensor's scale block, in rows and columns of its row-major +// storage. Columnwise scales transpose the tensor's logical dimensions. +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, mirroring what the uniform-shape path computes on +// the host. The kernel walks this twice -- once to measure the workload, once to +// map a linear block id back onto a tensor -- and 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; + + // An expert that was routed no tokens has an empty scale block and contributes + // no work. Returning here also keeps the zero out of the modulo below, which + // would otherwise 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, @@ -2220,90 +2286,64 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) int num_tensors, bool rowwise, size_t scale_elem_size, size_t common_m, size_t common_k) { - extern __shared__ int s_metadata[]; - int* s_total_blocks = &s_metadata[0]; - - // Warp reduction to compute total workload - if (threadIdx.x < 32 && threadIdx.y == 0) { - int local_blocks = 0; - for (int i = threadIdx.x; i < num_tensors; i += 32) { - size_t m = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k); - size_t k = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m); - - size_t padded_m = round_up_to_multiple(m, 128); - size_t padded_k = round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); - - int num_tiles_m = padded_m / SF_TILE_DIM_M; - int num_tiles_k = padded_k / SF_TILE_DIM_K; - - int vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1)); - if (vec_load_size == 3) vec_load_size = 1; - int n_tiles_in_tb = TB_DIM * vec_load_size; - - int grid_dim_x = rowwise ? DIVUP(num_tiles_k, n_tiles_in_tb) : DIVUP(num_tiles_k, TB_DIM); - int grid_dim_y = rowwise ? num_tiles_m : DIVUP(num_tiles_m, vec_load_size); - local_blocks += grid_dim_x * grid_dim_y; - } - - for (int offset = 16; offset > 0; offset /= 2) { - local_blocks += __shfl_down_sync(0xffffffff, local_blocks, offset); + // The per-tensor shapes live only in device memory, so the grid cannot be + // sized on the host. A persistent grid instead walks a linear block space that + // the kernel measures for itself. One thread scans rather than a warp + // reducing: num_tensors is the expert count, so this costs less than the + // per-iteration scan below, and it leaves the kernel free of any assumption + // about how many lanes execute in lockstep. + extern __shared__ int s_total_blocks[]; + 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(); } - if (threadIdx.x == 0) *s_total_blocks = local_blocks; + s_total_blocks[0] = total; } __syncthreads(); + // Read into a register before the loop: the tile staging area starts at this + // same word, and the barrier below is what hands it over. + const int total_blocks = s_total_blocks[0]; - const int total_blocks = *s_total_blocks; - - // Persistent-grid loop for (int linear_block_id = blockIdx.x; linear_block_id < total_blocks; linear_block_id += gridDim.x) { - // Discover tensor_id and local_block_id via linear scan - int tensor_id = 0; - int current_block_base = 0; - size_t current_scale_base = 0; - int grid_dim_x = 0; - int grid_dim_y = 0; - size_t M = 0, K = 0; - int vec_load_size = 0; - + // The tile implementation leaves no trailing barrier, so without this one + // the next iteration's writes into the staging area would overlap this + // iteration's reads out of it. The loop bound is block-uniform, so every + // thread reaches it. + __syncthreads(); + + // Map the linear block id onto its tensor. Tensors contributing no blocks + // 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) { - M = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k); - K = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m); - - size_t padded_m = round_up_to_multiple(M, 128); - size_t padded_k = round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4); - - int num_tiles_m = padded_m / SF_TILE_DIM_M; - int num_tiles_k = padded_k / SF_TILE_DIM_K; - - vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1)); - if (vec_load_size == 3) vec_load_size = 1; - int n_tiles_in_tb = TB_DIM * vec_load_size; - - grid_dim_x = rowwise ? DIVUP(num_tiles_k, n_tiles_in_tb) : DIVUP(num_tiles_k, TB_DIM); - grid_dim_y = rowwise ? num_tiles_m : DIVUP(num_tiles_m, vec_load_size); - int blocks_i = grid_dim_x * grid_dim_y; - - if (linear_block_id < current_block_base + blocks_i) { - tensor_id = i; - break; - } - current_block_base += blocks_i; - current_scale_base += padded_m * padded_k * scale_elem_size; + 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; } - int local_block_id = linear_block_id - current_block_base; - int block_x = local_block_id % grid_dim_x; - int block_y = local_block_id / grid_dim_x; + 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) + current_scale_base; - uint8_t* output_base = reinterpret_cast(output) + current_scale_base; + 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(round_up_to_multiple(M, 128)); - const int padded_k = - static_cast(round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4)); - const int original_M = static_cast(M); - const int original_K = static_cast(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE))); + const 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); @@ -2346,9 +2386,8 @@ int grouped_swizzle_variable_max_active_blocks_per_sm(int device_id) { NVTE_CHECK(0 <= device_id && device_id < cuda::num_devices(), "invalid CUDA device ID"); auto init = [&]() { - constexpr int metadata_shmem = sizeof(int); // s_total_blocks constexpr int dynamic_smem_size = - TB_DIM * 4 * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t) + metadata_shmem; + 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, @@ -2360,7 +2399,6 @@ int grouped_swizzle_variable_max_active_blocks_per_sm(int device_id) { std::call_once(flags[device_id], init); return cache[device_id]; } -#endif void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* output, cudaStream_t stream) { @@ -2507,16 +2545,13 @@ 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); - 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; + const int dynamic_smem_size = grouped_variable_shape_smem_size(SF_TILE_DIM_M, SF_TILE_DIM_K); 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 +2588,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 } } From d86ae68f18d53f4033ad0363bac4d5ca922bf308 Mon Sep 17 00:00:00 2001 From: root Date: Wed, 5 Aug 2026 00:05:46 +0000 Subject: [PATCH 2/3] Tighten comments in the variable-shape grouped swizzle Blocks ran four to six lines where this file's convention is one or two. The reasoning they carried belongs in review notes rather than the source. Each why is kept -- the gfx942 LDS ceiling, the missing trailing barrier, the unpadded ROCm scale stride, the divide-by-zero -- and the justification around it is not. No code change. Co-Authored-By: Claude Opus 5 (1M context) --- tests/cpp/operator/test_swizzle.cu | 32 +++++---------- transformer_engine/common/swizzle/swizzle.cu | 41 +++++++------------- 2 files changed, 23 insertions(+), 50 deletions(-) diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index 7ef7576fb..aa665c085 100644 --- a/tests/cpp/operator/test_swizzle.cu +++ b/tests/cpp/operator/test_swizzle.cu @@ -521,10 +521,8 @@ void performTestGroupedSwizzleUnswizzleRoundtrip(const int num_tensors, const si #endif // !__HIP_PLATFORM_AMD__ (grouped unswizzle / roundtrip suites) -// Geometry of one tensor's MXFP8 scale block in the per-tensor padded layout the -// grouped swizzle contracts on. `dim_128` is the extent the swizzle tiles by 128 -// and `dim_4` the one it tiles by 4; columnwise scales swap which logical -// dimension supplies each, and store transposed. +// 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; @@ -549,15 +547,8 @@ static VariableScaleBlock variable_scale_block(size_t M, size_t K, bool rowwise) return block; } -// Variable-shape grouped swizzle, against the CPU reference. -// -// The scale buffers are built here rather than gathered from test::Tensor. This -// entry point contracts on the per-tensor *padded* layout, which is not what -// test::Tensor allocates on both platforms -- ROCm quantizers emit unpadded -// scales, so scale_tensor_alignment_* is 1 there (test_common.h) and a gathered -// buffer would carry a per-tensor stride the kernel does not expect. Building -// the layout under test directly also makes a zero-token expert expressible, -// which is the case variable shapes exist for. +// 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 performTestGroupedSwizzleMXFP8Variable(const std::vector>& shapes) { using namespace transformer_engine; using namespace test; @@ -576,9 +567,8 @@ void performTestGroupedSwizzleMXFP8Variable(const std::vector byte_dist(0, 255); - // Host-side input blocks and their expected swizzled form. Padding stays zero - // on the input so the reference, which permutes whatever it is given, agrees - // with the kernel, which zeroes anything past the valid extent. + // 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; @@ -655,9 +645,8 @@ void performTestGroupedSwizzleMXFP8Variable(const std::vector data_offsets(num_tensors + 1, 0); for (size_t i = 0; i < num_tensors; ++i) { @@ -769,9 +758,8 @@ INSTANTIATE_TEST_SUITE_P( // all three vectorized load widths are selected within one launch. std::vector>{{256, 128}, {130, 256}, {512, 512}, {64, 96}}, - // Case 8-10: an expert routed no tokens, in leading, middle and trailing - // position. Its scale block is empty and must contribute no work without - // disturbing the offsets of its neighbours. + // 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}}, diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index acfb649f8..897d9d2f0 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -2212,17 +2212,13 @@ void nvte_multi_tensor_unswizzle_scaling_factors(const NVTETensor* inputs, NVTET namespace transformer_engine { -// Dynamic shared memory the variable-shape kernel reserves. The vectorized load -// width is a per-tensor runtime value here, so every block must be able to stage -// the widest specialization. This comes to exactly 64 KiB, which is also the -// per-workgroup LDS ceiling on gfx942-class parts, so the measured block count -// shares the first word of this buffer rather than being reserved alongside it. +// Sized for the widest vectorized load, which is chosen per tensor at runtime. +// Exactly 64 KiB, the gfx942 LDS ceiling, so the block count aliases word 0. 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)); } -// Extents of one tensor's scale block, in rows and columns of its row-major -// storage. Columnwise scales transpose the tensor's logical dimensions. +// One tensor's scale-block extents; columnwise transposes the logical dims. struct ScaleBlockDims { size_t m; size_t k; @@ -2236,10 +2232,8 @@ __device__ __forceinline__ ScaleBlockDims variable_shape_dims(int i, const int64 return rowwise ? ScaleBlockDims{first_dim, last_dim} : ScaleBlockDims{last_dim, first_dim}; } -// Per-tensor launch geometry, mirroring what the uniform-shape path computes on -// the host. The kernel walks this twice -- once to measure the workload, once to -// map a linear block id back onto a tensor -- and the two passes must agree, so -// the arithmetic has a single definition. +// 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; @@ -2260,9 +2254,8 @@ __device__ __forceinline__ VariableShapeTiling variable_shape_tiling(ScaleBlockD 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; - // An expert that was routed no tokens has an empty scale block and contributes - // no work. Returning here also keeps the zero out of the modulo below, which - // would otherwise leave vec_load_size at 0 and then divide by it. + // 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; @@ -2286,12 +2279,8 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) int num_tensors, bool rowwise, size_t scale_elem_size, size_t common_m, size_t common_k) { - // The per-tensor shapes live only in device memory, so the grid cannot be - // sized on the host. A persistent grid instead walks a linear block space that - // the kernel measures for itself. One thread scans rather than a warp - // reducing: num_tensors is the expert count, so this costs less than the - // per-iteration scan below, and it leaves the kernel free of any assumption - // about how many lanes execute in lockstep. + // Shapes are device-side, so the grid cannot be sized on the host and the + // kernel measures its own block space. Scanned serially to assume no warp width. extern __shared__ int s_total_blocks[]; if (threadIdx.x == 0 && threadIdx.y == 0) { int total = 0; @@ -2303,20 +2292,16 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) s_total_blocks[0] = total; } __syncthreads(); - // Read into a register before the loop: the tile staging area starts at this - // same word, and the barrier below is what hands it over. + // 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 implementation leaves no trailing barrier, so without this one - // the next iteration's writes into the staging area would overlap this - // iteration's reads out of it. The loop bound is block-uniform, so every - // thread reaches it. + // The tile impl has no trailing barrier, so consecutive iterations would + // overlap in the staging area. The loop bound is block-uniform. __syncthreads(); - // Map the linear block id onto its tensor. Tensors contributing no blocks - // are skipped naturally: the strict comparison never selects them. + // Empty tensors are skipped naturally: the strict comparison never selects them. ScaleBlockDims dims{}; VariableShapeTiling tiling{}; size_t scale_base_bytes = 0; From 473d268262421811a523cbfc64ff757ce9333a65 Mon Sep 17 00:00:00 2001 From: root Date: Wed, 5 Aug 2026 00:31:58 +0000 Subject: [PATCH 3/3] Keep the CUDA variable-shape swizzle path on upstream's code The previous commits changed code CUDA compiles: the block-count reduction, the shared memory request, and the variable-shape test body. None of that was needed to enable the feature on ROCm, and it put the fork ahead of upstream on a path upstream owns and tests. The kernel is now split. CUDA keeps upstream's implementation verbatim, wave32 shuffle and all, including its persistent-loop barrier omission and its zero-tile divide -- those are upstream defects to report, not to patch here. ROCm gets its own arm with the serial scan, the 64 KiB shared memory request, and both fixes. The test follows the same split: upstream's suite is restored unchanged for CUDA, and the ROCm layout builder becomes a separate suite. test::Tensor emits unpadded MXFP8 scales on ROCm, so the two cannot share a body. Every deletion against upstream is now guard scaffolding or the ROCm NVTE_ERROR. Co-Authored-By: Claude Opus 5 (1M context) --- tests/cpp/operator/test_swizzle.cu | 153 ++++++++++++++++- transformer_engine/common/swizzle/swizzle.cu | 168 ++++++++++++++++++- 2 files changed, 311 insertions(+), 10 deletions(-) diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index aa665c085..787a83f8b 100644 --- a/tests/cpp/operator/test_swizzle.cu +++ b/tests/cpp/operator/test_swizzle.cu @@ -519,8 +519,145 @@ void performTestGroupedSwizzleUnswizzleRoundtrip(const int num_tensors, const si num_tensors * col_numel); } -#endif // !__HIP_PLATFORM_AMD__ (grouped unswizzle / roundtrip suites) +void performTestGroupedSwizzleMXFP8Variable(const std::vector>& shapes) { + using namespace transformer_engine; + using namespace test; + + int num_tensors = shapes.size(); + std::vector> input_tensors; + std::vector> output_tensors; + std::vector input_ptrs; + std::vector output_ptrs; + input_tensors.reserve(num_tensors); + output_tensors.reserve(num_tensors); + input_ptrs.reserve(num_tensors); + output_ptrs.reserve(num_tensors); + + constexpr size_t BLOCK_SIZE = 32; + for (int i = 0; i < num_tensors; ++i) { + const std::vector shape{shapes[i].first, shapes[i].second}; + auto input = std::make_unique("input_" + std::to_string(i), shape, + DType::kFloat8E4M3, true, true, + NVTE_MXFP8_1D_SCALING); + auto output = std::make_unique("output_" + std::to_string(i), shape, + DType::kFloat8E4M3, true, true, + NVTE_MXFP8_1D_SCALING); + fillUniform(input.get()); + fillUniform(output.get()); + + // Zero padding + input->to_cpu(); + const NVTEShape rs = input->rowwise_scale_inv_shape(); + zero_scale_inv_padding(input->rowwise_cpu_scale_inv_ptr(), + rs.data[0], rs.data[1], + shapes[i].first, (shapes[i].second + BLOCK_SIZE - 1) / BLOCK_SIZE); + const NVTEShape cs = input->columnwise_scale_inv_shape(); + zero_scale_inv_padding(input->columnwise_cpu_scale_inv_ptr(), + cs.data[0], cs.data[1], + (shapes[i].first + BLOCK_SIZE - 1) / BLOCK_SIZE, shapes[i].second); + input->from_cpu(); + + input_ptrs.push_back(input.get()); + output_ptrs.push_back(output.get()); + input_tensors.emplace_back(std::move(input)); + output_tensors.emplace_back(std::move(output)); + } + + GroupedBuffers grouped_input = build_grouped_tensor(input_ptrs, NVTE_MXFP8_1D_SCALING); + GroupedBuffers grouped_output = build_grouped_tensor(output_ptrs, NVTE_MXFP8_1D_SCALING); + + const uint8_t input_swizzled = 0; + nvte_set_grouped_tensor_param(grouped_input.get_handle(), + kNVTEGroupedWithGEMMSwizzledScales, + &input_swizzled, sizeof(input_swizzled)); + const uint8_t output_swizzled = 1; + nvte_set_grouped_tensor_param(grouped_output.get_handle(), + kNVTEGroupedWithGEMMSwizzledScales, + &output_swizzled, sizeof(output_swizzled)); + + nvte_swizzle_grouped_scaling_factors(grouped_input.get_handle(), + grouped_output.get_handle(), + 0); + cudaDeviceSynchronize(); + NVTE_CHECK_CUDA(cudaGetLastError()); + + // Verification + size_t row_offset = 0; + size_t col_offset = 0; + for (int i = 0; i < num_tensors; ++i) { + const NVTEShape row_shape = input_tensors[i]->rowwise_scale_inv_shape(); + const NVTEShape col_shape = input_tensors[i]->columnwise_scale_inv_shape(); + const size_t row_numel = row_shape.data[0] * row_shape.data[1]; + const size_t col_numel = col_shape.data[0] * col_shape.data[1]; + + std::vector output_row_host(row_numel); + std::vector output_col_host(col_numel); + NVTE_CHECK_CUDA(cudaMemcpy(output_row_host.data(), + static_cast(grouped_output.scale_inv.get()) + row_offset, + row_numel, cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaMemcpy(output_col_host.data(), + static_cast(grouped_output.columnwise_scale_inv.get()) + col_offset, + col_numel, cudaMemcpyDeviceToHost)); + + std::vector ref_row(row_numel); + std::vector ref_col(col_numel); + compute_ref_swizzle<128, 4, true>(input_tensors[i]->rowwise_cpu_scale_inv_ptr(), + ref_row.data(), + row_shape.data[0], row_shape.data[1]); + compute_ref_swizzle<128, 4, false>( + input_tensors[i]->columnwise_cpu_scale_inv_ptr(), + ref_col.data(), + col_shape.data[1], col_shape.data[0]); + + compareResults("grouped_swizzle_variable_rowwise_" + std::to_string(i), + output_row_host.data(), ref_row.data(), row_numel); + compareResults("grouped_swizzle_variable_colwise_" + std::to_string(i), + output_col_host.data(), ref_col.data(), col_numel); + + row_offset += row_numel; + col_offset += col_numel; + } +} + +class SwizzleGroupedVariableTestSuite + : public ::testing::TestWithParam>> {}; + +TEST_P(SwizzleGroupedVariableTestSuite, TestGroupedSwizzleMXFP8Variable) { + const auto shapes = GetParam(); + performTestGroupedSwizzleMXFP8Variable(shapes); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + SwizzleGroupedVariableTestSuite, + ::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}} + ), + [](const testing::TestParamInfo& info) { + return "VariableShapes_" + std::to_string(info.index) + "_N" + std::to_string(info.param.size()); + } +); + +#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 { @@ -549,7 +686,7 @@ static VariableScaleBlock variable_scale_block(size_t M, size_t K, bool rowwise) // 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 performTestGroupedSwizzleMXFP8Variable(const std::vector>& shapes) { +void performTestGroupedSwizzleMXFP8VariableRocm(const std::vector>& shapes) { using namespace transformer_engine; using namespace test; @@ -724,17 +861,17 @@ void performTestGroupedSwizzleMXFP8Variable(const std::vector>> {}; -TEST_P(SwizzleGroupedVariableTestSuite, TestGroupedSwizzleMXFP8Variable) { +TEST_P(SwizzleGroupedVariableRocmTestSuite, TestGroupedSwizzleMXFP8Variable) { const auto shapes = GetParam(); - performTestGroupedSwizzleMXFP8Variable(shapes); + performTestGroupedSwizzleMXFP8VariableRocm(shapes); } INSTANTIATE_TEST_SUITE_P( OperatorTest, - SwizzleGroupedVariableTestSuite, + SwizzleGroupedVariableRocmTestSuite, ::testing::Values( // Case 1: num_tensors = 1 (n+3 = 4, even). Check simple alignment. std::vector>{{1024, 1024}}, @@ -767,11 +904,13 @@ INSTANTIATE_TEST_SUITE_P( // Case 11: a zero last dim, which empties the columnwise block instead. std::vector>{{128, 0}, {128, 256}} ), - [](const testing::TestParamInfo& info) { + [](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/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 897d9d2f0..91d1c3824 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -2212,8 +2212,160 @@ void nvte_multi_tensor_unswizzle_scaling_factors(const NVTETensor* inputs, NVTET namespace transformer_engine { -// Sized for the widest vectorized load, which is chosen per tensor at runtime. -// Exactly 64 KiB, the gfx942 LDS ceiling, so the block count aliases word 0. +#ifndef __HIP_PLATFORM_AMD__ +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) { + extern __shared__ int s_metadata[]; + int* s_total_blocks = &s_metadata[0]; + + // Warp reduction to compute total workload + if (threadIdx.x < 32 && threadIdx.y == 0) { + int local_blocks = 0; + for (int i = threadIdx.x; i < num_tensors; i += 32) { + size_t m = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k); + size_t k = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m); + + size_t padded_m = round_up_to_multiple(m, 128); + size_t padded_k = round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); + + int num_tiles_m = padded_m / SF_TILE_DIM_M; + int num_tiles_k = padded_k / SF_TILE_DIM_K; + + int vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1)); + if (vec_load_size == 3) vec_load_size = 1; + int n_tiles_in_tb = TB_DIM * vec_load_size; + + int grid_dim_x = rowwise ? DIVUP(num_tiles_k, n_tiles_in_tb) : DIVUP(num_tiles_k, TB_DIM); + int grid_dim_y = rowwise ? num_tiles_m : DIVUP(num_tiles_m, vec_load_size); + local_blocks += grid_dim_x * grid_dim_y; + } + + for (int offset = 16; offset > 0; offset /= 2) { + local_blocks += __shfl_down_sync(0xffffffff, local_blocks, offset); + } + if (threadIdx.x == 0) *s_total_blocks = local_blocks; + } + __syncthreads(); + + const int total_blocks = *s_total_blocks; + + // Persistent-grid loop + for (int linear_block_id = blockIdx.x; linear_block_id < total_blocks; + linear_block_id += gridDim.x) { + // Discover tensor_id and local_block_id via linear scan + int tensor_id = 0; + int current_block_base = 0; + size_t current_scale_base = 0; + int grid_dim_x = 0; + int grid_dim_y = 0; + size_t M = 0, K = 0; + int vec_load_size = 0; + + for (int i = 0; i < num_tensors; ++i) { + M = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k); + K = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m); + + size_t padded_m = round_up_to_multiple(M, 128); + size_t padded_k = round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4); + + int num_tiles_m = padded_m / SF_TILE_DIM_M; + int num_tiles_k = padded_k / SF_TILE_DIM_K; + + vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1)); + if (vec_load_size == 3) vec_load_size = 1; + int n_tiles_in_tb = TB_DIM * vec_load_size; + + grid_dim_x = rowwise ? DIVUP(num_tiles_k, n_tiles_in_tb) : DIVUP(num_tiles_k, TB_DIM); + grid_dim_y = rowwise ? num_tiles_m : DIVUP(num_tiles_m, vec_load_size); + int blocks_i = grid_dim_x * grid_dim_y; + + if (linear_block_id < current_block_base + blocks_i) { + tensor_id = i; + break; + } + current_block_base += blocks_i; + current_scale_base += padded_m * padded_k * scale_elem_size; + } + + int local_block_id = linear_block_id - current_block_base; + int block_x = local_block_id % grid_dim_x; + int block_y = local_block_id / grid_dim_x; + + const uint8_t* input_base = reinterpret_cast(input) + current_scale_base; + uint8_t* output_base = reinterpret_cast(output) + current_scale_base; + + const int padded_m = static_cast(round_up_to_multiple(M, 128)); + const int padded_k = + static_cast(round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4)); + const int original_M = static_cast(M); + const int original_K = static_cast(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE))); + const bool padding_m = (block_y == grid_dim_y - 1) && (original_M < padded_m); + 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 metadata_shmem = sizeof(int); // s_total_blocks + constexpr int dynamic_smem_size = + TB_DIM * 4 * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t) + metadata_shmem; + 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]; +} + +#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)); } @@ -2280,8 +2432,10 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) 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. Scanned serially to assume no warp width. + // 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) { @@ -2385,6 +2539,8 @@ int grouped_swizzle_variable_max_active_blocks_per_sm(int device_id) { return cache[device_id]; } +#endif // __HIP_PLATFORM_AMD__ + void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* output, cudaStream_t stream) { // Check scaling mode @@ -2536,7 +2692,13 @@ void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* 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;