From 3be50bc7219d5f333c01d261e5afe324848892ec Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 21 Oct 2026 21:56:18 +0000 Subject: [PATCH 1/8] Add gfx1250 Tensor Data Mover helper util/rocm_tdm.cuh wraps the gfx1250 TDM builtins for asynchronous 2D tile copies between global memory and LDS (tdm::load_2d, tdm::store_2d, tdm::wait). It packs the descriptors itself instead of using hip/amd_detail/amd_gfx1250_TDM.h, whose stride setters truncate, and is only defined when compiling for gfx1250 with a compiler that provides the unified TDM builtins (NVTE_ROCM_HAS_TDM). test_rocm_tdm covers per-wave load/store round trips for 1-8 byte elements, partial edge tiles and a row pitch different from the width. --- tests/cpp/util/CMakeLists.txt | 3 + tests/cpp/util/test_rocm_tdm.cu | 118 ++++++++++++++++++++ transformer_engine/common/util/rocm_tdm.cuh | 116 +++++++++++++++++++ 3 files changed, 237 insertions(+) create mode 100644 tests/cpp/util/test_rocm_tdm.cu create mode 100644 transformer_engine/common/util/rocm_tdm.cuh diff --git a/tests/cpp/util/CMakeLists.txt b/tests/cpp/util/CMakeLists.txt index 9d25c9297..d8cf3aa8b 100644 --- a/tests/cpp/util/CMakeLists.txt +++ b/tests/cpp/util/CMakeLists.txt @@ -8,6 +8,9 @@ list(APPEND test_cuda_sources test_nvrtc.cpp test_string.cpp ../test_common.cu) +if(USE_ROCM) + list(APPEND test_cuda_sources test_rocm_tdm.cu) +endif() if(USE_CUDA) add_executable(test_util ${test_cuda_sources}) else() diff --git a/tests/cpp/util/test_rocm_tdm.cu b/tests/cpp/util/test_rocm_tdm.cu new file mode 100644 index 000000000..cc8530b01 --- /dev/null +++ b/tests/cpp/util/test_rocm_tdm.cu @@ -0,0 +1,118 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ + +#include +#include + +#include +#include + +#include "util/rocm_tdm.cuh" +#include "../test_common.h" + +namespace { + +constexpr uint32_t kWaves = 4; +constexpr uint32_t kTileRows = 8; +constexpr uint32_t kTileCols = 64; + +// Each wave loads its own tile into LDS, dumps it, adds 1 and stores it back. +template +__global__ void __launch_bounds__(kWaves * 32) + tdm_roundtrip_kernel(const T *in, T *out, T *lds_dump, uint32_t rows, uint32_t cols, + uint32_t ld_in, uint32_t ld_out) { +#ifdef NVTE_ROCM_HAS_TDM + __shared__ alignas(128) T tile[kWaves][kTileRows][kTileCols]; + const uint32_t wave = threadIdx.x / 32; + const uint32_t row = (blockIdx.y * kWaves + wave) * kTileRows; + const uint32_t col = blockIdx.x * kTileCols; + transformer_engine::tdm::load_2d(tile[wave], in, rows, cols, ld_in, row, col); + transformer_engine::tdm::wait<0>(); + T *dump = lds_dump + ((blockIdx.y * gridDim.x + blockIdx.x) * kWaves + wave) * kTileRows * kTileCols; + for (uint32_t i = threadIdx.x % 32; i < kTileRows * kTileCols; i += 32) { + T &v = tile[wave][i / kTileCols][i % kTileCols]; + dump[i] = v; + v = v + 1; + } + transformer_engine::tdm::store_2d(out, tile[wave], rows, cols, ld_out, row, col); + transformer_engine::tdm::wait<0>(); +#endif +} + +template +void run_roundtrip(uint32_t rows, uint32_t cols, uint32_t ld_in, uint32_t ld_out) { + SCOPED_TRACE(testing::Message() << "rows=" << rows << " cols=" << cols << " ld_in=" << ld_in + << " ld_out=" << ld_out); + const dim3 grid((cols + kTileCols - 1) / kTileCols, + (rows + kWaves * kTileRows - 1) / (kWaves * kTileRows)); + // Every row a tile can touch, so stray stores are detected. + const size_t out_rows = grid.y * kWaves * kTileRows; + const T sentinel = static_cast(0xA5A5A5A5A5A5A5A5ull); + const T pad = static_cast(0xEEEEEEEEEEEEEEEEull); + + std::vector h_in(static_cast(rows) * ld_in); + std::vector h_out(out_rows * ld_out, sentinel); + std::vector h_dump(static_cast(grid.x) * grid.y * kWaves * kTileRows * kTileCols); + for (size_t r = 0; r < rows; ++r) { + for (size_t c = 0; c < ld_in; ++c) { + h_in[r * ld_in + c] = c < cols ? static_cast((r * 37 + c * 11) % 150 + 1) : pad; + } + } + + T *d_in, *d_out, *d_dump; + NVTE_CHECK_CUDA(cudaMalloc(&d_in, h_in.size() * sizeof(T))); + NVTE_CHECK_CUDA(cudaMalloc(&d_out, h_out.size() * sizeof(T))); + NVTE_CHECK_CUDA(cudaMalloc(&d_dump, h_dump.size() * sizeof(T))); + NVTE_CHECK_CUDA(cudaMemcpy(d_in, h_in.data(), h_in.size() * sizeof(T), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(d_out, h_out.data(), h_out.size() * sizeof(T), cudaMemcpyHostToDevice)); + tdm_roundtrip_kernel<<>>(d_in, d_out, d_dump, rows, cols, ld_in, ld_out); + NVTE_CHECK_CUDA(cudaGetLastError()); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + NVTE_CHECK_CUDA(cudaMemcpy(h_out.data(), d_out, h_out.size() * sizeof(T), cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaMemcpy(h_dump.data(), d_dump, h_dump.size() * sizeof(T), cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaFree(d_in)); + NVTE_CHECK_CUDA(cudaFree(d_out)); + NVTE_CHECK_CUDA(cudaFree(d_dump)); + + size_t load_mismatches = 0; + for (size_t tile = 0; tile < h_dump.size() / (kTileRows * kTileCols); ++tile) { + const size_t block = tile / kWaves, wave = tile % kWaves; + const size_t row0 = ((block / grid.x) * kWaves + wave) * kTileRows; + const size_t col0 = (block % grid.x) * kTileCols; + for (size_t i = 0; i < kTileRows * kTileCols; ++i) { + const size_t r = row0 + i / kTileCols, c = col0 + i % kTileCols; + const T expected = (r < rows && c < cols) ? h_in[r * ld_in + c] : T(0); + load_mismatches += h_dump[tile * kTileRows * kTileCols + i] != expected; + } + } + size_t store_mismatches = 0; + for (size_t r = 0; r < out_rows; ++r) { + for (size_t c = 0; c < ld_out; ++c) { + const T expected = + (r < rows && c < cols) ? static_cast(h_in[r * ld_in + c] + 1) : sentinel; + store_mismatches += h_out[r * ld_out + c] != expected; + } + } + EXPECT_EQ(load_mismatches, 0u) << "LDS contents after load_2d (out of bounds must read 0)"; + EXPECT_EQ(store_mismatches, 0u) << "global contents after store_2d (out of bounds untouched)"; +} + +} // namespace + +template +class RocmTdmTest : public ::testing::Test {}; + +using TdmElementTypes = ::testing::Types; +TYPED_TEST_SUITE(RocmTdmTest, TdmElementTypes); + +TYPED_TEST(RocmTdmTest, RoundTrip2D) { + if (test::getDeviceComputeCapability() != 125) { + GTEST_SKIP() << "TDM requires gfx1250"; + } + run_roundtrip(64, 128, 128, 128); // whole tiles + run_roundtrip(100, 200, 256, 216); // partial tiles, row pitch != cols + run_roundtrip(1, 1, 16, 16); +} diff --git a/transformer_engine/common/util/rocm_tdm.cuh b/transformer_engine/common/util/rocm_tdm.cuh new file mode 100644 index 000000000..2cec11178 --- /dev/null +++ b/transformer_engine/common/util/rocm_tdm.cuh @@ -0,0 +1,116 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ +#pragma once +//#include "hip/hip_runtime.h" // prevent hipification of this rocm_ file + +// gfx1250 Tensor Data Mover (TDM): asynchronous 2D tile copies between global memory and LDS. +// Descriptors are packed here rather than with hip/amd_detail/amd_gfx1250_TDM.h, whose stride +// setters truncate (ROCm/rocm-systems#13274). + +#include +#include + +#if defined(__gfx1250__) && __has_builtin(__builtin_amdgcn_tensor_load_to_lds) && \ + __has_builtin(__builtin_amdgcn_tensor_store_from_lds) && \ + __has_builtin(__builtin_amdgcn_s_wait_tensorcnt) +#define NVTE_ROCM_HAS_TDM 1 +#endif + +#ifdef NVTE_ROCM_HAS_TDM +namespace transformer_engine { +namespace tdm { +namespace detail { + +// The builtins' operand signedness has changed between compiler versions; integer vectors of +// either signedness convert implicitly. +typedef uint32_t u32x4 __attribute__((ext_vector_type(4))); +typedef uint32_t u32x8 __attribute__((ext_vector_type(8))); + +template +__device__ __forceinline__ void copy_2d(const T (&lds)[TILE_ROWS][TILE_COLS], const T *gmem, + uint32_t rows, uint32_t cols, uint64_t ld, uint32_t row, + uint32_t col) { + static_assert(sizeof(T) == 1 || sizeof(T) == 2 || sizeof(T) == 4 || sizeof(T) == 8, + "TDM elements are 1, 2, 4 or 8 bytes"); + static_assert(TILE_ROWS > 0 && TILE_ROWS <= 0xFFFF && TILE_COLS > 0 && TILE_COLS <= 0xFFFF, + "TDM tile dimensions are 16-bit"); + constexpr uint32_t data_size = sizeof(T) == 1 ? 0 : sizeof(T) == 2 ? 1 : sizeof(T) == 4 ? 2 : 3; + + // The descriptor addresses the tile origin and the tensor extents remaining from there, so the + // TDM zero-fills (load) or skips (store) everything outside the tensor. + const uint64_t addr = reinterpret_cast(gmem) + (row * ld + col) * sizeof(T); + const uint32_t dim0 = col < cols ? cols - col : 0; + const uint32_t dim1 = row < rows ? rows - row : 0; + + u32x4 g0; + g0[0] = 1; // count + g0[1] = static_cast(reinterpret_cast(&lds)); // low bits: LDS offset + g0[2] = static_cast(addr); + g0[3] = (static_cast(addr >> 32) & 0x1FFFFFF) | (2u << 30); // type = 2 + + u32x8 g1; + g1[0] = data_size << 16; + g1[1] = (dim0 & 0xFFFF) << 16; + g1[2] = (dim0 >> 16) | ((dim1 & 0xFFFF) << 16); + g1[3] = (dim1 >> 16) | (static_cast(TILE_COLS) << 16); + g1[4] = static_cast(TILE_ROWS); + g1[5] = static_cast(ld); + g1[6] = static_cast(ld >> 32) & 0xFFFF; + g1[7] = 0; + + // Descriptors are scalar operands. +#pragma unroll + for (int i = 0; i < 4; ++i) { + g0[i] = __builtin_amdgcn_readfirstlane(g0[i]); + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + g1[i] = __builtin_amdgcn_readfirstlane(g1[i]); + } + const u32x4 z4 = {0, 0, 0, 0}; + const u32x8 z8 = {0, 0, 0, 0, 0, 0, 0, 0}; + + // The TDM does not order against this wave's outstanding LDS accesses. + asm volatile("s_wait_dscnt 0x0" ::: "memory"); + if constexpr (STORE) { + __builtin_amdgcn_tensor_store_from_lds(g0, g1, z4, z4, z8, 0); + } else { + __builtin_amdgcn_tensor_load_to_lds(g0, g1, z4, z4, z8, 0); + } +} + +} // namespace detail + +// Starts copying the tile at (row, col) of a rows x cols row-major tensor with a row pitch of ld +// elements into the dense LDS array lds; elements outside the tensor read as zero. Arguments must +// be wave-uniform, and every calling wave issues its own copy. Complete it with wait(). +template +__device__ __forceinline__ void load_2d(T (&lds)[TILE_ROWS][TILE_COLS], const T *gmem, + uint32_t rows, uint32_t cols, uint64_t ld, uint32_t row, + uint32_t col) { + detail::copy_2d(lds, gmem, rows, cols, ld, row, col); +} + +// Inverse of load_2d: elements outside the tensor are not written. lds must stay unmodified +// until wait() returns. +template +__device__ __forceinline__ void store_2d(T *gmem, const T (&lds)[TILE_ROWS][TILE_COLS], + uint32_t rows, uint32_t cols, uint64_t ld, uint32_t row, + uint32_t col) { + detail::copy_2d(lds, gmem, rows, cols, ld, row, col); +} + +// Waits until at most N copies issued by this wave are in flight. Other waves of the block see +// the loaded data only after a subsequent barrier. +template +__device__ __forceinline__ void wait() { + __builtin_amdgcn_s_wait_tensorcnt(N); + asm volatile("" ::: "memory"); // the compiler does not see the TDM's LDS writes +} + +} // namespace tdm +} // namespace transformer_engine +#endif // NVTE_ROCM_HAS_TDM From 38c0d8d20a9a2a122437d73f27277e56a4c32746 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 21 Oct 2026 21:56:18 +0000 Subject: [PATCH 2/8] use TDM for MXFP8 quant --- .../common/cast/dispatch/quantize.cuh | 4 + .../common/cast/mxfp8/dequantize_mxfp8.cuh | 8 + .../common/cast/mxfp8/gated_mxfp8.cuh | 8 + .../cast/mxfp8/group_quantize_mxfp8.cuh | 5 + .../common/cast/mxfp8/quantize_mxfp8.cuh | 9 + .../cast/mxfp8/rocm_dequantize_mxfp8.cuh | 14 +- .../cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh | 150 ++++ .../cast/mxfp8/rocm_gated_mxfp8_tdm.cuh | 168 +++++ .../cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh | 669 ++++++++++++++++++ transformer_engine/common/util/rocm_tdm.cuh | 40 +- 10 files changed, 1070 insertions(+), 5 deletions(-) create mode 100644 transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh create mode 100644 transformer_engine/common/cast/mxfp8/rocm_gated_mxfp8_tdm.cuh create mode 100644 transformer_engine/common/cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 78a49c516..81676eed8 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -729,6 +729,10 @@ inline void multi_quantize_mxfp8(const std::vector &input_list, "multi_quantize_mxfp8: num_tensors (", num_tensors, ") exceeds maximum (", mxfp8::quantize_kernel::kMultiQuantizeMXFP8MaxTensors, ")."); + if (mxfp8::multi_quantize_tdm(input_list, output_list, stream)) { + return; + } + DType itype = input_list[0]->data.dtype; DType otype = output_list[0]->dtype(); const bool use_rowwise = output_list[0]->has_data(); diff --git a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh index 3da9a7106..595001a14 100644 --- a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh @@ -28,6 +28,9 @@ #ifdef __HIP_PLATFORM_AMD__ #include "../../util/cuda_runtime.h" +#include "../../util/ptx_arch_spec.cuh" +#include "../../util/system.h" +#include "./rocm_quantize_mxfp8_tdm.cuh" #include "./rocm_vectorized_2d.cuh" #endif @@ -37,6 +40,7 @@ namespace mxfp8 { namespace dequantize_kernel { #ifdef __HIP_PLATFORM_AMD__ #include "rocm_dequantize_mxfp8.cuh" +#include "rocm_dequantize_mxfp8_tdm.cuh" #else constexpr size_t CHUNK_DIM_Y = 128; constexpr size_t CHUNK_DIM_X = 128; @@ -323,6 +327,10 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) // rows, colwise uses the columnwise-scale cols. const size_t mx_swizzle_padded_dim = use_rowwise_scaling ? input.scale_inv.shape[0] : input.columnwise_scale_inv.shape[1]; + if (dequantize_tdm(input, input_data, output, scales_ptr, scales_stride, mx_swizzle_padded_dim, + stream)) { + return; + } #endif TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh index 30ba08439..cfac5cfef 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh @@ -28,7 +28,10 @@ #ifdef __HIP_PLATFORM_AMD__ #include "./rocm_vectorized_2d.cuh" +#include "../../util/cuda_runtime.h" #include "../../util/rocm_device_utils.cuh" +#include "../../util/system.h" +#include "./rocm_quantize_mxfp8_tdm.cuh" #endif namespace transformer_engine { @@ -37,6 +40,7 @@ namespace mxfp8 { namespace gated_kernel { #ifdef __HIP_PLATFORM_AMD__ #include "rocm_gated_mxfp8.cuh" +#include "rocm_gated_mxfp8_tdm.cuh" #else constexpr size_t CHUNK_DIM_Y = 64; constexpr size_t CHUNK_DIM_X = 64; @@ -761,6 +765,10 @@ void quantize_gated(const Tensor &gated_input, const Tensor &grad, Tensor *outpu const size_t output_cols = (IS_BWD ? 2 : 1) * cols; #ifdef __HIP_PLATFORM_AMD__ + if (quantize_gated_tdm(gated_input, grad, output, p, stream)) { + return; + } + constexpr size_t TMA_SHMEM_ALIGNMENT = ALIGNMENT_SIZE; constexpr size_t BUFF_DIM_Y = BUFFER_DIM_Y; diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 1da3a3bd2..82ea039e7 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -1193,6 +1193,11 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations } } + if (group_quantize_tdm(input, activations, noop, output, + dbias, workspace, stream)) { + return; + } + TRANSFORMER_ENGINE_SWITCH_CONDITION( use_large_chunks, USE_LARGE_CHUNKS, diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index de0d54ebe..d4524a8c8 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -32,7 +32,10 @@ #ifdef __HIP_PLATFORM_AMD__ #include "./rocm_vectorized_2d.cuh" +#include "../../util/cuda_runtime.h" #include "../../util/rocm_device_utils.cuh" +#include "../../util/system.h" +#include "./rocm_quantize_mxfp8_tdm.cuh" #endif namespace transformer_engine { @@ -691,6 +694,12 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, const auto [rows, cols] = input.flat_2d_dims(); #ifdef __HIP_PLATFORM_AMD__ + if (!use_2d_quantization && + quantize_tdm(input, act_input, noop, output, dbias, + workspace, stream)) { + return; + } + // Choose chunk size based on tensor dimensions const bool use_large_chunks = (rows * cols > 32 * 1024 * 1024); diff --git a/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8.cuh index 7c4eead42..f2dcfbf1a 100644 --- a/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8.cuh @@ -32,6 +32,16 @@ __device__ __forceinline__ size_t mx_preswizzle_scale_idx(size_t i, size_t j, si return (j / GROUP) * (padded_dim * GROUP) + i * GROUP + (j % GROUP); } +// hip_bfloat16 rounds with per-element branches; __bf16 is branch-free (v_cvt_pk_bf16_f32 on gfx950+). +template +__device__ __forceinline__ T cvt_from_float(const float x) { + if constexpr (std::is_same_v) { + return __builtin_bit_cast(T, static_cast<__bf16>(x)); + } else { + return static_cast(x); + } +} + template __device__ __forceinline__ void @@ -124,14 +134,14 @@ __device__ __forceinline__ void #pragma unroll for (int j = 0; j < ELEMS_PER_THREAD; j++) { - out.data.elt[j] = static_cast(block_scale * static_cast(in.data.elt[j])); + out.data.elt[j] = cvt_from_float(block_scale * static_cast(in.data.elt[j])); } out.store_to(&out_sh[shmem_offset_y][shmem_offset_x]); } else { #pragma unroll for (int i = 0; i < BUFFER_DIM_Y; i++) { const float elt = static_cast(in_sh[i][tid_colwise_X]); - out_sh[i][tid_colwise_X] = static_cast(block_scale * elt); + out_sh[i][tid_colwise_X] = cvt_from_float(block_scale * elt); } } diff --git a/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh b/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh new file mode 100644 index 000000000..fb27953f3 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh @@ -0,0 +1,150 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ + +#pragma once +// gfx1250 MXFP8 dequantize with the Tensor Data Mover, included into dequantize_kernel after +// rocm_dequantize_mxfp8.cuh, whose grid, scale indexing and fallback chunk it shares. +//#include "hip/hip_runtime.h" //dummy include to prevent hipification adding this header + +// Each wave moves its own TDM_BAND_Y-row band of the chunk global->LDS->global, so waves never +// wait for each other. +constexpr size_t TDM_WAVE_SIZE = 32; +constexpr size_t TDM_WAVES = THREADS_PER_CHUNK / TDM_WAVE_SIZE; // 4 +constexpr size_t TDM_BAND_Y = CHUNK_DIM_Y / TDM_WAVES; // 32 + +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + dequantize_mxfp8_tdm_kernel(const IType *input_ptr, OType *output_ptr, + const e8m0_t *const scales_ptr, const size_t rows, + const size_t cols, const size_t scales_stride, + const size_t mx_swizzle_padded_dim) { +#ifdef NVTE_ROCM_HAS_TDM + constexpr bool USE_ROWWISE_SCALING = SCALE_DIM_X > 1; + + __shared__ alignas(128) IType in_sh[TDM_WAVES][TDM_BAND_Y][CHUNK_DIM_X]; + __shared__ alignas(128) OType out_sh[TDM_WAVES][TDM_BAND_Y][CHUNK_DIM_X]; + + const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / TDM_WAVE_SIZE); + const int lane = threadIdx.x % TDM_WAVE_SIZE; + const size_t band_offset_Y = blockIdx.y * CHUNK_DIM_Y + wave * TDM_BAND_Y; + const size_t chunk_offset_X = blockIdx.x * CHUNK_DIM_X; + if (band_offset_Y >= rows) { + return; + } + tdm::load_2d(in_sh[wave], input_ptr, rows, cols, cols, band_offset_Y, chunk_offset_X); + + // Scales are fetched while the copy is in flight. Outside the scale tensor a lane reads a dummy + // in-bounds scale and uses 2^0 instead, so the loads need no branches and are issued together. + const size_t scales_rows = USE_ROWWISE_SCALING ? rows : DIVUP(rows, SCALE_DIM_Y); + const size_t scales_cols = USE_ROWWISE_SCALING ? DIVUP(cols, SCALE_DIM_X) : cols; + const auto load_scale = [&](const size_t scale_offset_Y, const size_t scale_offset_X) { + const bool in_bounds = scale_offset_Y < scales_rows && scale_offset_X < scales_cols; + size_t scale_idx; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + scale_idx = USE_ROWWISE_SCALING + ? mx_preswizzle_scale_idx(scale_offset_Y, scale_offset_X, mx_swizzle_padded_dim) + : mx_preswizzle_scale_idx(scale_offset_X, scale_offset_Y, mx_swizzle_padded_dim); + } else { + scale_idx = scale_offset_Y * scales_stride + scale_offset_X; + } + const e8m0_t biased_exponent = scales_ptr[in_bounds ? scale_idx : 0]; + return in_bounds ? biased_exponent : static_cast(127); + }; + + if constexpr (USE_ROWWISE_SCALING) { + // Per step, THREADS_PER_CHUNK_X_ROWWISE lanes cover one row of the band. + constexpr size_t ROWS_PER_STEP = TDM_WAVE_SIZE / THREADS_PER_CHUNK_X_ROWWISE; // 4 + constexpr size_t STEPS = TDM_BAND_Y / ROWS_PER_STEP; // 8 + const int lane_offset_Y = lane / THREADS_PER_CHUNK_X_ROWWISE; + const int lane_offset_X = (lane % THREADS_PER_CHUNK_X_ROWWISE) * ELEMS_PER_THREAD; + e8m0_t biased_exponent[STEPS]; +#pragma unroll + for (int step = 0; step < STEPS; step++) { + biased_exponent[step] = load_scale(band_offset_Y + step * ROWS_PER_STEP + lane_offset_Y, + (chunk_offset_X + lane_offset_X) / SCALE_DIM_X); + } + tdm::wait<0>(); +#pragma unroll + for (int step = 0; step < STEPS; step++) { + const int y = step * ROWS_PER_STEP + lane_offset_Y; + const float block_scale = ptx::exp2f(biased_exponent[step]); + Vec in; + Vec out; + in.load_from(&in_sh[wave][y][lane_offset_X]); +#pragma unroll + for (int j = 0; j < ELEMS_PER_THREAD; j++) { + out.data.elt[j] = cvt_from_float(block_scale * static_cast(in.data.elt[j])); + } + out.store_to(&out_sh[wave][y][lane_offset_X]); + } + } else { + // The band is one scale block tall; each lane owns COLS_PER_LANE adjacent columns. + constexpr size_t COLS_PER_LANE = CHUNK_DIM_X / TDM_WAVE_SIZE; // 4 + const int lane_offset_X = lane * COLS_PER_LANE; + e8m0_t biased_exponent[COLS_PER_LANE]; +#pragma unroll + for (int i = 0; i < COLS_PER_LANE; i++) { + biased_exponent[i] = load_scale(band_offset_Y / SCALE_DIM_Y, chunk_offset_X + lane_offset_X + i); + } + tdm::wait<0>(); + float block_scale[COLS_PER_LANE]; +#pragma unroll + for (int i = 0; i < COLS_PER_LANE; i++) { + block_scale[i] = ptx::exp2f(biased_exponent[i]); + } +#pragma unroll + for (int y = 0; y < TDM_BAND_Y; y++) { + Vec in; + Vec out; + in.load_from(&in_sh[wave][y][lane_offset_X]); +#pragma unroll + for (int i = 0; i < COLS_PER_LANE; i++) { + out.data.elt[i] = cvt_from_float(block_scale[i] * static_cast(in.data.elt[i])); + } + out.store_to(&out_sh[wave][y][lane_offset_X]); + } + } + + tdm::store_2d(output_ptr, out_sh[wave], rows, cols, cols, band_offset_Y, chunk_offset_X); + tdm::wait<0>(); +#elif defined(__gfx1250__) + // Compiler without the TDM builtins + dequantize_mxfp8_chunk( + input_ptr, output_ptr, scales_ptr, rows, cols, scales_stride, mx_swizzle_padded_dim, blockIdx.y, + blockIdx.x); +#endif // NVTE_ROCM_HAS_TDM +} + +// Returns false, leaving the work to the generic kernel, where TDM does not apply. The scale +// arguments are the ones of the generic kernel. +inline bool dequantize_tdm(const Tensor &input, const SimpleTensor &input_data, Tensor *output, + const e8m0_t *const scales_ptr, const size_t scales_stride, + const size_t mx_swizzle_padded_dim, hipStream_t stream) { + const auto [rows, cols] = input.flat_2d_dims(); + if (!tdm_kernel::use_tdm(rows, cols, {input_data.dptr, output->data.dptr})) { + return false; + } + const dim3 grid(DIVUP(cols, CHUNK_DIM_X), DIVUP(rows, CHUNK_DIM_Y)); + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + (input.has_columnwise_data() ? 32 : 1), SCALE_DIM_Y, + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + (input.has_data() ? 32 : 1), SCALE_DIM_X, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + input.dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + output->dtype(), OType, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + input.with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, + dequantize_mxfp8_tdm_kernel + <<>>( + reinterpret_cast(input_data.dptr), + reinterpret_cast(output->data.dptr), scales_ptr, rows, cols, + scales_stride, mx_swizzle_padded_dim););)))); // NOLINT(*) + NVTE_CHECK_CUDA(hipGetLastError()); + return true; +} diff --git a/transformer_engine/common/cast/mxfp8/rocm_gated_mxfp8_tdm.cuh b/transformer_engine/common/cast/mxfp8/rocm_gated_mxfp8_tdm.cuh new file mode 100644 index 000000000..32072c74a --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/rocm_gated_mxfp8_tdm.cuh @@ -0,0 +1,168 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ + +#pragma once +// gfx1250 gated MXFP8 quantize with the Tensor Data Mover, included into gated_kernel after +// rocm_gated_mxfp8.cuh for its compute_gated_activation. +//#include "hip/hip_runtime.h" //dummy include to prevent hipification adding this header + +// Each wave moves its own band of a 128x128 chunk of the act (and gate) half global->LDS->global, +// so the waves never synchronize. +template +__global__ void __launch_bounds__(tdm_kernel::THREADS_PER_CHUNK) + quantize_gated_mxfp8_tdm_kernel( + const IType *grad_ptr, const IType *input_act, const IType *input_gate, + OType *output_act_rowwise, OType *output_gate_rowwise, OType *output_act_colwise, + OType *output_gate_colwise, e8m0_t *const scales_rowwise, e8m0_t *const scales_colwise, + const size_t rows, const size_t cols, const size_t scale_stride_rowwise, + const size_t scale_stride_colwise, const ParamOP p) { +#ifdef __gfx1250__ + // These hide the chunk constants of the generic gated kernel. + using tdm_kernel::BAND_Y; + using tdm_kernel::band_t; + using tdm_kernel::CHUNK_DIM_X; + using tdm_kernel::CHUNK_DIM_Y; + using tdm_kernel::COLS_PER_LANE; + using tdm_kernel::WAVE_SIZE; + using tdm_kernel::WAVES; + constexpr bool USE_ROWWISE_SCALING = SCALE_DIM_X > 1; + constexpr bool USE_COLWISE_SCALING = SCALE_DIM_Y > 1; + constexpr size_t HALVES = IS_DGATED ? 2 : 1; + constexpr size_t OUT_BAND_BYTES = sizeof(band_t); + constexpr size_t HALF_OUT_BYTES = (USE_ROWWISE_SCALING + USE_COLWISE_SCALING) * OUT_BAND_BYTES; + constexpr size_t AUX_BYTES = std::max(sizeof(band_t), HALVES * HALF_OUT_BYTES); + static_assert(!USE_COLWISE_SCALING || SCALE_DIM_Y == BAND_Y, + "A band must be one colwise scale block"); + + // act (and for dgated gate) bands, which hold the activations after the first pass + __shared__ alignas(128) band_t in_sh[WAVES][HALVES]; + // The band only the first pass reads (gate, or grad for dgated), then the outputs + __shared__ alignas(128) unsigned char aux_sh[WAVES][AUX_BYTES]; + + const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / WAVE_SIZE); + const int lane_offset_X = (threadIdx.x % WAVE_SIZE) * COLS_PER_LANE; + const size_t band_offset_Y = blockIdx.y * CHUNK_DIM_Y + wave * BAND_Y; + const size_t chunk_offset_X = blockIdx.x * CHUNK_DIM_X; + if (band_offset_Y >= rows) { + return; + } + auto &act = in_sh[wave][0]; + auto &aux_in = *reinterpret_cast *>(aux_sh[wave]); + auto &gate = IS_DGATED ? in_sh[wave][HALVES - 1] : aux_in; + + tdm::load_2d(act, input_act, rows, cols, 2 * cols, band_offset_Y, chunk_offset_X); + tdm::load_2d(gate, input_gate, rows, cols, 2 * cols, band_offset_Y, chunk_offset_X); + if constexpr (IS_DGATED) { + tdm::load_2d(aux_in, grad_ptr, rows, cols, cols, band_offset_Y, chunk_offset_X); + } + tdm::wait<0>(); + + // Out-of-bounds elements are zero and every gated activation keeps them zero. +#pragma unroll + for (int y = 0; y < BAND_Y; y++) { + Vec a, g, d; + a.load_from(&act[y][lane_offset_X]); + g.load_from(&gate[y][lane_offset_X]); + if constexpr (IS_DGATED) { + d.load_from(&aux_in[y][lane_offset_X]); + } +#pragma unroll + for (int i = 0; i < COLS_PER_LANE; i++) { + float grad_elt = 0.0f; + if constexpr (IS_DGATED) { + grad_elt = static_cast(d.data.elt[i]); + } + float result_act, result_gate; + compute_gated_activation( + static_cast(a.data.elt[i]), static_cast(g.data.elt[i]), grad_elt, p, + result_act, result_gate); + a.data.elt[i] = tdm_kernel::cvt_from_float(result_act); + g.data.elt[i] = tdm_kernel::cvt_from_float(result_gate); + } + a.store_to(&act[y][lane_offset_X]); + if constexpr (IS_DGATED) { + g.store_to(&gate[y][lane_offset_X]); + } + } + asm volatile("" ::: "memory"); // the outputs overwrite aux_in through another type + + const size_t out_ld = HALVES * cols; + const size_t scale_cols = DIVUP(cols, SCALE_DIM_X); + // Gate scales follow the act scales of each row (rowwise) or block row (colwise). + auto quantize_half = [&](const band_t &in, OType *out_rowwise, OType *out_colwise, + const size_t scale_offset_rowwise, const size_t scale_offset_colwise, + unsigned char *out_sh) { + if constexpr (USE_ROWWISE_SCALING) { + auto &out = *reinterpret_cast *>(out_sh); + const size_t scale_offset_X = scale_offset_rowwise + chunk_offset_X / SCALE_DIM_X; + tdm_kernel::quantize_band_rowwise( + in, out, scales_rowwise + band_offset_Y * scale_stride_rowwise + scale_offset_X, + scale_stride_rowwise, rows - band_offset_Y, scale_cols - chunk_offset_X / SCALE_DIM_X); + tdm::store_2d(out_rowwise, out, rows, cols, out_ld, band_offset_Y, chunk_offset_X); + } + if constexpr (USE_COLWISE_SCALING) { + auto &out = *reinterpret_cast *>( + out_sh + (USE_ROWWISE_SCALING ? OUT_BAND_BYTES : 0)); + tdm_kernel::quantize_band_colwise( + in, out, + scales_colwise + band_offset_Y / SCALE_DIM_Y * scale_stride_colwise + + scale_offset_colwise + chunk_offset_X, + cols - chunk_offset_X); + tdm::store_2d(out_colwise, out, rows, cols, out_ld, band_offset_Y, chunk_offset_X); + } + }; + quantize_half(act, output_act_rowwise, output_act_colwise, 0, 0, aux_sh[wave]); + if constexpr (IS_DGATED) { + quantize_half(gate, output_gate_rowwise, output_gate_colwise, scale_cols, cols, + aux_sh[wave] + HALF_OUT_BYTES); + } + tdm::wait<0>(); +#endif // __gfx1250__ +} + +// Returns false, leaving the work to the generic kernel, where TDM does not apply. +template +bool quantize_gated_tdm(const Tensor &gated_input, const Tensor &grad, Tensor *output, ParamOP &p, + hipStream_t stream) { + const bool rowwise = output->has_data(); + const bool colwise = output->has_columnwise_data(); + const size_t rows = gated_input.flat_first_dim(); + const size_t cols = gated_input.flat_last_dim() / 2; + if (!tdm_kernel::use_tdm(rows, cols, + {gated_input.data.dptr, IS_BWD ? grad.data.dptr : nullptr, + output->data.dptr, output->columnwise_data.dptr})) { + return false; + } + const dim3 grid(DIVUP(cols, tdm_kernel::CHUNK_DIM_X), DIVUP(rows, tdm_kernel::CHUNK_DIM_Y)); + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + gated_input.dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output->dtype(), OType, + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + (colwise ? 32 : 1), SCALE_DIM_Y, + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + (rowwise ? 32 : 1), SCALE_DIM_X, + // The gate half follows the act half in every row. + const IType *const input = reinterpret_cast(gated_input.data.dptr); + OType *const output_rowwise = reinterpret_cast(output->data.dptr); + OType *const output_colwise = + reinterpret_cast(output->columnwise_data.dptr); + quantize_gated_mxfp8_tdm_kernel + <<>>( + IS_BWD ? reinterpret_cast(grad.data.dptr) : nullptr, input, + input + cols, output_rowwise, rowwise ? output_rowwise + cols : nullptr, + output_colwise, colwise ? output_colwise + cols : nullptr, + reinterpret_cast(output->scale_inv.dptr), + reinterpret_cast(output->columnwise_scale_inv.dptr), rows, cols, + rowwise ? output->scale_inv.shape[1] : 1, + colwise ? output->columnwise_scale_inv.shape[1] : 1, p);)))); // NOLINT(*) + NVTE_CHECK_CUDA(hipGetLastError()); + return true; +} diff --git a/transformer_engine/common/cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh b/transformer_engine/common/cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh new file mode 100644 index 000000000..b4049131c --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh @@ -0,0 +1,669 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ +#pragma once +//#include "hip/hip_runtime.h" //dummy include to prevent hipification adding this header + +// gfx1250 MXFP8 quantize with the Tensor Data Mover: each wave moves its own BAND_Y-row band of a +// 128x128 chunk global->LDS->global and derives both scalings from LDS, so the waves only +// synchronize to combine dbias. The launchers at the end return false where TDM does not apply, +// leaving the work to the generic kernels. Like the other rocm_ headers, this expects its includer +// to have included the common headers. + +#include +#include + +#include "../../util/rocm_device_utils.cuh" +#include "../../util/rocm_tdm.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace tdm_kernel { + +constexpr size_t CHUNK_DIM_Y = 128; +constexpr size_t CHUNK_DIM_X = 128; +constexpr size_t THREADS_PER_CHUNK = 128; +constexpr size_t WAVE_SIZE = 32; +constexpr size_t WAVES = THREADS_PER_CHUNK / WAVE_SIZE; // 4 +constexpr size_t BAND_Y = CHUNK_DIM_Y / WAVES; // 32 +constexpr size_t COLS_PER_LANE = CHUNK_DIM_X / WAVE_SIZE; // 4 +constexpr size_t DBIAS_WAVES = 8; +constexpr size_t MULTI_MAX_TENSORS = 1024; + +template +using band_t = T[BAND_Y][CHUNK_DIM_X]; + +// TDM tiles need 16-byte aligned rows and 32-bit tensor extents. +inline bool use_tdm(const size_t rows, const size_t cols, + const std::initializer_list tensors) { + static const bool enabled = getenv("NVTE_USE_TDM", true); + const auto aligned = [](const void *ptr) { return reinterpret_cast(ptr) % 16 == 0; }; + return enabled && cuda::sm_arch() == 125 && cols % 16 == 0 && rows <= UINT32_MAX && + cols <= UINT32_MAX && std::all_of(tensors.begin(), tensors.end(), aligned); +} + +#ifdef __gfx1250__ +// hip_bfloat16 rounds with per-element branches; __bf16 converts with one instruction. +template +__device__ __forceinline__ T cvt_from_float(const float x) { + if constexpr (std::is_same_v) { + return __builtin_bit_cast(T, static_cast<__bf16>(x)); + } else { + return static_cast(x); + } +} + +// The e8m0 scale rounds up, so scaled finite values stay within the FP8 range and need no clamp: +// this matches static_cast bit for bit, including Inf and NaN, at two values per instruction. +template +__device__ __forceinline__ uint32_t cvt_4x_to_fp8(const float x0, const float x1, const float x2, + const float x3) { + int packed = 0; + if constexpr (std::is_same_v) { + packed = __builtin_amdgcn_cvt_pk_fp8_f32(x0, x1, packed, false); + packed = __builtin_amdgcn_cvt_pk_fp8_f32(x2, x3, packed, true); + } else { + packed = __builtin_amdgcn_cvt_pk_bf8_f32(x0, x1, packed, false); + packed = __builtin_amdgcn_cvt_pk_bf8_f32(x2, x3, packed, true); + } + return static_cast(packed); +} + +// Rowwise MXFP8 of a band in LDS: per step, 8 lanes cover a row with 16 elements each and a lane +// pair shares a scale. Writes the scales of the first valid_rows rows and valid_scales scale +// columns, and returns the lane's amax. +template +__device__ __forceinline__ float quantize_band_rowwise(const band_t &in, band_t &out, + e8m0_t *const scales, + const size_t scale_stride, + const size_t valid_rows, + const size_t valid_scales) { + constexpr size_t ELEMS = 16; + constexpr size_t LANES_PER_ROW = CHUNK_DIM_X / ELEMS; // 8 + constexpr size_t ROWS_PER_STEP = WAVE_SIZE / LANES_PER_ROW; // 4 + constexpr size_t LANES_PER_SCALE = SCALE_DIM_X / ELEMS; // 2 + const int lane = threadIdx.x % WAVE_SIZE; + const int lane_offset_Y = lane / LANES_PER_ROW; + const int lane_offset_X = (lane % LANES_PER_ROW) * ELEMS; + const size_t scale_offset_X = lane_offset_X / SCALE_DIM_X; + float amax = 0; +#pragma unroll + for (int step = 0; step < BAND_Y / ROWS_PER_STEP; step++) { + const size_t y = step * ROWS_PER_STEP + lane_offset_Y; + Vec x; + x.load_from(&in[y][lane_offset_X]); + float thread_amax = 0; +#pragma unroll + for (int j = 0; j < ELEMS; j++) { + thread_amax = fmaxf(thread_amax, fabsf(static_cast(x.data.elt[j]))); + } + amax = fmaxf(amax, thread_amax); + const float scale_amax = rocm_subwarp_allreduce(thread_amax, rocm_op::max{}); + const e8m0_t biased_exponent = + ptx::float_to_e8m0(scale_amax * Quantized_Limits::max_norm_rcp); + if (lane % LANES_PER_SCALE == 0 && y < valid_rows && scale_offset_X < valid_scales) { + scales[y * scale_stride + scale_offset_X] = biased_exponent; + } + const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + uint32_t packed[ELEMS / 4]; +#pragma unroll + for (int p = 0; p < ELEMS / 4; p++) { + packed[p] = cvt_4x_to_fp8(static_cast(x.data.elt[4 * p]) * block_scale_inverse, + static_cast(x.data.elt[4 * p + 1]) * block_scale_inverse, + static_cast(x.data.elt[4 * p + 2]) * block_scale_inverse, + static_cast(x.data.elt[4 * p + 3]) * block_scale_inverse); + } + Vec o; + memcpy(&o.data, packed, sizeof(packed)); + o.store_to(&out[y][lane_offset_X]); + } + return amax; +} + +// Colwise MXFP8 of a band in LDS: each lane owns COLS_PER_LANE adjacent columns, with one pass for +// the amaxes and one to convert. Writes the scales of the first valid_cols columns and returns the +// lane's amax. +template +__device__ __forceinline__ float quantize_band_colwise(const band_t &in, band_t &out, + e8m0_t *const scales, + const size_t valid_cols) { + const int lane_offset_X = (threadIdx.x % WAVE_SIZE) * COLS_PER_LANE; + float col_amax[COLS_PER_LANE] = {}; +#pragma unroll + for (int y = 0; y < BAND_Y; y++) { + Vec x; + x.load_from(&in[y][lane_offset_X]); +#pragma unroll + for (int i = 0; i < COLS_PER_LANE; i++) { + col_amax[i] = fmaxf(col_amax[i], fabsf(static_cast(x.data.elt[i]))); + } + } + float amax = 0; + float block_scale_inverse[COLS_PER_LANE]; +#pragma unroll + for (int i = 0; i < COLS_PER_LANE; i++) { + amax = fmaxf(amax, col_amax[i]); + const e8m0_t biased_exponent = + ptx::float_to_e8m0(col_amax[i] * Quantized_Limits::max_norm_rcp); + if (static_cast(lane_offset_X + i) < valid_cols) { + scales[lane_offset_X + i] = biased_exponent; + } + block_scale_inverse[i] = ptx::exp2f_rcp(biased_exponent); + } +#pragma unroll + for (int y = 0; y < BAND_Y; y++) { + Vec x; + x.load_from(&in[y][lane_offset_X]); + Vec o; + o.data.vec = cvt_4x_to_fp8(static_cast(x.data.elt[0]) * block_scale_inverse[0], + static_cast(x.data.elt[1]) * block_scale_inverse[1], + static_cast(x.data.elt[2]) * block_scale_inverse[2], + static_cast(x.data.elt[3]) * block_scale_inverse[3]); + o.store_to(&out[y][lane_offset_X]); + } + return amax; +} +#endif // __gfx1250__ + +// Quantizes chunk (chunk_Y, chunk_X) of a rows x cols tensor and writes its dbias partials to row +// dbias_row of dbias_workspace. +template +__device__ __forceinline__ void quantize_chunk( + const IType *input_ptr, const IType *act_input_ptr, OType *output_rowwise, + OType *output_colwise, e8m0_t *scales_rowwise, e8m0_t *scales_colwise, + float *dbias_workspace, float *amax_ptr, const size_t rows, const size_t cols, + const size_t scale_stride_rowwise, const size_t scale_stride_colwise, const size_t chunk_Y, + const size_t chunk_X, const size_t dbias_row) { +#ifdef __gfx1250__ + constexpr bool USE_ROWWISE_SCALING = SCALE_DIM_X > 1; + constexpr bool USE_COLWISE_SCALING = SCALE_DIM_Y > 1; + constexpr bool COMPUTE_ACTIVATIONS = IS_ACT || IS_DACT; + constexpr size_t OUT_BAND_BYTES = sizeof(band_t); + constexpr size_t OUT_BYTES = (USE_ROWWISE_SCALING + USE_COLWISE_SCALING) * OUT_BAND_BYTES; + constexpr size_t ACT_BYTES = IS_DACT ? sizeof(band_t) : 0; + constexpr size_t AUX_BYTES = std::max({OUT_BYTES, ACT_BYTES, size_t{1}}); + static_assert(!USE_COLWISE_SCALING || SCALE_DIM_Y == BAND_Y, + "A band must be one colwise scale block"); + + __shared__ alignas(128) band_t in_sh[WAVES]; + // Stages the activation input of DACT, then holds the outputs once in_sh has the activations. + __shared__ alignas(128) unsigned char aux_sh[WAVES][AUX_BYTES]; + + const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / WAVE_SIZE); + const int lane = threadIdx.x % WAVE_SIZE; + const int lane_offset_X = lane * COLS_PER_LANE; + const size_t band_offset_Y = chunk_Y * CHUNK_DIM_Y + wave * BAND_Y; + const size_t chunk_offset_X = chunk_X * CHUNK_DIM_X; + auto &in = in_sh[wave]; + auto &act_in = *reinterpret_cast *>(aux_sh[wave]); + auto &out_rowwise = *reinterpret_cast *>(aux_sh[wave]); + auto &out_colwise = + *reinterpret_cast *>(aux_sh[wave] + (USE_ROWWISE_SCALING ? OUT_BAND_BYTES : 0)); + // Out-of-bounds elements are zero in LDS and every activation (or derivative times a zero + // gradient) keeps them zero, so they leave the amaxes and dbias unchanged. + float amax = 0; + float partial_dbias[COLS_PER_LANE] = {}; + + if (band_offset_Y < rows) { + tdm::load_2d(in, input_ptr, rows, cols, cols, band_offset_Y, chunk_offset_X); + if constexpr (IS_DACT) { + tdm::load_2d(act_in, act_input_ptr, rows, cols, cols, band_offset_Y, chunk_offset_X); + } + tdm::wait<0>(); + + if constexpr (COMPUTE_ACTIVATIONS || IS_DBIAS) { + // Activations in place; a lane owns COLS_PER_LANE columns, so it also sums their dbias. +#pragma unroll + for (int y = 0; y < BAND_Y; y++) { + Vec x; + Vec a; + x.load_from(&in[y][lane_offset_X]); + if constexpr (IS_DACT) { + a.load_from(&act_in[y][lane_offset_X]); + } +#pragma unroll + for (int i = 0; i < COLS_PER_LANE; i++) { + float elt = static_cast(x.data.elt[i]); + if constexpr (IS_ACT) { + elt = OP(elt, {}); + } + if constexpr (IS_DACT) { + elt *= OP(static_cast(a.data.elt[i]), {}); + } + if constexpr (IS_DBIAS) { + partial_dbias[i] += elt; + } + if constexpr (COMPUTE_ACTIVATIONS) { + x.data.elt[i] = cvt_from_float(elt); + } + } + if constexpr (COMPUTE_ACTIVATIONS) { + x.store_to(&in[y][lane_offset_X]); + } + } + if constexpr (IS_DACT) { + asm volatile("" ::: "memory"); // the outputs overwrite act_in through another type + } + } + + if constexpr (USE_ROWWISE_SCALING) { + const size_t scale_offset_X = chunk_offset_X / SCALE_DIM_X; + // Like the 128-wide chunks of the generic kernel, this also fills the scale padding. + amax = fmaxf(amax, quantize_band_rowwise( + in, out_rowwise, + scales_rowwise + band_offset_Y * scale_stride_rowwise + scale_offset_X, + scale_stride_rowwise, rows - band_offset_Y, + scale_stride_rowwise - scale_offset_X)); + tdm::store_2d(output_rowwise, out_rowwise, rows, cols, cols, band_offset_Y, chunk_offset_X); + } + if constexpr (USE_COLWISE_SCALING) { + amax = fmaxf(amax, quantize_band_colwise( + in, out_colwise, + scales_colwise + band_offset_Y / SCALE_DIM_Y * scale_stride_colwise + + chunk_offset_X, + cols - chunk_offset_X)); + tdm::store_2d(output_colwise, out_colwise, rows, cols, cols, band_offset_Y, chunk_offset_X); + } + } + + if constexpr (IS_DBIAS) { + // Every wave leaves its column sums in its consumed in_sh band. + asm volatile("" ::: "memory"); // in_sh is reused as float + Vec partial; +#pragma unroll + for (int i = 0; i < COLS_PER_LANE; i++) { + partial.data.elt[i] = partial_dbias[i]; + } + partial.store_to(reinterpret_cast(in_sh[wave]) + lane_offset_X); + __syncthreads(); + const size_t col = chunk_offset_X + threadIdx.x; + if (col < cols) { + float dbias = 0; +#pragma unroll + for (int w = 0; w < WAVES; w++) { + dbias += reinterpret_cast(in_sh[w])[threadIdx.x]; + } + dbias_workspace[dbias_row * cols + col] = dbias; + } + } + + if (amax_ptr != nullptr) { + const float wave_amax = rocm_subwarp_allreduce(amax, rocm_op::max{}); + // amax only grows, so a stale read can only cause a redundant atomic. + if (lane == 0 && wave_amax > *amax_ptr) { + atomicMaxFloat(amax_ptr, wave_amax); + } + } + tdm::wait<0>(); +#endif // __gfx1250__ +} + +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + quantize_mxfp8_tdm_kernel(const IType *input, const IType *act_input, OType *output_rowwise, + OType *output_colwise, e8m0_t *scales_rowwise, + e8m0_t *scales_colwise, const float *noop, float *dbias_workspace, + float *amax, const size_t rows, const size_t cols, + const size_t scale_stride_rowwise, + const size_t scale_stride_colwise) { + if constexpr (!IS_DBIAS && !IS_DACT && !IS_ACT) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + } + quantize_chunk( + input, act_input, output_rowwise, output_colwise, scales_rowwise, scales_colwise, + dbias_workspace, amax, rows, cols, scale_stride_rowwise, scale_stride_colwise, blockIdx.y, + blockIdx.x, blockIdx.y); +} + +// The tensors share the last dimension and each has a multiple of CHUNK_DIM_Y rows, so no chunk +// straddles two of them. The dbias partials of chunk row blockIdx.y go to that workspace row. +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + grouped_quantize_mxfp8_tdm_kernel(const IType *input, const IType *act_input, + OType *output_rowwise, OType *output_colwise, + e8m0_t *scales_rowwise, e8m0_t *scales_colwise, + const float *noop, float *dbias_workspace, float *amax, + const size_t first_logical_dim, const size_t cols, + const size_t num_tensors, const int64_t *offsets_ptr, + const int64_t *first_dims_ptr) { + if constexpr (!IS_DBIAS && !IS_DACT && !IS_ACT) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + } + const size_t chunk_offset_Y = static_cast(blockIdx.y) * CHUNK_DIM_Y; + size_t tensor_id; + size_t tensor_offset; + if (offsets_ptr == nullptr) { + const size_t rows_per_tensor = first_logical_dim / num_tensors; + tensor_id = chunk_offset_Y / rows_per_tensor; + tensor_offset = tensor_id * rows_per_tensor * cols; + } else { + // The last tensor starting at or before the chunk + size_t lo = 1; + size_t hi = num_tensors; + while (lo < hi) { + const size_t mid = lo + (hi - lo) / 2; + if (static_cast(offsets_ptr[mid]) <= chunk_offset_Y * cols) { + lo = mid + 1; + } else { + hi = mid; + } + } + tensor_id = lo - 1; + tensor_offset = static_cast(offsets_ptr[tensor_id]); + } + const size_t rows = first_dims_ptr != nullptr ? static_cast(first_dims_ptr[tensor_id]) + : first_logical_dim / num_tensors; + const size_t tensor_offset_Y = tensor_offset / cols; + const size_t chunk_Y = blockIdx.y - tensor_offset_Y / CHUNK_DIM_Y; + if (chunk_Y * CHUNK_DIM_Y >= rows) { + return; + } + constexpr bool ROWWISE = SCALE_DIM_X > 1; + constexpr bool COLWISE = SCALE_DIM_Y > 1; + const size_t scale_stride_rowwise = DIVUP(cols, SCALE_DIM_X); + quantize_chunk( + input + tensor_offset, IS_DACT ? act_input + tensor_offset : nullptr, + ROWWISE ? output_rowwise + tensor_offset : nullptr, + COLWISE ? output_colwise + tensor_offset : nullptr, + ROWWISE ? scales_rowwise + tensor_offset_Y * scale_stride_rowwise : nullptr, + COLWISE ? scales_colwise + tensor_offset_Y / SCALE_DIM_Y * cols : nullptr, dbias_workspace, + amax, rows, cols, scale_stride_rowwise, cols, chunk_Y, blockIdx.x, blockIdx.y); +} + +struct MultiQuantizeArgs { + const void *input[MULTI_MAX_TENSORS]; + void *output_rowwise[MULTI_MAX_TENSORS]; + void *output_colwise[MULTI_MAX_TENSORS]; + void *scales_rowwise[MULTI_MAX_TENSORS]; + void *scales_colwise[MULTI_MAX_TENSORS]; + float *amax[MULTI_MAX_TENSORS]; + uint32_t rows[MULTI_MAX_TENSORS]; + uint32_t cols[MULTI_MAX_TENSORS]; + // Tensor t owns chunk rows [first_chunk_Y[t], first_chunk_Y[t + 1]). + uint32_t first_chunk_Y[MULTI_MAX_TENSORS + 1]; + int num_tensors; +}; + +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + multi_quantize_mxfp8_tdm_kernel(const MultiQuantizeArgs args) { + int t = 0; + int hi = args.num_tensors - 1; + while (t < hi) { + const int mid = t + (hi - t) / 2; + if (args.first_chunk_Y[mid + 1] <= blockIdx.y) { + t = mid + 1; + } else { + hi = mid; + } + } + const size_t rows = args.rows[t]; + const size_t cols = args.cols[t]; + const size_t chunk_Y = blockIdx.y - args.first_chunk_Y[t]; + if (chunk_Y * CHUNK_DIM_Y >= rows || blockIdx.x * CHUNK_DIM_X >= cols) { + return; + } + quantize_chunk( + static_cast(args.input[t]), nullptr, + static_cast(args.output_rowwise[t]), static_cast(args.output_colwise[t]), + static_cast(args.scales_rowwise[t]), static_cast(args.scales_colwise[t]), + nullptr, args.amax[t], rows, cols, DIVUP(cols, SCALE_DIM_X), cols, chunk_Y, blockIdx.x, 0); +} + +// Sums each tensor's dbias partials, one row per chunk, into its row of dbias. A block covers +// DBIAS_WAVES / waves_per_slab 128-column slabs, whose waves split the BAND_Y-row bands. +template +__global__ void __launch_bounds__(DBIAS_WAVES * WAVE_SIZE) + reduce_dbias_tdm_kernel(OType *const dbias, const float *const partial, const size_t cols, + const size_t partial_rows, const int64_t *const offsets_ptr, + const int64_t *const first_dims_ptr, const int waves_per_slab) { +#ifdef __gfx1250__ + __shared__ alignas(128) band_t band_sh[DBIAS_WAVES]; + + const size_t tensor_id = blockIdx.y; + size_t tensor_rows = partial_rows / gridDim.y; + size_t tensor_offset_Y = tensor_id * tensor_rows; + if (offsets_ptr != nullptr) { + tensor_rows = static_cast(first_dims_ptr[tensor_id]) / CHUNK_DIM_Y; + tensor_offset_Y = static_cast(offsets_ptr[tensor_id]) / cols / CHUNK_DIM_Y; + } + const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / WAVE_SIZE); + const size_t block_cols = DBIAS_WAVES / waves_per_slab * CHUNK_DIM_X; + const size_t block_offset_X = blockIdx.x * block_cols; + const size_t slab_offset_X = block_offset_X + wave / waves_per_slab * CHUNK_DIM_X; + const int lane_offset_X = (threadIdx.x % WAVE_SIZE) * COLS_PER_LANE; + const float *const tensor_partial = partial + tensor_offset_Y * cols; + auto &band = band_sh[wave]; + + Vec sum; + sum.clear(); + if (slab_offset_X < cols) { + for (size_t band_offset_Y = wave % waves_per_slab * BAND_Y; band_offset_Y < tensor_rows; + band_offset_Y += waves_per_slab * BAND_Y) { + tdm::load_2d(band, tensor_partial, tensor_rows, cols, cols, band_offset_Y, slab_offset_X); + tdm::wait<0>(); +#pragma unroll + for (int y = 0; y < BAND_Y; y++) { + Vec x; + x.load_from(&band[y][lane_offset_X]); +#pragma unroll + for (int i = 0; i < COLS_PER_LANE; i++) { + sum.data.elt[i] += x.data.elt[i]; + } + } + } + } + sum.store_to(&band[0][lane_offset_X]); + __syncthreads(); + for (size_t x = threadIdx.x; x < block_cols; x += blockDim.x) { + const size_t col = block_offset_X + x; + if (col < cols) { + const int first_wave = x / CHUNK_DIM_X * waves_per_slab; + float dbias_sum = 0; + for (int w = first_wave; w < first_wave + waves_per_slab; w++) { + dbias_sum += band_sh[w][0][x % CHUNK_DIM_X]; + } + dbias[tensor_id * cols + col] = static_cast(dbias_sum); + } + } +#endif // __gfx1250__ +} + +// partial_rows counts the partials of all num_tensors tensors. Without offsets, the tensors share +// them equally; otherwise tensor t has first_dims[t] / CHUNK_DIM_Y of them. +template +void reduce_dbias(OType *dbias, const float *partial, const size_t cols, const size_t num_tensors, + const size_t partial_rows, const int64_t *offsets_ptr, + const int64_t *first_dims_ptr, hipStream_t stream) { + // About one band per wave; short tensors leave the spare waves to more slabs. + const size_t bands = DIVUP(partial_rows / num_tensors, BAND_Y); + size_t waves_per_slab = 1; + while (waves_per_slab < DBIAS_WAVES && waves_per_slab < bands) { + waves_per_slab *= 2; + } + const dim3 grid(DIVUP(cols, DBIAS_WAVES / waves_per_slab * CHUNK_DIM_X), num_tensors); + reduce_dbias_tdm_kernel<<>>( + dbias, partial, cols, partial_rows, offsets_ptr, first_dims_ptr, + static_cast(waves_per_slab)); + NVTE_CHECK_CUDA(hipGetLastError()); +} + +} // namespace tdm_kernel + +template +bool quantize_tdm(const Tensor &input, const Tensor *act_input, const Tensor *noop, Tensor *output, + Tensor *dbias, Tensor *workspace, hipStream_t stream) { + using namespace tdm_kernel; + const bool rowwise = output->has_data(); + const bool colwise = output->has_columnwise_data(); + const auto [rows, cols] = input.flat_2d_dims(); + // The dbias workspace query takes the generic path, whose workspace has at least as many rows. + if (!(rowwise || colwise) || + (IS_DBIAS && (workspace == nullptr || workspace->data.dptr == nullptr)) || + !use_tdm(rows, cols, + {input.data.dptr, IS_DACT ? act_input->data.dptr : nullptr, output->data.dptr, + output->columnwise_data.dptr, IS_DBIAS ? workspace->data.dptr : nullptr})) { + return false; + } + if constexpr (IS_DBIAS) { + NVTE_CHECK(dbias->data.dtype == input.dtype(), "DBias must have the same type as input."); + NVTE_CHECK(dbias->data.shape == std::vector{cols}, "Wrong shape of DBias."); + } + float *const workspace_ptr = IS_DBIAS ? reinterpret_cast(workspace->data.dptr) : nullptr; + const dim3 grid(DIVUP(cols, CHUNK_DIM_X), DIVUP(rows, CHUNK_DIM_Y)); + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input.dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output->dtype(), OType, + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + (colwise ? 32 : 1), SCALE_DIM_Y, + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + (rowwise ? 32 : 1), SCALE_DIM_X, + quantize_mxfp8_tdm_kernel + <<>>( + reinterpret_cast(input.data.dptr), + IS_DACT ? reinterpret_cast(act_input->data.dptr) : nullptr, + reinterpret_cast(output->data.dptr), + reinterpret_cast(output->columnwise_data.dptr), + reinterpret_cast(output->scale_inv.dptr), + reinterpret_cast(output->columnwise_scale_inv.dptr), + reinterpret_cast(noop->data.dptr), workspace_ptr, + reinterpret_cast(output->amax.dptr), rows, cols, + rowwise ? output->scale_inv.shape[1] : 1, + colwise ? output->columnwise_scale_inv.shape[1] : 1);)))); // NOLINT(*) + NVTE_CHECK_CUDA(hipGetLastError()); + if constexpr (IS_DBIAS) { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input.dtype(), IType, + reduce_dbias(reinterpret_cast(dbias->data.dptr), workspace_ptr, cols, 1, + grid.y, nullptr, nullptr, stream);); + } + return true; +} + +// Called after the dbias workspace query: the generic chunks it sized for are not taller than TDM's. +template +bool group_quantize_tdm(const GroupedTensor *input, const GroupedTensor *activations, + const Tensor *noop, GroupedTensor *output, GroupedTensor *dbias, + Tensor *workspace, hipStream_t stream) { + using namespace tdm_kernel; + const bool rowwise = output->has_data(); + const bool colwise = output->has_columnwise_data(); + const size_t first_logical_dim = input->logical_shape.data[0]; + const size_t cols = input->logical_shape.data[1]; + const size_t num_tensors = input->num_tensors; + const bool same_shape = output->all_same_shape(); + // Per-tensor rows of a varying first dim are device-side and, as upstream requires, 128-aligned. + const size_t tensor_rows = same_shape ? first_logical_dim / num_tensors : first_logical_dim; + if (!output->all_same_last_dim() || tensor_rows % CHUNK_DIM_Y != 0 || + !use_tdm(first_logical_dim, cols, + {input->data.dptr, IS_DACT ? activations->data.dptr : nullptr, + rowwise ? output->data.dptr : nullptr, + colwise ? output->columnwise_data.dptr : nullptr, + IS_DBIAS ? workspace->data.dptr : nullptr})) { + return false; + } + const int64_t *const offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr); + const int64_t *const first_dims_ptr = reinterpret_cast(output->first_dims.dptr); + float *const workspace_ptr = IS_DBIAS ? reinterpret_cast(workspace->data.dptr) : nullptr; + const dim3 grid(DIVUP(cols, CHUNK_DIM_X), DIVUP(first_logical_dim, CHUNK_DIM_Y)); + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input->dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output->dtype(), OType, + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + (colwise ? 32 : 1), SCALE_DIM_Y, + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + (rowwise ? 32 : 1), SCALE_DIM_X, + grouped_quantize_mxfp8_tdm_kernel + <<>>( + reinterpret_cast(input->data.dptr), + IS_DACT ? reinterpret_cast(activations->data.dptr) : nullptr, + reinterpret_cast(output->data.dptr), + reinterpret_cast(output->columnwise_data.dptr), + reinterpret_cast(output->scale_inv.dptr), + reinterpret_cast(output->columnwise_scale_inv.dptr), + reinterpret_cast(noop->data.dptr), workspace_ptr, + reinterpret_cast(output->amax.dptr), first_logical_dim, cols, + num_tensors, offsets_ptr, first_dims_ptr);)))); // NOLINT(*) + NVTE_CHECK_CUDA(hipGetLastError()); + if constexpr (IS_DBIAS) { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input->dtype(), IType, + reduce_dbias(reinterpret_cast(dbias->data.dptr), workspace_ptr, cols, + num_tensors, grid.y, same_shape ? nullptr : offsets_ptr, + same_shape ? nullptr : first_dims_ptr, stream);); + } + return true; +} + +// TDM only if every tensor qualifies, as one launch covers them all. +inline bool multi_quantize_tdm(const std::vector &input_list, + std::vector &output_list, hipStream_t stream) { + using namespace tdm_kernel; + const size_t num_tensors = input_list.size(); + const bool rowwise = output_list[0]->has_data(); + const bool colwise = output_list[0]->has_columnwise_data(); + if (num_tensors > MULTI_MAX_TENSORS) { + return false; + } + for (size_t i = 0; i < num_tensors; i++) { + if (!use_tdm(input_list[i]->flat_first_dim(), input_list[i]->flat_last_dim(), + {input_list[i]->data.dptr, rowwise ? output_list[i]->data.dptr : nullptr, + colwise ? output_list[i]->columnwise_data.dptr : nullptr})) { + return false; + } + } + MultiQuantizeArgs args; + args.num_tensors = num_tensors; + args.first_chunk_Y[0] = 0; + size_t chunks_X = 0; + for (size_t i = 0; i < num_tensors; i++) { + const size_t rows = input_list[i]->flat_first_dim(); + const size_t cols = input_list[i]->flat_last_dim(); + args.input[i] = input_list[i]->data.dptr; + args.output_rowwise[i] = rowwise ? output_list[i]->data.dptr : nullptr; + args.output_colwise[i] = colwise ? output_list[i]->columnwise_data.dptr : nullptr; + args.scales_rowwise[i] = rowwise ? output_list[i]->scale_inv.dptr : nullptr; + args.scales_colwise[i] = colwise ? output_list[i]->columnwise_scale_inv.dptr : nullptr; + args.amax[i] = reinterpret_cast(output_list[i]->amax.dptr); + args.rows[i] = rows; + args.cols[i] = cols; + args.first_chunk_Y[i + 1] = args.first_chunk_Y[i] + DIVUP(rows, CHUNK_DIM_Y); + chunks_X = std::max(chunks_X, DIVUP(cols, CHUNK_DIM_X)); + } + const dim3 grid(chunks_X, args.first_chunk_Y[num_tensors]); + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input_list[0]->data.dtype, IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output_list[0]->dtype(), OType, + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + (colwise ? 32 : 1), SCALE_DIM_Y, + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + (rowwise ? 32 : 1), SCALE_DIM_X, + multi_quantize_mxfp8_tdm_kernel + <<>>(args);)))); // NOLINT(*) + NVTE_CHECK_CUDA(hipGetLastError()); + return true; +} + +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine diff --git a/transformer_engine/common/util/rocm_tdm.cuh b/transformer_engine/common/util/rocm_tdm.cuh index 2cec11178..c40f90a4e 100644 --- a/transformer_engine/common/util/rocm_tdm.cuh +++ b/transformer_engine/common/util/rocm_tdm.cuh @@ -8,10 +8,12 @@ // gfx1250 Tensor Data Mover (TDM): asynchronous 2D tile copies between global memory and LDS. // Descriptors are packed here rather than with hip/amd_detail/amd_gfx1250_TDM.h, whose stride -// setters truncate (ROCm/rocm-systems#13274). +// setters truncate (ROCm/rocm-systems#13274). With a compiler that lacks the TDM builtins, the +// calling wave copies the tile itself. #include #include +#include #if defined(__gfx1250__) && __has_builtin(__builtin_amdgcn_tensor_load_to_lds) && \ __has_builtin(__builtin_amdgcn_tensor_store_from_lds) && \ @@ -19,11 +21,12 @@ #define NVTE_ROCM_HAS_TDM 1 #endif -#ifdef NVTE_ROCM_HAS_TDM +#ifdef __gfx1250__ namespace transformer_engine { namespace tdm { namespace detail { +#ifdef NVTE_ROCM_HAS_TDM // The builtins' operand signedness has changed between compiler versions; integer vectors of // either signedness convert implicitly. typedef uint32_t u32x4 __attribute__((ext_vector_type(4))); @@ -81,6 +84,35 @@ __device__ __forceinline__ void copy_2d(const T (&lds)[TILE_ROWS][TILE_COLS], co __builtin_amdgcn_tensor_load_to_lds(g0, g1, z4, z4, z8, 0); } } +#else +template +__device__ __forceinline__ void copy_2d(const T (&lds)[TILE_ROWS][TILE_COLS], const T *gmem, + uint32_t rows, uint32_t cols, uint64_t ld, uint32_t row, + uint32_t col) { + static_assert(sizeof(T) == 1 || sizeof(T) == 2 || sizeof(T) == 4 || sizeof(T) == 8, + "TDM elements are 1, 2, 4 or 8 bytes"); + using U = std::conditional_t< + sizeof(T) == 1, uint8_t, + std::conditional_t>>; + U *const tile = reinterpret_cast(const_cast(&lds[0][0])); + U *const tensor = reinterpret_cast(const_cast(gmem)); + // The tile is accessed as U here and as T by the caller. + asm volatile("" ::: "memory"); + for (uint32_t i = __lane_id(); i < TILE_ROWS * TILE_COLS; i += warpSize) { + const uint64_t r = row + static_cast(i / TILE_COLS); + const uint64_t c = col + static_cast(i % TILE_COLS); + const bool inside = r < rows && c < cols; + if constexpr (STORE) { + if (inside) { + tensor[r * ld + c] = tile[i]; + } + } else { + tile[i] = inside ? tensor[r * ld + c] : U{0}; + } + } +} +#endif // NVTE_ROCM_HAS_TDM } // namespace detail @@ -107,10 +139,12 @@ __device__ __forceinline__ void store_2d(T *gmem, const T (&lds)[TILE_ROWS][TILE // the loaded data only after a subsequent barrier. template __device__ __forceinline__ void wait() { +#ifdef NVTE_ROCM_HAS_TDM __builtin_amdgcn_s_wait_tensorcnt(N); +#endif asm volatile("" ::: "memory"); // the compiler does not see the TDM's LDS writes } } // namespace tdm } // namespace transformer_engine -#endif // NVTE_ROCM_HAS_TDM +#endif // __gfx1250__ From deaafd8296224103086eb5c8404f9f0439163ae2 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Thu, 22 Oct 2026 23:12:12 +0000 Subject: [PATCH 3/8] add fp8 --- .../common/cast/core/rocm_tdm_common.cuh | 202 ++++++++ .../common/cast/dispatch/quantize.cuh | 15 + .../common/cast/fp8/rocm_quantize_fp8_tdm.cuh | 431 ++++++++++++++++++ .../cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh | 107 +---- .../common/recipe/current_scaling.cu | 10 + .../common/transpose/multi_cast_transpose.cu | 9 + 6 files changed, 678 insertions(+), 96 deletions(-) create mode 100644 transformer_engine/common/cast/core/rocm_tdm_common.cuh create mode 100644 transformer_engine/common/cast/fp8/rocm_quantize_fp8_tdm.cuh diff --git a/transformer_engine/common/cast/core/rocm_tdm_common.cuh b/transformer_engine/common/cast/core/rocm_tdm_common.cuh new file mode 100644 index 000000000..a7410e031 --- /dev/null +++ b/transformer_engine/common/cast/core/rocm_tdm_common.cuh @@ -0,0 +1,202 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ +#pragma once +//#include "hip/hip_runtime.h" //dummy include to prevent hipification adding this header + +// gfx1250 Tensor Data Mover pieces shared by the TDM quantize paths: the check whether TDM applies, +// the reduction of per-chunk dbias partials, and amax. This expects its includer to have included +// the common headers, util/system.h and util/cuda_runtime.h. + +#include +#include + +#include "../../util/rocm_device_utils.cuh" +#include "../../util/rocm_tdm.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace common { + +// TDM tiles need 16-byte aligned rows and 32-bit tensor extents. +inline bool use_tdm(const size_t rows, const size_t cols, + const std::initializer_list tensors) { + static const bool enabled = getenv("NVTE_USE_TDM", true); + const auto aligned = [](const void *ptr) { return reinterpret_cast(ptr) % 16 == 0; }; + return enabled && cuda::sm_arch() == 125 && cols % 16 == 0 && rows <= UINT32_MAX && + cols <= UINT32_MAX && std::all_of(tensors.begin(), tensors.end(), aligned); +} + +namespace tdm_dbias_kernel { + +constexpr size_t CHUNK_DIM_Y = 128; // input rows per partial row +constexpr size_t SLAB_DIM_X = 128; +constexpr size_t BAND_Y = 32; +constexpr size_t WAVE_SIZE = 32; +constexpr size_t WAVES = 8; +constexpr size_t COLS_PER_LANE = SLAB_DIM_X / WAVE_SIZE; // 4 + +// Sums each tensor's dbias partials, one row per chunk, into its row of dbias. A block covers +// WAVES / waves_per_slab 128-column slabs, whose waves split the BAND_Y-row bands. +template +__global__ void __launch_bounds__(WAVES * WAVE_SIZE) + reduce_dbias_tdm_kernel(OType *const dbias, const float *const partial, const size_t cols, + const size_t partial_rows, const int64_t *const offsets_ptr, + const int64_t *const first_dims_ptr, const int waves_per_slab) { +#ifdef __gfx1250__ + __shared__ alignas(128) float band_sh[WAVES][BAND_Y][SLAB_DIM_X]; + + const size_t tensor_id = blockIdx.y; + size_t tensor_rows = partial_rows / gridDim.y; + size_t tensor_offset_Y = tensor_id * tensor_rows; + if (offsets_ptr != nullptr) { + tensor_rows = static_cast(first_dims_ptr[tensor_id]) / CHUNK_DIM_Y; + tensor_offset_Y = static_cast(offsets_ptr[tensor_id]) / cols / CHUNK_DIM_Y; + } + const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / WAVE_SIZE); + const size_t block_cols = WAVES / waves_per_slab * SLAB_DIM_X; + const size_t block_offset_X = blockIdx.x * block_cols; + const size_t slab_offset_X = block_offset_X + wave / waves_per_slab * SLAB_DIM_X; + const int lane_offset_X = (threadIdx.x % WAVE_SIZE) * COLS_PER_LANE; + const float *const tensor_partial = partial + tensor_offset_Y * cols; + auto &band = band_sh[wave]; + + Vec sum; + sum.clear(); + if (slab_offset_X < cols) { + for (size_t band_offset_Y = wave % waves_per_slab * BAND_Y; band_offset_Y < tensor_rows; + band_offset_Y += waves_per_slab * BAND_Y) { + tdm::load_2d(band, tensor_partial, tensor_rows, cols, cols, band_offset_Y, slab_offset_X); + tdm::wait<0>(); +#pragma unroll + for (int y = 0; y < BAND_Y; y++) { + Vec x; + x.load_from(&band[y][lane_offset_X]); +#pragma unroll + for (int i = 0; i < COLS_PER_LANE; i++) { + sum.data.elt[i] += x.data.elt[i]; + } + } + } + } + sum.store_to(&band[0][lane_offset_X]); + __syncthreads(); + for (size_t x = threadIdx.x; x < block_cols; x += blockDim.x) { + const size_t col = block_offset_X + x; + if (col < cols) { + const int first_wave = x / SLAB_DIM_X * waves_per_slab; + float dbias_sum = 0; + for (int w = first_wave; w < first_wave + waves_per_slab; w++) { + dbias_sum += band_sh[w][0][x % SLAB_DIM_X]; + } + dbias[tensor_id * cols + col] = static_cast(dbias_sum); + } + } +#endif // __gfx1250__ +} + +} // namespace tdm_dbias_kernel + +// Reduces the dbias partials of TDM chunks, one row per CHUNK_DIM_Y input rows. partial_rows counts +// the partials of all num_tensors tensors. Without offsets, the tensors share them equally; +// otherwise tensor t has first_dims[t] / CHUNK_DIM_Y of them. +template +void reduce_dbias_tdm(OType *dbias, const float *partial, const size_t cols, + const size_t num_tensors, const size_t partial_rows, + const int64_t *offsets_ptr, const int64_t *first_dims_ptr, + hipStream_t stream) { + using namespace tdm_dbias_kernel; + // About one band per wave; short tensors leave the spare waves to more slabs. + const size_t bands = DIVUP(partial_rows / num_tensors, BAND_Y); + size_t waves_per_slab = 1; + while (waves_per_slab < WAVES && waves_per_slab < bands) { + waves_per_slab *= 2; + } + const dim3 grid(DIVUP(cols, WAVES / waves_per_slab * SLAB_DIM_X), num_tensors); + reduce_dbias_tdm_kernel<<>>( + dbias, partial, cols, partial_rows, offsets_ptr, first_dims_ptr, + static_cast(waves_per_slab)); + NVTE_CHECK_CUDA(hipGetLastError()); +} + +namespace tdm_amax_kernel { + +constexpr size_t TILE_BYTES = 16384; +constexpr size_t WAVE_SIZE = 32; +constexpr size_t WAVES = 8; + +// Max |x| over the n elements of input, which the waves stream a row tile at a time. Atomics on +// one address serialize, so a block reduces its waves' amaxes first and grid has a block per CU. +template +__global__ void __launch_bounds__(WAVES * WAVE_SIZE) + amax_tdm_kernel(const IType *input, float *amax_ptr, const size_t n, const float *noop) { +#ifdef __gfx1250__ + constexpr size_t TILE = TILE_BYTES / sizeof(IType); + constexpr size_t VEC = 16 / sizeof(IType); + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + __shared__ alignas(128) IType tile_sh[WAVES][1][TILE]; + __shared__ float amax_sh[WAVES]; + + const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / WAVE_SIZE); + const int lane = threadIdx.x % WAVE_SIZE; + auto &tile = tile_sh[wave]; + const size_t step = static_cast(gridDim.x) * WAVES * TILE; + // The single-row tiles never step by the pitch. + const size_t ld = DIVUP(n, size_t{16}) * 16; + float amax = 0; + for (size_t offset = (blockIdx.x * WAVES + wave) * TILE; offset < n; offset += step) { + tdm::load_2d(tile, input, 1, n, ld, 0, offset); + tdm::wait<0>(); +#pragma unroll 4 + for (size_t i = lane * VEC; i < TILE; i += WAVE_SIZE * VEC) { + Vec x; + x.load_from(&tile[0][i]); +#pragma unroll + for (int j = 0; j < VEC; j++) { + __builtin_assume(amax >= 0); + amax = fmaxf(fabsf(static_cast(x.data.elt[j])), amax); + } + } + } + const float wave_amax = rocm_subwarp_allreduce(amax, rocm_op::max{}); + if (lane == 0) { + amax_sh[wave] = wave_amax; + } + __syncthreads(); + if (threadIdx.x == 0) { + float block_amax = 0; +#pragma unroll + for (int w = 0; w < WAVES; w++) { + block_amax = fmaxf(block_amax, amax_sh[w]); + } + rocm_atomicMaxFloat(amax_ptr, block_amax); + } +#endif // __gfx1250__ +} + +} // namespace tdm_amax_kernel + +// Raises *amax to the max |x| of input's n elements, so the caller zeroes it first. +template +bool compute_amax_tdm(const IType *input, float *amax, const size_t n, const float *noop, + hipStream_t stream) { + using namespace tdm_amax_kernel; + // A single row has no row alignment to keep. + if (n > UINT32_MAX || !use_tdm(1, 0, {input})) { + return false; + } + constexpr size_t TILE = TILE_BYTES / sizeof(IType); + const size_t blocks = + std::min(DIVUP(n, TILE * WAVES), static_cast(cuda::sm_count())); + amax_tdm_kernel<<>>(input, amax, n, noop); + NVTE_CHECK_CUDA(hipGetLastError()); + return true; +} + +} // namespace common +} // namespace dispatch +} // namespace transformer_engine diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 81676eed8..e53679352 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -35,6 +35,7 @@ #include "../mxfp8/group_scaled_swiglu_mxfp8.cuh" #include "../mxfp8/quantize_mxfp8.cuh" #ifdef __HIP_PLATFORM_AMD__ +#include "../fp8/rocm_quantize_fp8_tdm.cuh" #include "../mxfp4/quantize_mxfp4.cuh" #endif //#ifdef __HIP_PLATFORM_AMD__ // The optimized NVFP4 kernels are Blackwell-only (CUtensorMap/tcgen05); ROCm uses @@ -87,6 +88,13 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, const Tensor *dummy_input_tensor = nullptr; Tensor *dummy_dbias_tensor = nullptr; Tensor *dummy_workspace_tensor = nullptr; +#ifdef __HIP_PLATFORM_AMD__ + if (fp8::quantize_tdm( + *input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, + dummy_workspace_tensor, stream)) { + break; + } +#endif if (output_tensor->has_columnwise_data()) { NVTE_CHECK(output_tensor->has_data(), "Quantizing in only the columnwise direction not supported yet!"); @@ -300,6 +308,13 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens // Dispatch to quantization kernel depending on data format switch (output_tensor->scaling_mode) { case NVTE_DELAYED_TENSOR_SCALING: { +#ifdef __HIP_PLATFORM_AMD__ + if (fp8::quantize_tdm( + *grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, + workspace_tensor, stream)) { + break; + } +#endif if (output_tensor->has_columnwise_data()) { NVTE_CHECK(output_tensor->has_data(), "Quantizing in only the columnwise direction not supported yet!"); diff --git a/transformer_engine/common/cast/fp8/rocm_quantize_fp8_tdm.cuh b/transformer_engine/common/cast/fp8/rocm_quantize_fp8_tdm.cuh new file mode 100644 index 000000000..b9c303769 --- /dev/null +++ b/transformer_engine/common/cast/fp8/rocm_quantize_fp8_tdm.cuh @@ -0,0 +1,431 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ +#pragma once +//#include "hip/hip_runtime.h" //dummy include to prevent hipification adding this header + +// gfx1250 FP8 tensor-scaling quantize with the Tensor Data Mover: each wave moves its own +// BAND_Y-row band of a 128x128 chunk global->LDS->global, and for the transposed output the waves +// assemble the chunk's transpose in LDS. The launchers return false where TDM does not apply, +// leaving the work to the generic kernels. Like the other rocm_ headers, this expects its includer +// to have included the common headers, util/system.h and util/cuda_runtime.h. + +#include + +#include "../../util/math.h" +#include "../../util/rocm_device_utils.cuh" +#include "../../util/rocm_tdm.cuh" +#include "../core/rocm_tdm_common.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace fp8 { +namespace tdm_kernel { + +using common::use_tdm; + +constexpr size_t CHUNK_DIM_Y = 128; +constexpr size_t CHUNK_DIM_X = 128; +constexpr size_t THREADS_PER_CHUNK = 128; +constexpr size_t WAVE_SIZE = 32; +constexpr size_t WAVES = THREADS_PER_CHUNK / WAVE_SIZE; // 4 +constexpr size_t BAND_Y = CHUNK_DIM_Y / WAVES; // 32 +constexpr size_t COLS_PER_LANE = CHUNK_DIM_X / WAVE_SIZE; // 4 +constexpr size_t QUADS = BAND_Y / 4; // 8 +constexpr size_t MULTI_MAX_TENSORS = 256; +static_assert(CHUNK_DIM_Y == common::tdm_dbias_kernel::CHUNK_DIM_Y, + "A dbias partial row covers a chunk"); + +template +using band_t = T[BAND_Y][CHUNK_DIM_X]; + +#ifdef __gfx1250__ +// Like the generic kernels, plain casts saturate as rocm_pack_4xfloat8 does, Inf included, and +// fused ones as static_cast, which keeps Inf and NaN. +template +__device__ __forceinline__ uint32_t cvt_4x_to_fp8(const float x0, const float x1, const float x2, + const float x3) { + if constexpr (SATURATE_INF) { + return rocm_pack_4xfloat8(x0, x1, x2, x3); + } else { + constexpr float max = transformer_engine::detail::TypeExtrema::max; + const auto saturate = [](const float x) { + return __builtin_isfinite(x) ? __builtin_amdgcn_fmed3f(x, max, -max) : x; + }; + int packed = 0; + if constexpr (std::is_same_v) { + packed = __builtin_amdgcn_cvt_pk_fp8_f32(saturate(x0), saturate(x1), packed, false); + packed = __builtin_amdgcn_cvt_pk_fp8_f32(saturate(x2), saturate(x3), packed, true); + } else { + packed = __builtin_amdgcn_cvt_pk_bf8_f32(saturate(x0), saturate(x1), packed, false); + packed = __builtin_amdgcn_cvt_pk_bf8_f32(saturate(x2), saturate(x3), packed, true); + } + return static_cast(packed); + } +} +#endif // __gfx1250__ + +// Quantizes chunk (chunk_Y, chunk_X) of a rows x cols tensor into output and, with TRANSPOSE, into +// its cols x rows transpose output_t. The input holds the first in_rows rows; the rest are zero. +// Writes the chunk's dbias partials to row chunk_Y of dbias_workspace. +template +__device__ __forceinline__ void quantize_chunk(const IType *input, const IType *act_input, + OType *output, OType *output_t, const float scale, + float *amax_ptr, float *dbias_workspace, + const size_t in_rows, const size_t rows, + const size_t cols, const size_t chunk_Y, + const size_t chunk_X) { +#ifdef __gfx1250__ + constexpr bool FUSED = IS_DBIAS || IS_DACT || IS_ACT; + __shared__ alignas(128) band_t in_sh[WAVES]; + __shared__ alignas(128) + IType act_sh[IS_DACT ? WAVES : 1][IS_DACT ? BAND_Y : 1][IS_DACT ? CHUNK_DIM_X : 1]; + __shared__ alignas(128) band_t out_sh[WAVES]; + // Wave w stores rows [w * BAND_Y, (w + 1) * BAND_Y) of the transposed chunk. TDM stores read LDS + // densely. + __shared__ alignas(128) OType t_sh[TRANSPOSE ? WAVES : 1][BAND_Y][CHUNK_DIM_Y]; + + const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / WAVE_SIZE); + const int lane = threadIdx.x % WAVE_SIZE; + const int lane_offset_X = lane * COLS_PER_LANE; + const size_t chunk_offset_Y = chunk_Y * CHUNK_DIM_Y; + const size_t chunk_offset_X = chunk_X * CHUNK_DIM_X; + const size_t band_offset_Y = chunk_offset_Y + wave * BAND_Y; + auto &in = in_sh[wave]; + auto &out = out_sh[wave]; + // Out-of-bounds elements are zero in LDS and every activation (or derivative times a zero + // gradient) keeps them zero, so they leave the amax and dbias unchanged. + float amax = 0; + float partial_dbias[COLS_PER_LANE] = {}; + + if (band_offset_Y < rows) { + tdm::load_2d(in, input, in_rows, cols, cols, band_offset_Y, chunk_offset_X); + if constexpr (IS_DACT) { + tdm::load_2d(act_sh[wave], act_input, in_rows, cols, cols, band_offset_Y, chunk_offset_X); + } + tdm::wait<0>(); + + // A lane converts its COLS_PER_LANE columns a quad of rows at a time and transposes each 4x4 + // block of bytes. A transposed dword's bank depends only on its quad, so the lanes start at + // different quads to spread their transposed writes over QUADS banks. + const int first_quad = TRANSPOSE ? lane % QUADS : 0; +#pragma unroll + for (int step = 0; step < QUADS; step++) { + const int quad = (first_quad + step) % QUADS; + uint32_t packed[4]; +#pragma unroll + for (int i = 0; i < 4; i++) { + const int y = quad * 4 + i; + Vec x; + Vec a; + x.load_from(&in[y][lane_offset_X]); + if constexpr (IS_DACT) { + a.load_from(&act_sh[wave][y][lane_offset_X]); + } + float elt[COLS_PER_LANE]; +#pragma unroll + for (int j = 0; j < COLS_PER_LANE; j++) { + elt[j] = static_cast(x.data.elt[j]); + if constexpr (IS_ACT) { + elt[j] = OP(elt[j], {}); + } + if constexpr (IS_DACT) { + elt[j] *= OP(static_cast(a.data.elt[j]), {}); + } + if constexpr (IS_DBIAS) { + partial_dbias[j] += elt[j]; + } + __builtin_assume(amax >= 0); + amax = fmaxf(fabsf(elt[j]), amax); + } + packed[i] = cvt_4x_to_fp8(elt[0] * scale, elt[1] * scale, elt[2] * scale, + elt[3] * scale); + *reinterpret_cast(&out[y][lane_offset_X]) = packed[i]; + } + if constexpr (TRANSPOSE) { + // Column j of the quad: byte j of each row + const uint32_t lo01 = __builtin_amdgcn_perm(packed[1], packed[0], 0x05010400); + const uint32_t hi01 = __builtin_amdgcn_perm(packed[1], packed[0], 0x07030602); + const uint32_t lo23 = __builtin_amdgcn_perm(packed[3], packed[2], 0x05010400); + const uint32_t hi23 = __builtin_amdgcn_perm(packed[3], packed[2], 0x07030602); + const uint32_t column[COLS_PER_LANE] = {__builtin_amdgcn_perm(lo23, lo01, 0x05040100), + __builtin_amdgcn_perm(lo23, lo01, 0x07060302), + __builtin_amdgcn_perm(hi23, hi01, 0x05040100), + __builtin_amdgcn_perm(hi23, hi01, 0x07060302)}; +#pragma unroll + for (int j = 0; j < COLS_PER_LANE; j++) { + const int x = lane_offset_X + j; + *reinterpret_cast(&t_sh[x / BAND_Y][x % BAND_Y][wave * BAND_Y + quad * 4]) = + column[j]; + } + } + } + tdm::store_2d(output, out, rows, cols, cols, band_offset_Y, chunk_offset_X); + } + + // Every wave leaves its column sums and amax in its consumed in_sh band. + asm volatile("" ::: "memory"); // in_sh is reused as float + float *const wave_scratch = reinterpret_cast(in_sh[wave]); + if constexpr (IS_DBIAS) { + Vec partial; +#pragma unroll + for (int i = 0; i < COLS_PER_LANE; i++) { + partial.data.elt[i] = partial_dbias[i]; + } + partial.store_to(wave_scratch + lane_offset_X); + } + if (amax_ptr != nullptr) { + const float wave_amax = rocm_subwarp_allreduce(amax, rocm_op::max{}); + if (lane == 0) { + wave_scratch[CHUNK_DIM_X] = wave_amax; + } + } + __syncthreads(); + if constexpr (TRANSPOSE) { + // Rows past the tensor's are garbage in t_sh but outside the transpose's columns. + const size_t t_offset_Y = chunk_offset_X + wave * BAND_Y; + if (t_offset_Y < cols) { + tdm::store_2d(output_t, t_sh[wave], cols, rows, rows, t_offset_Y, chunk_offset_Y); + } + } + if constexpr (IS_DBIAS) { + const size_t col = chunk_offset_X + threadIdx.x; + if (col < cols) { + float dbias = 0; +#pragma unroll + for (int w = 0; w < WAVES; w++) { + dbias += reinterpret_cast(in_sh[w])[threadIdx.x]; + } + dbias_workspace[chunk_Y * cols + col] = dbias; + } + } + + if (amax_ptr != nullptr && threadIdx.x == 0) { + float block_amax = 0; +#pragma unroll + for (int w = 0; w < WAVES; w++) { + block_amax = fmaxf(block_amax, reinterpret_cast(in_sh[w])[CHUNK_DIM_X]); + } + // Atomics on one address serialize, so only those that may raise the amax are issued. It only + // grows during the kernel, so a stale read can only cause a redundant atomic. + if (block_amax > *amax_ptr) { + rocm_atomicMaxFloat(amax_ptr, block_amax); + } + } + tdm::wait<0>(); +#endif // __gfx1250__ +} + +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + quantize_fp8_tdm_kernel(const IType *input, const IType *act_input, OType *output, + OType *output_t, const float *noop, const float *scale_ptr, + float *amax_ptr, float *scale_inv_ptr, float *dbias_workspace, + const size_t rows, const size_t cols) { + // The launcher leaves fused quantizations with a noop flag to the generic kernels. + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + const float scale = scale_ptr != nullptr ? *scale_ptr : 1.0f; + quantize_chunk( + input, act_input, output, output_t, scale, amax_ptr, dbias_workspace, rows, rows, cols, + blockIdx.y, blockIdx.x); + if (blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0 && scale_inv_ptr != nullptr) { + *scale_inv_ptr = __frcp_rn(scale); + } +} + +struct MultiCastTransposeArgs { + const void *input[MULTI_MAX_TENSORS]; + void *output[MULTI_MAX_TENSORS]; + void *output_t[MULTI_MAX_TENSORS]; + const float *scale[MULTI_MAX_TENSORS]; + float *amax[MULTI_MAX_TENSORS]; + float *scale_inv[MULTI_MAX_TENSORS]; + uint32_t in_rows[MULTI_MAX_TENSORS]; + uint32_t rows[MULTI_MAX_TENSORS]; + uint32_t cols[MULTI_MAX_TENSORS]; + // Tensor t owns chunk rows [first_chunk_Y[t], first_chunk_Y[t + 1]). + uint32_t first_chunk_Y[MULTI_MAX_TENSORS + 1]; + int num_tensors; +}; + +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + multi_cast_transpose_fp8_tdm_kernel(const MultiCastTransposeArgs args) { + int t = 0; + int hi = args.num_tensors - 1; + while (t < hi) { + const int mid = t + (hi - t) / 2; + if (args.first_chunk_Y[mid + 1] <= blockIdx.y) { + t = mid + 1; + } else { + hi = mid; + } + } + const size_t cols = args.cols[t]; + const size_t chunk_Y = blockIdx.y - args.first_chunk_Y[t]; + if (blockIdx.x * CHUNK_DIM_X >= cols) { + return; + } + const float scale = args.scale[t] != nullptr ? *args.scale[t] : 1.0f; + quantize_chunk( + static_cast(args.input[t]), nullptr, static_cast(args.output[t]), + static_cast(args.output_t[t]), scale, args.amax[t], nullptr, args.in_rows[t], + args.rows[t], cols, chunk_Y, blockIdx.x); + if (chunk_Y == 0 && blockIdx.x == 0 && threadIdx.x == 0 && args.scale_inv[t] != nullptr) { + *args.scale_inv[t] = __frcp_rn(scale); + } +} + +} // namespace tdm_kernel + +// FP8 tensor scaling of input, after an activation or its derivative, into output's rowwise data +// and, if allocated, its transposed columnwise data. +template +bool quantize_tdm(const Tensor &input, const Tensor *act_input, const Tensor *noop, Tensor *output, + Tensor *dbias, Tensor *workspace, hipStream_t stream) { + using namespace tdm_kernel; + constexpr bool FUSED = IS_DBIAS || IS_DACT || IS_ACT; + const bool transpose = output->has_columnwise_data(); + const auto [rows, cols] = input.flat_2d_dims(); + const bool has_noop = noop != nullptr && noop->data.dptr != nullptr; + // The transpose's rows are rows long. + if (!output->has_data() || !is_tensor_scaling(output->scaling_mode) || + is_fp8_dtype(input.dtype()) || !is_fp8_dtype(output->dtype()) || rows == 0 || + (transpose && rows % 16 != 0) || (FUSED && has_noop) || + (IS_DACT && (act_input == nullptr || !act_input->has_data())) || + (IS_DBIAS && (dbias == nullptr || workspace == nullptr)) || + !use_tdm(rows, cols, + {input.data.dptr, IS_DACT ? act_input->data.dptr : nullptr, output->data.dptr, + output->columnwise_data.dptr})) { + return false; + } + // One row of dbias partials per chunk row + const size_t chunks_Y = DIVUP(rows, CHUNK_DIM_Y); + if constexpr (IS_DBIAS) { + if (workspace->data.dptr == nullptr) { + workspace->data.shape = {chunks_Y, cols}; + workspace->data.dtype = DType::kFloat32; + return true; + } + NVTE_CHECK(workspace->data.dtype == DType::kFloat32 && + workspace->data.numel() >= chunks_Y * cols && + reinterpret_cast(workspace->data.dptr) % 16 == 0, + "Invalid dbias workspace (expected an aligned FP32 buffer of ", chunks_Y * cols, + " elements)"); + CheckOutputTensor(*dbias, "quantize_fp8_tdm_dbias"); + NVTE_CHECK(dbias->data.dtype == input.dtype(), "DBias must have the same type as input."); + NVTE_CHECK(dbias->data.shape == std::vector{cols}, "Wrong shape of DBias."); + } + if constexpr (IS_DACT) { + CheckInputTensor(*act_input, "quantize_fp8_tdm_act_input"); + NVTE_CHECK(input.dtype() == act_input->dtype(), "Types of both inputs must match."); + NVTE_CHECK(input.data.shape == act_input->data.shape, "Shapes of both inputs must match."); + } + if (has_noop) { + CheckNoopTensor(*noop, "quantize_fp8_tdm_noop"); + } + CheckInputTensor(input, "quantize_fp8_tdm_input"); + CheckOutputTensor(*output, "quantize_fp8_tdm_output"); + NVTE_CHECK(input.data.shape == output->data.shape, "Input and output shapes do not match (input=", + input.data.shape, ", output=", output->data.shape, ")"); + if (transpose) { + NVTE_CHECK(output->data.dtype == output->columnwise_data.dtype, + "Cast and transposed output types must match."); + NVTE_CHECK(output->scale_inv.dptr == output->columnwise_scale_inv.dptr, + "Cast and transposed outputs need to share scale-inverse tensor."); + } + + float *const workspace_ptr = IS_DBIAS ? reinterpret_cast(workspace->data.dptr) : nullptr; + const dim3 grid(DIVUP(cols, CHUNK_DIM_X), chunks_Y); + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input.dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output->dtype(), OType, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + transpose, TRANSPOSE, + quantize_fp8_tdm_kernel<<>>( + reinterpret_cast(input.data.dptr), + IS_DACT ? reinterpret_cast(act_input->data.dptr) : nullptr, + reinterpret_cast(output->data.dptr), + reinterpret_cast(output->columnwise_data.dptr), + has_noop ? reinterpret_cast(noop->data.dptr) : nullptr, + reinterpret_cast(output->scale.dptr), + reinterpret_cast(output->amax.dptr), + reinterpret_cast(output->scale_inv.dptr), workspace_ptr, rows, + cols);))); // NOLINT(*) + NVTE_CHECK_CUDA(hipGetLastError()); + if constexpr (IS_DBIAS) { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input.dtype(), IType, + common::reduce_dbias_tdm(reinterpret_cast(dbias->data.dptr), workspace_ptr, + cols, 1, chunks_Y, nullptr, nullptr, stream);); + } + return true; +} + +// FP8 cast and transpose of each input into its output, whose rows past the input's (with +// valid_num_rows, valid_num_rows[i] of them) are zero padding. TDM only if every tensor qualifies. +inline bool multi_cast_transpose_tdm(const std::vector &input_list, + const std::vector &output_list, + const int *valid_num_rows, hipStream_t stream) { + using namespace tdm_kernel; + const size_t num_tensors = input_list.size(); + for (size_t i = 0; i < num_tensors; i++) { + const Tensor &output = *output_list[i]; + const size_t rows = output.data.shape[0]; + if (!is_tensor_scaling(output.scaling_mode) || is_fp8_dtype(input_list[i]->dtype()) || + !is_fp8_dtype(output.dtype()) || rows % 16 != 0 || + !use_tdm(rows, output.data.shape[1], + {input_list[i]->data.dptr, output.data.dptr, output.columnwise_data.dptr})) { + return false; + } + } + for (size_t first = 0; first < num_tensors; first += MULTI_MAX_TENSORS) { + const size_t count = std::min(num_tensors - first, MULTI_MAX_TENSORS); + MultiCastTransposeArgs args; + args.num_tensors = count; + args.first_chunk_Y[0] = 0; + size_t chunks_X = 0; + for (size_t j = 0; j < count; j++) { + const Tensor &input = *input_list[first + j]; + const Tensor &output = *output_list[first + j]; + const size_t rows = output.data.shape[0]; + const size_t cols = output.data.shape[1]; + args.input[j] = input.data.dptr; + args.output[j] = output.data.dptr; + args.output_t[j] = output.columnwise_data.dptr; + args.scale[j] = reinterpret_cast(output.scale.dptr); + args.amax[j] = reinterpret_cast(output.amax.dptr); + args.scale_inv[j] = reinterpret_cast(output.scale_inv.dptr); + args.in_rows[j] = valid_num_rows != nullptr ? valid_num_rows[first + j] : rows; + args.rows[j] = rows; + args.cols[j] = cols; + args.first_chunk_Y[j + 1] = args.first_chunk_Y[j] + DIVUP(rows, CHUNK_DIM_Y); + chunks_X = std::max(chunks_X, DIVUP(cols, CHUNK_DIM_X)); + } + if (args.first_chunk_Y[count] == 0 || chunks_X == 0) { + continue; + } + const dim3 grid(chunks_X, args.first_chunk_Y[count]); + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input_list[0]->dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output_list[0]->dtype(), OType, + multi_cast_transpose_fp8_tdm_kernel + <<>>(args);)); // NOLINT(*) + NVTE_CHECK_CUDA(hipGetLastError()); + } + return true; +} + +} // namespace fp8 +} // namespace dispatch +} // namespace transformer_engine diff --git a/transformer_engine/common/cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh b/transformer_engine/common/cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh index b4049131c..53b6684d9 100644 --- a/transformer_engine/common/cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh +++ b/transformer_engine/common/cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh @@ -12,17 +12,18 @@ // leaving the work to the generic kernels. Like the other rocm_ headers, this expects its includer // to have included the common headers. -#include -#include - #include "../../util/rocm_device_utils.cuh" #include "../../util/rocm_tdm.cuh" +#include "../core/rocm_tdm_common.cuh" namespace transformer_engine { namespace dispatch { namespace mxfp8 { namespace tdm_kernel { +using common::reduce_dbias_tdm; +using common::use_tdm; + constexpr size_t CHUNK_DIM_Y = 128; constexpr size_t CHUNK_DIM_X = 128; constexpr size_t THREADS_PER_CHUNK = 128; @@ -30,21 +31,13 @@ constexpr size_t WAVE_SIZE = 32; constexpr size_t WAVES = THREADS_PER_CHUNK / WAVE_SIZE; // 4 constexpr size_t BAND_Y = CHUNK_DIM_Y / WAVES; // 32 constexpr size_t COLS_PER_LANE = CHUNK_DIM_X / WAVE_SIZE; // 4 -constexpr size_t DBIAS_WAVES = 8; constexpr size_t MULTI_MAX_TENSORS = 1024; +static_assert(CHUNK_DIM_Y == common::tdm_dbias_kernel::CHUNK_DIM_Y, + "A dbias partial row covers a chunk"); template using band_t = T[BAND_Y][CHUNK_DIM_X]; -// TDM tiles need 16-byte aligned rows and 32-bit tensor extents. -inline bool use_tdm(const size_t rows, const size_t cols, - const std::initializer_list tensors) { - static const bool enabled = getenv("NVTE_USE_TDM", true); - const auto aligned = [](const void *ptr) { return reinterpret_cast(ptr) % 16 == 0; }; - return enabled && cuda::sm_arch() == 125 && cols % 16 == 0 && rows <= UINT32_MAX && - cols <= UINT32_MAX && std::all_of(tensors.begin(), tensors.end(), aligned); -} - #ifdef __gfx1250__ // hip_bfloat16 rounds with per-element branches; __bf16 converts with one instruction. template @@ -423,84 +416,6 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) nullptr, args.amax[t], rows, cols, DIVUP(cols, SCALE_DIM_X), cols, chunk_Y, blockIdx.x, 0); } -// Sums each tensor's dbias partials, one row per chunk, into its row of dbias. A block covers -// DBIAS_WAVES / waves_per_slab 128-column slabs, whose waves split the BAND_Y-row bands. -template -__global__ void __launch_bounds__(DBIAS_WAVES * WAVE_SIZE) - reduce_dbias_tdm_kernel(OType *const dbias, const float *const partial, const size_t cols, - const size_t partial_rows, const int64_t *const offsets_ptr, - const int64_t *const first_dims_ptr, const int waves_per_slab) { -#ifdef __gfx1250__ - __shared__ alignas(128) band_t band_sh[DBIAS_WAVES]; - - const size_t tensor_id = blockIdx.y; - size_t tensor_rows = partial_rows / gridDim.y; - size_t tensor_offset_Y = tensor_id * tensor_rows; - if (offsets_ptr != nullptr) { - tensor_rows = static_cast(first_dims_ptr[tensor_id]) / CHUNK_DIM_Y; - tensor_offset_Y = static_cast(offsets_ptr[tensor_id]) / cols / CHUNK_DIM_Y; - } - const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / WAVE_SIZE); - const size_t block_cols = DBIAS_WAVES / waves_per_slab * CHUNK_DIM_X; - const size_t block_offset_X = blockIdx.x * block_cols; - const size_t slab_offset_X = block_offset_X + wave / waves_per_slab * CHUNK_DIM_X; - const int lane_offset_X = (threadIdx.x % WAVE_SIZE) * COLS_PER_LANE; - const float *const tensor_partial = partial + tensor_offset_Y * cols; - auto &band = band_sh[wave]; - - Vec sum; - sum.clear(); - if (slab_offset_X < cols) { - for (size_t band_offset_Y = wave % waves_per_slab * BAND_Y; band_offset_Y < tensor_rows; - band_offset_Y += waves_per_slab * BAND_Y) { - tdm::load_2d(band, tensor_partial, tensor_rows, cols, cols, band_offset_Y, slab_offset_X); - tdm::wait<0>(); -#pragma unroll - for (int y = 0; y < BAND_Y; y++) { - Vec x; - x.load_from(&band[y][lane_offset_X]); -#pragma unroll - for (int i = 0; i < COLS_PER_LANE; i++) { - sum.data.elt[i] += x.data.elt[i]; - } - } - } - } - sum.store_to(&band[0][lane_offset_X]); - __syncthreads(); - for (size_t x = threadIdx.x; x < block_cols; x += blockDim.x) { - const size_t col = block_offset_X + x; - if (col < cols) { - const int first_wave = x / CHUNK_DIM_X * waves_per_slab; - float dbias_sum = 0; - for (int w = first_wave; w < first_wave + waves_per_slab; w++) { - dbias_sum += band_sh[w][0][x % CHUNK_DIM_X]; - } - dbias[tensor_id * cols + col] = static_cast(dbias_sum); - } - } -#endif // __gfx1250__ -} - -// partial_rows counts the partials of all num_tensors tensors. Without offsets, the tensors share -// them equally; otherwise tensor t has first_dims[t] / CHUNK_DIM_Y of them. -template -void reduce_dbias(OType *dbias, const float *partial, const size_t cols, const size_t num_tensors, - const size_t partial_rows, const int64_t *offsets_ptr, - const int64_t *first_dims_ptr, hipStream_t stream) { - // About one band per wave; short tensors leave the spare waves to more slabs. - const size_t bands = DIVUP(partial_rows / num_tensors, BAND_Y); - size_t waves_per_slab = 1; - while (waves_per_slab < DBIAS_WAVES && waves_per_slab < bands) { - waves_per_slab *= 2; - } - const dim3 grid(DIVUP(cols, DBIAS_WAVES / waves_per_slab * CHUNK_DIM_X), num_tensors); - reduce_dbias_tdm_kernel<<>>( - dbias, partial, cols, partial_rows, offsets_ptr, first_dims_ptr, - static_cast(waves_per_slab)); - NVTE_CHECK_CUDA(hipGetLastError()); -} - } // namespace tdm_kernel template (reinterpret_cast(dbias->data.dptr), workspace_ptr, cols, 1, - grid.y, nullptr, nullptr, stream);); + reduce_dbias_tdm(reinterpret_cast(dbias->data.dptr), workspace_ptr, cols, + 1, grid.y, nullptr, nullptr, stream);); } return true; } @@ -607,9 +522,9 @@ bool group_quantize_tdm(const GroupedTensor *input, const GroupedTensor *activat if constexpr (IS_DBIAS) { TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input->dtype(), IType, - reduce_dbias(reinterpret_cast(dbias->data.dptr), workspace_ptr, cols, - num_tensors, grid.y, same_shape ? nullptr : offsets_ptr, - same_shape ? nullptr : first_dims_ptr, stream);); + reduce_dbias_tdm(reinterpret_cast(dbias->data.dptr), workspace_ptr, cols, + num_tensors, grid.y, same_shape ? nullptr : offsets_ptr, + same_shape ? nullptr : first_dims_ptr, stream);); } return true; } diff --git a/transformer_engine/common/recipe/current_scaling.cu b/transformer_engine/common/recipe/current_scaling.cu index 7cb0c6695..0981e8e46 100644 --- a/transformer_engine/common/recipe/current_scaling.cu +++ b/transformer_engine/common/recipe/current_scaling.cu @@ -17,6 +17,11 @@ #include "../util/logging.h" #include "../util/vectorized_pointwise.h" #include "recipe_common.cuh" +#ifdef __HIP_PLATFORM_AMD__ +#include "../util/cuda_runtime.h" +#include "../util/system.h" +#include "../cast/core/rocm_tdm_common.cuh" +#endif //__HIP_PLATFORM_AMD__ namespace transformer_engine { namespace { @@ -133,6 +138,11 @@ void launch_amax_kernel(const InputType *input, float *amax, const size_t N, if (N == 0) { return; } +#ifdef __HIP_PLATFORM_AMD__ + if (dispatch::common::compute_amax_tdm(input, amax, N, noop_ptr, stream)) { + return; + } +#endif //__HIP_PLATFORM_AMD__ // Figure out alignment auto align = CheckAlignment(N, nvec, input); diff --git a/transformer_engine/common/transpose/multi_cast_transpose.cu b/transformer_engine/common/transpose/multi_cast_transpose.cu index ce1f68b93..d40952c0c 100644 --- a/transformer_engine/common/transpose/multi_cast_transpose.cu +++ b/transformer_engine/common/transpose/multi_cast_transpose.cu @@ -15,6 +15,11 @@ #include "../common.h" #include "../utils.cuh" +#ifdef __HIP_PLATFORM_AMD__ +#include "../util/cuda_runtime.h" +#include "../util/system.h" +#include "../cast/fp8/rocm_quantize_fp8_tdm.cuh" +#endif // #ifdef __HIP_PLATFORM_AMD__ namespace transformer_engine { @@ -261,6 +266,10 @@ void multi_cast_transpose(const std::vector input_list, std::vector Date: Tue, 27 Oct 2026 19:58:37 +0000 Subject: [PATCH 4/8] assume gfx150 has TDM --- tests/cpp/util/test_rocm_tdm.cu | 2 +- .../cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh | 11 ++--- transformer_engine/common/util/rocm_tdm.cuh | 42 +------------------ 3 files changed, 5 insertions(+), 50 deletions(-) diff --git a/tests/cpp/util/test_rocm_tdm.cu b/tests/cpp/util/test_rocm_tdm.cu index cc8530b01..6fc59f5fa 100644 --- a/tests/cpp/util/test_rocm_tdm.cu +++ b/tests/cpp/util/test_rocm_tdm.cu @@ -24,7 +24,7 @@ template __global__ void __launch_bounds__(kWaves * 32) tdm_roundtrip_kernel(const T *in, T *out, T *lds_dump, uint32_t rows, uint32_t cols, uint32_t ld_in, uint32_t ld_out) { -#ifdef NVTE_ROCM_HAS_TDM +#ifdef __gfx1250__ __shared__ alignas(128) T tile[kWaves][kTileRows][kTileCols]; const uint32_t wave = threadIdx.x / 32; const uint32_t row = (blockIdx.y * kWaves + wave) * kTileRows; diff --git a/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh b/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh index fb27953f3..959ee97ae 100644 --- a/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh +++ b/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh @@ -6,7 +6,7 @@ #pragma once // gfx1250 MXFP8 dequantize with the Tensor Data Mover, included into dequantize_kernel after -// rocm_dequantize_mxfp8.cuh, whose grid, scale indexing and fallback chunk it shares. +// rocm_dequantize_mxfp8.cuh, whose grid and scale indexing it shares. //#include "hip/hip_runtime.h" //dummy include to prevent hipification adding this header // Each wave moves its own TDM_BAND_Y-row band of the chunk global->LDS->global, so waves never @@ -22,7 +22,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const e8m0_t *const scales_ptr, const size_t rows, const size_t cols, const size_t scales_stride, const size_t mx_swizzle_padded_dim) { -#ifdef NVTE_ROCM_HAS_TDM +#ifdef __gfx1250__ constexpr bool USE_ROWWISE_SCALING = SCALE_DIM_X > 1; __shared__ alignas(128) IType in_sh[TDM_WAVES][TDM_BAND_Y][CHUNK_DIM_X]; @@ -111,12 +111,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) tdm::store_2d(output_ptr, out_sh[wave], rows, cols, cols, band_offset_Y, chunk_offset_X); tdm::wait<0>(); -#elif defined(__gfx1250__) - // Compiler without the TDM builtins - dequantize_mxfp8_chunk( - input_ptr, output_ptr, scales_ptr, rows, cols, scales_stride, mx_swizzle_padded_dim, blockIdx.y, - blockIdx.x); -#endif // NVTE_ROCM_HAS_TDM +#endif // __gfx1250__ } // Returns false, leaving the work to the generic kernel, where TDM does not apply. The scale diff --git a/transformer_engine/common/util/rocm_tdm.cuh b/transformer_engine/common/util/rocm_tdm.cuh index c40f90a4e..d40ea8b7c 100644 --- a/transformer_engine/common/util/rocm_tdm.cuh +++ b/transformer_engine/common/util/rocm_tdm.cuh @@ -8,25 +8,16 @@ // gfx1250 Tensor Data Mover (TDM): asynchronous 2D tile copies between global memory and LDS. // Descriptors are packed here rather than with hip/amd_detail/amd_gfx1250_TDM.h, whose stride -// setters truncate (ROCm/rocm-systems#13274). With a compiler that lacks the TDM builtins, the -// calling wave copies the tile itself. +// setters truncate (ROCm/rocm-systems#13274). The 6-operand TDM builtins need clang 23 or newer. #include #include -#include - -#if defined(__gfx1250__) && __has_builtin(__builtin_amdgcn_tensor_load_to_lds) && \ - __has_builtin(__builtin_amdgcn_tensor_store_from_lds) && \ - __has_builtin(__builtin_amdgcn_s_wait_tensorcnt) -#define NVTE_ROCM_HAS_TDM 1 -#endif #ifdef __gfx1250__ namespace transformer_engine { namespace tdm { namespace detail { -#ifdef NVTE_ROCM_HAS_TDM // The builtins' operand signedness has changed between compiler versions; integer vectors of // either signedness convert implicitly. typedef uint32_t u32x4 __attribute__((ext_vector_type(4))); @@ -84,35 +75,6 @@ __device__ __forceinline__ void copy_2d(const T (&lds)[TILE_ROWS][TILE_COLS], co __builtin_amdgcn_tensor_load_to_lds(g0, g1, z4, z4, z8, 0); } } -#else -template -__device__ __forceinline__ void copy_2d(const T (&lds)[TILE_ROWS][TILE_COLS], const T *gmem, - uint32_t rows, uint32_t cols, uint64_t ld, uint32_t row, - uint32_t col) { - static_assert(sizeof(T) == 1 || sizeof(T) == 2 || sizeof(T) == 4 || sizeof(T) == 8, - "TDM elements are 1, 2, 4 or 8 bytes"); - using U = std::conditional_t< - sizeof(T) == 1, uint8_t, - std::conditional_t>>; - U *const tile = reinterpret_cast(const_cast(&lds[0][0])); - U *const tensor = reinterpret_cast(const_cast(gmem)); - // The tile is accessed as U here and as T by the caller. - asm volatile("" ::: "memory"); - for (uint32_t i = __lane_id(); i < TILE_ROWS * TILE_COLS; i += warpSize) { - const uint64_t r = row + static_cast(i / TILE_COLS); - const uint64_t c = col + static_cast(i % TILE_COLS); - const bool inside = r < rows && c < cols; - if constexpr (STORE) { - if (inside) { - tensor[r * ld + c] = tile[i]; - } - } else { - tile[i] = inside ? tensor[r * ld + c] : U{0}; - } - } -} -#endif // NVTE_ROCM_HAS_TDM } // namespace detail @@ -139,9 +101,7 @@ __device__ __forceinline__ void store_2d(T *gmem, const T (&lds)[TILE_ROWS][TILE // the loaded data only after a subsequent barrier. template __device__ __forceinline__ void wait() { -#ifdef NVTE_ROCM_HAS_TDM __builtin_amdgcn_s_wait_tensorcnt(N); -#endif asm volatile("" ::: "memory"); // the compiler does not see the TDM's LDS writes } From 9938df03215d6df5c9e887e039eadbde7f04f6e7 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 27 Oct 2026 22:16:33 +0000 Subject: [PATCH 5/8] spread TDM amax atomics over slots --- transformer_engine/common/CMakeLists.txt | 3 +- .../common/cast/fp8/rocm_quantize_fp8_tdm.cuh | 43 ++++++----- .../common/util/rocm_amax_slots.cu | 73 +++++++++++++++++++ .../common/util/rocm_amax_slots.cuh | 33 +++++++++ 4 files changed, 134 insertions(+), 18 deletions(-) create mode 100644 transformer_engine/common/util/rocm_amax_slots.cu create mode 100644 transformer_engine/common/util/rocm_amax_slots.cuh diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 4dbcc9e7f..b01cf511b 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -406,7 +406,8 @@ if(USE_ROCM) gemm/ck_grouped_gemm/ck_grouped_gemm_fp16_tn.cpp gemm/ck_grouped_gemm/ck_grouped_gemm_fp16_tt.cpp gemm/ck_grouped_gemm/ck_mx_grouped_gemm.cpp - gemm/rocm_gemm.cu) + gemm/rocm_gemm.cu + util/rocm_amax_slots.cu) endif() if(USE_CUDA) diff --git a/transformer_engine/common/cast/fp8/rocm_quantize_fp8_tdm.cuh b/transformer_engine/common/cast/fp8/rocm_quantize_fp8_tdm.cuh index b9c303769..50a6048a7 100644 --- a/transformer_engine/common/cast/fp8/rocm_quantize_fp8_tdm.cuh +++ b/transformer_engine/common/cast/fp8/rocm_quantize_fp8_tdm.cuh @@ -15,6 +15,7 @@ #include #include "../../util/math.h" +#include "../../util/rocm_amax_slots.cuh" #include "../../util/rocm_device_utils.cuh" #include "../../util/rocm_tdm.cuh" #include "../core/rocm_tdm_common.cuh" @@ -69,15 +70,16 @@ __device__ __forceinline__ uint32_t cvt_4x_to_fp8(const float x0, const float x1 // Quantizes chunk (chunk_Y, chunk_X) of a rows x cols tensor into output and, with TRANSPOSE, into // its cols x rows transpose output_t. The input holds the first in_rows rows; the rest are zero. -// Writes the chunk's dbias partials to row chunk_Y of dbias_workspace. +// Writes the chunk's dbias partials to row chunk_Y of dbias_workspace and raises the amax at its +// rocm_amax_target. template __device__ __forceinline__ void quantize_chunk(const IType *input, const IType *act_input, OType *output, OType *output_t, const float scale, - float *amax_ptr, float *dbias_workspace, - const size_t in_rows, const size_t rows, - const size_t cols, const size_t chunk_Y, - const size_t chunk_X) { + float *amax_ptr, float *amax_slots, + float *dbias_workspace, const size_t in_rows, + const size_t rows, const size_t cols, + const size_t chunk_Y, const size_t chunk_X) { #ifdef __gfx1250__ constexpr bool FUSED = IS_DBIAS || IS_DACT || IS_ACT; __shared__ alignas(128) band_t in_sh[WAVES]; @@ -209,10 +211,12 @@ __device__ __forceinline__ void quantize_chunk(const IType *input, const IType * for (int w = 0; w < WAVES; w++) { block_amax = fmaxf(block_amax, reinterpret_cast(in_sh[w])[CHUNK_DIM_X]); } + float *const amax_dst = + rocm_amax_target(amax_ptr, amax_slots, blockIdx.y * gridDim.x + blockIdx.x); // Atomics on one address serialize, so only those that may raise the amax are issued. It only // grows during the kernel, so a stale read can only cause a redundant atomic. - if (block_amax > *amax_ptr) { - rocm_atomicMaxFloat(amax_ptr, block_amax); + if (block_amax > *amax_dst) { + rocm_atomicMaxFloat(amax_dst, block_amax); } } tdm::wait<0>(); @@ -224,16 +228,16 @@ template ( - input, act_input, output, output_t, scale, amax_ptr, dbias_workspace, rows, rows, cols, - blockIdx.y, blockIdx.x); + input, act_input, output, output_t, scale, amax_ptr, amax_slots, dbias_workspace, rows, rows, + cols, blockIdx.y, blockIdx.x); if (blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0 && scale_inv_ptr != nullptr) { *scale_inv_ptr = __frcp_rn(scale); } @@ -275,8 +279,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const float scale = args.scale[t] != nullptr ? *args.scale[t] : 1.0f; quantize_chunk( static_cast(args.input[t]), nullptr, static_cast(args.output[t]), - static_cast(args.output_t[t]), scale, args.amax[t], nullptr, args.in_rows[t], - args.rows[t], cols, chunk_Y, blockIdx.x); + static_cast(args.output_t[t]), scale, args.amax[t], nullptr, nullptr, + args.in_rows[t], args.rows[t], cols, chunk_Y, blockIdx.x); if (chunk_Y == 0 && blockIdx.x == 0 && threadIdx.x == 0 && args.scale_inv[t] != nullptr) { *args.scale_inv[t] = __frcp_rn(scale); } @@ -343,7 +347,11 @@ bool quantize_tdm(const Tensor &input, const Tensor *act_input, const Tensor *no } float *const workspace_ptr = IS_DBIAS ? reinterpret_cast(workspace->data.dptr) : nullptr; + const float *const noop_ptr = + has_noop ? reinterpret_cast(noop->data.dptr) : nullptr; + float *const amax = reinterpret_cast(output->amax.dptr); const dim3 grid(DIVUP(cols, CHUNK_DIM_X), chunks_Y); + float *const amax_slots = amax != nullptr ? rocm_amax_slots(grid.x * grid.y, stream) : nullptr; TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input.dtype(), IType, TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( @@ -355,13 +363,14 @@ bool quantize_tdm(const Tensor &input, const Tensor *act_input, const Tensor *no reinterpret_cast(input.data.dptr), IS_DACT ? reinterpret_cast(act_input->data.dptr) : nullptr, reinterpret_cast(output->data.dptr), - reinterpret_cast(output->columnwise_data.dptr), - has_noop ? reinterpret_cast(noop->data.dptr) : nullptr, - reinterpret_cast(output->scale.dptr), - reinterpret_cast(output->amax.dptr), + reinterpret_cast(output->columnwise_data.dptr), noop_ptr, + reinterpret_cast(output->scale.dptr), amax, amax_slots, reinterpret_cast(output->scale_inv.dptr), workspace_ptr, rows, cols);))); // NOLINT(*) NVTE_CHECK_CUDA(hipGetLastError()); + if (amax_slots != nullptr) { + rocm_amax_finalize(amax_slots, amax, noop_ptr, stream); + } if constexpr (IS_DBIAS) { TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input.dtype(), IType, diff --git a/transformer_engine/common/util/rocm_amax_slots.cu b/transformer_engine/common/util/rocm_amax_slots.cu new file mode 100644 index 000000000..a89621a07 --- /dev/null +++ b/transformer_engine/common/util/rocm_amax_slots.cu @@ -0,0 +1,73 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ + +#include +#include +#include + +#include "../common.h" +#include "cuda_runtime.h" +#include "logging.h" +#include "rocm_amax_slots.cuh" +#include "rocm_device_utils.cuh" +#include "system.h" + +namespace transformer_engine { +namespace { + +constexpr int FINALIZE_THREADS = 32; + +__global__ void __launch_bounds__(FINALIZE_THREADS) + rocm_amax_finalize_kernel(float *slots, float *amax, const float *noop) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + // The maximum of the bits as integers, like the atomics take it + int amax_bits = 0; + for (int i = threadIdx.x; i < ROCM_AMAX_SLOTS; i += FINALIZE_THREADS) { + float *const slot = slots + i * ROCM_AMAX_SLOT_STRIDE; + amax_bits = max(amax_bits, __float_as_int(*slot)); + *slot = 0.0f; + } + amax_bits = rocm_subwarp_allreduce( + amax_bits, [](const int a, const int b) { return max(a, b); }); + if (threadIdx.x == 0) { + atomicMax(reinterpret_cast(amax), amax_bits); + } +} + +} // namespace + +float *rocm_amax_slots(const size_t num_blocks, hipStream_t stream) { + static const bool enabled = !getenv("NVTE_USE_ATOMIC_AMAX"); + if (!enabled || num_blocks <= ROCM_AMAX_SLOTS) { + return nullptr; + } + // One buffer per stream, so that concurrent kernels never share slots, kept for the graphs + // captured with it. Leaked deliberately: device memory must not be freed at static teardown. + static std::mutex &mutex = *new std::mutex; + static auto &buffers = *new std::map, float *>; + std::lock_guard lock(mutex); + float *&slots = buffers[{cuda::current_device(), stream}]; + if (slots == nullptr) { + hipStreamCaptureStatus capture_status; + NVTE_CHECK_CUDA(hipStreamIsCapturing(stream, &capture_status)); + if (capture_status != hipStreamCaptureStatusNone) { + return nullptr; + } + constexpr size_t bytes = ROCM_AMAX_SLOTS * ROCM_AMAX_SLOT_STRIDE * sizeof(float); + NVTE_CHECK_CUDA(hipMalloc(&slots, bytes)); + NVTE_CHECK_CUDA(hipMemsetAsync(slots, 0, bytes, stream)); + } + return slots; +} + +void rocm_amax_finalize(float *slots, float *amax, const float *noop, hipStream_t stream) { + rocm_amax_finalize_kernel<<<1, FINALIZE_THREADS, 0, stream>>>(slots, amax, noop); + NVTE_CHECK_CUDA(hipGetLastError()); +} + +} // namespace transformer_engine diff --git a/transformer_engine/common/util/rocm_amax_slots.cuh b/transformer_engine/common/util/rocm_amax_slots.cuh new file mode 100644 index 000000000..a3e12bc5c --- /dev/null +++ b/transformer_engine/common/util/rocm_amax_slots.cuh @@ -0,0 +1,33 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ +#pragma once +//#include "hip/hip_runtime.h" // prevent hipification of this rocm_ file + +// Atomics on one address serialize, so a kernel with many blocks can reduce its amax in two stages: +// each block raises one of ROCM_AMAX_SLOTS slots, and rocm_amax_finalize then folds the slots into +// the amax and zeroes them for the next kernel. + +#include + +constexpr int ROCM_AMAX_SLOTS = 64; +constexpr int ROCM_AMAX_SLOT_STRIDE = 32; // floats: a 128-byte line per slot + +// The address a block raises the amax at: its slot, or the amax itself without slots. +__device__ __forceinline__ float *rocm_amax_target(float *amax, float *slots, unsigned int block) { + return slots != nullptr ? slots + (block % ROCM_AMAX_SLOTS) * ROCM_AMAX_SLOT_STRIDE : amax; +} + +namespace transformer_engine { + +// The zeroed slots for a kernel of num_blocks blocks on stream, or nullptr to raise the amax +// directly: for few blocks, with NVTE_USE_ATOMIC_AMAX=1, or while capturing a stream that has no +// slots yet. +float *rocm_amax_slots(size_t num_blocks, hipStream_t stream); + +// Folds the slots into *amax and zeroes them, unless noop is set. +void rocm_amax_finalize(float *slots, float *amax, const float *noop, hipStream_t stream); + +} // namespace transformer_engine From 68daf1fd764a0a68a408e2fd6a3de5d107382deb Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 28 Oct 2026 23:24:53 +0000 Subject: [PATCH 6/8] remaining FP8/MXFP8 ops --- .../common/cast/core/rocm_tdm_common.cuh | 67 +- .../common/cast/dispatch/dequantize.cuh | 10 + .../common/cast/dispatch/gated.cuh | 14 + .../common/cast/fp8/rocm_quantize_fp8_tdm.cuh | 672 +++++++++++++++++- .../cast/mxfp8/group_dequantize_mxfp8.cuh | 8 +- .../cast/mxfp8/group_quantize_mxfp8.cuh | 5 +- .../common/cast/mxfp8/quantize_mxfp8.cuh | 6 +- .../cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh | 104 ++- .../cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh | 299 +++++--- .../quantize_transpose_square_blockwise.cu | 13 + .../quantize_transpose_vector_blockwise.cu | 18 + 11 files changed, 1069 insertions(+), 147 deletions(-) diff --git a/transformer_engine/common/cast/core/rocm_tdm_common.cuh b/transformer_engine/common/cast/core/rocm_tdm_common.cuh index a7410e031..55a21c258 100644 --- a/transformer_engine/common/cast/core/rocm_tdm_common.cuh +++ b/transformer_engine/common/cast/core/rocm_tdm_common.cuh @@ -7,8 +7,8 @@ //#include "hip/hip_runtime.h" //dummy include to prevent hipification adding this header // gfx1250 Tensor Data Mover pieces shared by the TDM quantize paths: the check whether TDM applies, -// the reduction of per-chunk dbias partials, and amax. This expects its includer to have included -// the common headers, util/system.h and util/cuda_runtime.h. +// the chunks of grouped tensors, the reduction of per-chunk dbias partials, and amax. This expects +// its includer to have included the common headers, util/system.h and util/cuda_runtime.h. #include #include @@ -29,6 +29,69 @@ inline bool use_tdm(const size_t rows, const size_t cols, cols <= UINT32_MAX && std::all_of(tensors.begin(), tensors.end(), aligned); } +#ifdef __gfx1250__ +// hip_bfloat16 rounds with per-element branches; __bf16 converts with one instruction. +template +__device__ __forceinline__ T cvt_from_float(const float x) { + if constexpr (std::is_same_v) { + return __builtin_bit_cast(T, static_cast<__bf16>(x)); + } else { + return static_cast(x); + } +} +#endif // __gfx1250__ + +// The last of num_tensors tensors that starts at or before element offset. +__device__ __forceinline__ size_t find_tensor(const int64_t *offsets, const size_t num_tensors, + const size_t offset) { + size_t lo = 1; + size_t hi = num_tensors; + while (lo < hi) { + const size_t mid = lo + (hi - lo) / 2; + if (static_cast(offsets[mid]) <= offset) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo - 1; +} + +struct GroupChunk { + size_t offset; // of the tensor's first element + size_t rows; // of the tensor, or 0 past the last one + size_t cols; + size_t chunk_Y; + size_t chunk_X; +}; + +// Chunk `chunk` of a grouped tensor whose tensors have dimensions divisible by the chunk's, so that +// tensor t owns the chunks from offsets[t] / (CHUNK_DIM_Y * CHUNK_DIM_X) on, row-major. Without +// first_dims, every tensor has first_logical_dim rows. +template +__device__ __forceinline__ GroupChunk find_group_chunk(const size_t chunk, const size_t num_tensors, + const size_t first_logical_dim, + const int64_t *offsets, + const int64_t *first_dims, + const int64_t *last_dims) { + constexpr size_t CHUNK_ELTS = CHUNK_DIM_Y * CHUNK_DIM_X; + const size_t chunk_offset = chunk * CHUNK_ELTS; + const size_t tensor_id = find_tensor(offsets, num_tensors, chunk_offset); + GroupChunk c{}; + c.offset = static_cast(offsets[tensor_id]); + c.rows = first_dims != nullptr ? static_cast(first_dims[tensor_id]) : first_logical_dim; + c.cols = static_cast(last_dims[tensor_id]); + if (chunk_offset >= c.offset + c.rows * c.cols) { + c.rows = 0; + return c; + } + const size_t chunks_X = c.cols / CHUNK_DIM_X; + const size_t tensor_chunk = (chunk_offset - c.offset) / CHUNK_ELTS; + c.chunk_Y = tensor_chunk / chunks_X; + c.chunk_X = tensor_chunk % chunks_X; + return c; +} + namespace tdm_dbias_kernel { constexpr size_t CHUNK_DIM_Y = 128; // input rows per partial row diff --git a/transformer_engine/common/cast/dispatch/dequantize.cuh b/transformer_engine/common/cast/dispatch/dequantize.cuh index 83494d953..81049a6bf 100644 --- a/transformer_engine/common/cast/dispatch/dequantize.cuh +++ b/transformer_engine/common/cast/dispatch/dequantize.cuh @@ -25,6 +25,11 @@ #include "../mxfp8/dequantize_mxfp8.cuh" #include "../mxfp8/group_dequantize_mxfp8.cuh" #include "../nvfp4/dequantize_nvfp4.cuh" +#ifdef __HIP_PLATFORM_AMD__ +#include "../../util/cuda_runtime.h" +#include "../../util/system.h" +#include "../fp8/rocm_quantize_fp8_tdm.cuh" +#endif // #ifdef __HIP_PLATFORM_AMD__ namespace transformer_engine { namespace dispatch { @@ -42,6 +47,11 @@ inline void dequantize_helper(const Tensor &input, Tensor *output, cudaStream_t NVTE_CHECK(is_fp8_dtype(input.dtype()), "Input must have FP8 type."); NVTE_CHECK(!is_fp8_dtype(output->dtype()), "Output must be in higher precision."); NVTE_CHECK(output->shape() == input.shape(), "Input and output shapes need to match."); +#ifdef __HIP_PLATFORM_AMD__ + if (fp8::dequantize_tdm(input, output, stream)) { + break; + } +#endif // #ifdef __HIP_PLATFORM_AMD__ fp8::dequantize(input, output, stream); break; } diff --git a/transformer_engine/common/cast/dispatch/gated.cuh b/transformer_engine/common/cast/dispatch/gated.cuh index d915a621c..d05ee673a 100644 --- a/transformer_engine/common/cast/dispatch/gated.cuh +++ b/transformer_engine/common/cast/dispatch/gated.cuh @@ -20,6 +20,11 @@ #include "../../utils.cuh" #include "../fp8/gated_fp8.cuh" #include "../mxfp8/gated_mxfp8.cuh" +#ifdef __HIP_PLATFORM_AMD__ +#include "../../util/cuda_runtime.h" +#include "../../util/system.h" +#include "../fp8/rocm_quantize_fp8_tdm.cuh" +#endif // #ifdef __HIP_PLATFORM_AMD__ namespace transformer_engine { namespace dispatch { @@ -49,6 +54,11 @@ void quantize_gated_fwd_helper(const NVTETensor nvte_input, NVTETensor nvte_outp switch (output->scaling_mode) { case NVTE_DELAYED_TENSOR_SCALING: { #ifdef __HIP_PLATFORM_AMD__ + Tensor dummy_grad_tensor; + if (fp8::quantize_gated_tdm( + input, dummy_grad_tensor, output, p, stream)) { + break; + } fp8::cast_gated_fwd(input, output, p, stream); #else const bool use_tma_kernels = (cols % 32 == 0) && is_supported_by_CC_100(); @@ -146,6 +156,10 @@ void quantize_gated_bwd_helper(const NVTETensor nvte_grad, const NVTETensor nvte switch (output->scaling_mode) { case NVTE_DELAYED_TENSOR_SCALING: { #ifdef __HIP_PLATFORM_AMD__ + if (fp8::quantize_gated_tdm(gated_input, grad, + output, p, stream)) { + break; + } fp8::cast_gated_bwd(gated_input, grad, output, p, stream); #else const bool use_tma_kernels = (cols % 32 == 0) && is_supported_by_CC_100(); diff --git a/transformer_engine/common/cast/fp8/rocm_quantize_fp8_tdm.cuh b/transformer_engine/common/cast/fp8/rocm_quantize_fp8_tdm.cuh index 50a6048a7..a4f10bc2c 100644 --- a/transformer_engine/common/cast/fp8/rocm_quantize_fp8_tdm.cuh +++ b/transformer_engine/common/cast/fp8/rocm_quantize_fp8_tdm.cuh @@ -8,12 +8,14 @@ // gfx1250 FP8 tensor-scaling quantize with the Tensor Data Mover: each wave moves its own // BAND_Y-row band of a 128x128 chunk global->LDS->global, and for the transposed output the waves -// assemble the chunk's transpose in LDS. The launchers return false where TDM does not apply, +// assemble the chunk's transpose in LDS. 2D block scaling quantizes each chunk with its own scale, +// and dequantize moves bands the same way. The launchers return false where TDM does not apply, // leaving the work to the generic kernels. Like the other rocm_ headers, this expects its includer // to have included the common headers, util/system.h and util/cuda_runtime.h. #include +#include "../../recipe/recipe_common.cuh" #include "../../util/math.h" #include "../../util/rocm_amax_slots.cuh" #include "../../util/rocm_device_utils.cuh" @@ -43,6 +45,8 @@ template using band_t = T[BAND_Y][CHUNK_DIM_X]; #ifdef __gfx1250__ +using common::cvt_from_float; + // Like the generic kernels, plain casts saturate as rocm_pack_4xfloat8 does, Inf included, and // fused ones as static_cast, which keeps Inf and NaN. template @@ -66,6 +70,30 @@ __device__ __forceinline__ uint32_t cvt_4x_to_fp8(const float x0, const float x1 return static_cast(packed); } } + +// Writes the lane's COLS_PER_LANE columns of a quad of packed rows of wave's band, transposed, into +// the chunk's transpose t. +template +__device__ __forceinline__ void store_transposed_quad(OType (&t)[N][BAND_Y][CHUNK_DIM_Y], + const uint32_t (&packed)[4], + const int lane_offset_X, const int wave, + const int quad) { + // Column j of the quad: byte j of each row + const uint32_t lo01 = __builtin_amdgcn_perm(packed[1], packed[0], 0x05010400); + const uint32_t hi01 = __builtin_amdgcn_perm(packed[1], packed[0], 0x07030602); + const uint32_t lo23 = __builtin_amdgcn_perm(packed[3], packed[2], 0x05010400); + const uint32_t hi23 = __builtin_amdgcn_perm(packed[3], packed[2], 0x07030602); + const uint32_t column[COLS_PER_LANE] = {__builtin_amdgcn_perm(lo23, lo01, 0x05040100), + __builtin_amdgcn_perm(lo23, lo01, 0x07060302), + __builtin_amdgcn_perm(hi23, hi01, 0x05040100), + __builtin_amdgcn_perm(hi23, hi01, 0x07060302)}; +#pragma unroll + for (int j = 0; j < COLS_PER_LANE; j++) { + const int x = lane_offset_X + j; + *reinterpret_cast(&t[x / BAND_Y][x % BAND_Y][wave * BAND_Y + quad * 4]) = + column[j]; + } +} #endif // __gfx1250__ // Quantizes chunk (chunk_Y, chunk_X) of a rows x cols tensor into output and, with TRANSPOSE, into @@ -148,21 +176,7 @@ __device__ __forceinline__ void quantize_chunk(const IType *input, const IType * *reinterpret_cast(&out[y][lane_offset_X]) = packed[i]; } if constexpr (TRANSPOSE) { - // Column j of the quad: byte j of each row - const uint32_t lo01 = __builtin_amdgcn_perm(packed[1], packed[0], 0x05010400); - const uint32_t hi01 = __builtin_amdgcn_perm(packed[1], packed[0], 0x07030602); - const uint32_t lo23 = __builtin_amdgcn_perm(packed[3], packed[2], 0x05010400); - const uint32_t hi23 = __builtin_amdgcn_perm(packed[3], packed[2], 0x07030602); - const uint32_t column[COLS_PER_LANE] = {__builtin_amdgcn_perm(lo23, lo01, 0x05040100), - __builtin_amdgcn_perm(lo23, lo01, 0x07060302), - __builtin_amdgcn_perm(hi23, hi01, 0x05040100), - __builtin_amdgcn_perm(hi23, hi01, 0x07060302)}; -#pragma unroll - for (int j = 0; j < COLS_PER_LANE; j++) { - const int x = lane_offset_X + j; - *reinterpret_cast(&t_sh[x / BAND_Y][x % BAND_Y][wave * BAND_Y + quad * 4]) = - column[j]; - } + store_transposed_quad(t_sh, packed, lane_offset_X, wave, quad); } } tdm::store_2d(output, out, rows, cols, cols, band_offset_Y, chunk_offset_X); @@ -286,6 +300,509 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } } +// FP8 tensor scaling of chunk (blockIdx.y, blockIdx.x) of a gated activation of input, whose rows +// hold the cols act inputs and then the cols gate inputs, or with IS_DGATED of its act and gate +// gradients for grad. The element math is that of gated_act_kernel and dgated_act_kernel. +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + quantize_gated_fp8_tdm_kernel(const IType *grad, const IType *input, OType *output, + OType *output_t, const float *scale_ptr, float *amax_ptr, + float *amax_slots, float *scale_inv_ptr, const size_t rows, + const size_t cols, const ParamOP p) { +#ifdef __gfx1250__ + // The act and gate gradient halves of each output row, or the activation alone + constexpr int HALVES = IS_DGATED ? 2 : 1; + // With TRANSPOSE, the packed outputs stay in registers until the consumed input bands can take + // them: a wave's output bands go into its act band, and the transposed chunk into the gate bands. + // Without it, a wave writes its rows into its own output bands as it computes them. + static_assert(HALVES * sizeof(band_t) <= sizeof(band_t)); + static_assert(HALVES * WAVES * BAND_Y * CHUNK_DIM_Y * sizeof(OType) <= + WAVES * sizeof(band_t)); + // The act, gate and, with IS_DGATED, grad bands of the waves + __shared__ alignas(128) band_t in_sh[IS_DGATED ? 3 : 2][WAVES]; + __shared__ alignas(128) band_t out_sh[TRANSPOSE ? 1 : HALVES][TRANSPOSE ? 1 : WAVES]; + auto &act_sh = in_sh[0]; + auto &gate_sh = in_sh[1]; + auto &grad_sh = in_sh[IS_DGATED ? 2 : 0]; + // A wave's amax goes after its output bands, or into its grad band if they fill the act band. + const auto wave_amax = [&](const int w) -> float * { + if constexpr (HALVES * sizeof(band_t) < sizeof(band_t)) { + return reinterpret_cast(reinterpret_cast(act_sh[w]) + + HALVES * sizeof(band_t)); + } else { + return reinterpret_cast(grad_sh[w]); + } + }; + + const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / WAVE_SIZE); + const int lane = threadIdx.x % WAVE_SIZE; + const int lane_offset_X = lane * COLS_PER_LANE; + const size_t chunk_offset_Y = blockIdx.y * CHUNK_DIM_Y; + const size_t chunk_offset_X = blockIdx.x * CHUNK_DIM_X; + const size_t band_offset_Y = chunk_offset_Y + wave * BAND_Y; + const bool has_band = band_offset_Y < rows; + const float scale = scale_ptr != nullptr ? *scale_ptr : 1.0f; + const auto out_band = [&](const int h) -> band_t & { + if constexpr (TRANSPOSE) { + return reinterpret_cast *>(act_sh[wave])[h]; + } else { + return out_sh[h][wave]; + } + }; + // Out-of-bounds elements are zero and every gated activation and gradient keeps them zero. + float amax = 0; + // As in quantize_chunk, the lanes start at different quads to spread the transposed writes. The + // packed rows of step s are those of quad (first_quad + s) % QUADS. + const int first_quad = TRANSPOSE ? lane % QUADS : 0; + uint32_t packed[QUADS][HALVES][4]; + + if (has_band) { + tdm::load_2d(act_sh[wave], input, rows, cols, 2 * cols, band_offset_Y, chunk_offset_X); + tdm::load_2d(gate_sh[wave], input + cols, rows, cols, 2 * cols, band_offset_Y, chunk_offset_X); + if constexpr (IS_DGATED) { + tdm::load_2d(grad_sh[wave], grad, rows, cols, cols, band_offset_Y, chunk_offset_X); + } + tdm::wait<0>(); + +#pragma unroll + for (int step = 0; step < QUADS; step++) { + const int quad = (first_quad + step) % QUADS; +#pragma unroll + for (int i = 0; i < 4; i++) { + const int y = quad * 4 + i; + Vec a, g, d; + a.load_from(&act_sh[wave][y][lane_offset_X]); + g.load_from(&gate_sh[wave][y][lane_offset_X]); + if constexpr (IS_DGATED) { + d.load_from(&grad_sh[wave][y][lane_offset_X]); + } + float elt[HALVES][COLS_PER_LANE]; +#pragma unroll + for (int j = 0; j < COLS_PER_LANE; j++) { + const float act_in = static_cast(a.data.elt[j]); + float gate_in = static_cast(g.data.elt[j]); + float dgate_in = 1.0f; + if constexpr (std::is_same_v) { + const float limit = p.limit; + dgate_in = gate_in <= limit && gate_in >= -limit; + gate_in = std::min(std::max(-limit, gate_in), limit) + p.glu_linear_offset; + } else if constexpr (std::is_same_v) { + if constexpr (IS_DGATED) { + dgate_in = dsitu_up(gate_in, p); + } + gate_in = situ_up(gate_in, p); + } + if constexpr (IS_DGATED) { + const float grad_val = static_cast(d.data.elt[j]); + elt[0][j] = DActOP(act_in, p) * grad_val * gate_in; + elt[1][j] = dgate_in * grad_val * ActOP(act_in, p); + } else { + elt[0][j] = ActOP(act_in, p) * gate_in; + } +#pragma unroll + for (int h = 0; h < HALVES; h++) { + __builtin_assume(amax >= 0); + amax = fmaxf(fabsf(elt[h][j]), amax); + } + } +#pragma unroll + for (int h = 0; h < HALVES; h++) { + packed[step][h][i] = cvt_4x_to_fp8(elt[h][0] * scale, elt[h][1] * scale, + elt[h][2] * scale, elt[h][3] * scale); + if constexpr (!TRANSPOSE) { + *reinterpret_cast(&out_band(h)[y][lane_offset_X]) = packed[step][h][i]; + } + } + } + } + + if constexpr (TRANSPOSE) { + asm volatile("" ::: "memory"); // act_sh is reused for the output bands +#pragma unroll + for (int step = 0; step < QUADS; step++) { + const int quad = (first_quad + step) % QUADS; +#pragma unroll + for (int i = 0; i < 4; i++) { +#pragma unroll + for (int h = 0; h < HALVES; h++) { + *reinterpret_cast(&out_band(h)[quad * 4 + i][lane_offset_X]) = + packed[step][h][i]; + } + } + } + } +#pragma unroll + for (int h = 0; h < HALVES; h++) { + tdm::store_2d(output + h * cols, out_band(h), rows, cols, HALVES * cols, band_offset_Y, + chunk_offset_X); + } + } + + asm volatile("" ::: "memory"); // the consumed bands are reused as float + if (amax_ptr != nullptr) { + const float reduced = rocm_subwarp_allreduce(amax, rocm_op::max{}); + if (lane == 0) { + *wave_amax(wave) = reduced; + } + } + if constexpr (TRANSPOSE) { + __syncthreads(); // every wave is done with the gate bands + asm volatile("" ::: "memory"); + auto &t_sh = *reinterpret_cast(gate_sh); + if (has_band) { +#pragma unroll + for (int step = 0; step < QUADS; step++) { +#pragma unroll + for (int h = 0; h < HALVES; h++) { + store_transposed_quad(t_sh[h], packed[step][h], lane_offset_X, wave, + (first_quad + step) % QUADS); + } + } + } + __syncthreads(); + // Rows past the tensor's are garbage in t_sh but outside the transpose's columns. + const size_t t_offset_Y = chunk_offset_X + wave * BAND_Y; + if (t_offset_Y < cols) { +#pragma unroll + for (int h = 0; h < HALVES; h++) { + tdm::store_2d(output_t + h * cols * rows, t_sh[h][wave], cols, rows, rows, t_offset_Y, + chunk_offset_Y); + } + } + } else { + __syncthreads(); + } + if (amax_ptr != nullptr && threadIdx.x == 0) { + float block_amax = 0; +#pragma unroll + for (int w = 0; w < WAVES; w++) { + block_amax = fmaxf(block_amax, *wave_amax(w)); + } + float *const amax_dst = + rocm_amax_target(amax_ptr, amax_slots, blockIdx.y * gridDim.x + blockIdx.x); + if (block_amax > *amax_dst) { + rocm_atomicMaxFloat(amax_dst, block_amax); + } + } + if (blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0 && scale_inv_ptr != nullptr) { + *scale_inv_ptr = __frcp_rn(scale); + } + tdm::wait<0>(); +#endif // __gfx1250__ +} + +// FP8 tensor-scaling dequantize of chunk (blockIdx.y, blockIdx.x): per step, 8 lanes cover a row of +// the wave's band with 16 elements each. +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + dequantize_fp8_tdm_kernel(const IType *input, OType *output, const float *scale_inv_ptr, + const size_t rows, const size_t cols) { +#ifdef __gfx1250__ + constexpr int ELEMS = 16; + constexpr int LANES_PER_ROW = CHUNK_DIM_X / ELEMS; // 8 + constexpr int ROWS_PER_STEP = WAVE_SIZE / LANES_PER_ROW; // 4 + __shared__ alignas(128) band_t in_sh[WAVES]; + __shared__ alignas(128) band_t out_sh[WAVES]; + + const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / WAVE_SIZE); + const int lane = threadIdx.x % WAVE_SIZE; + const size_t band_offset_Y = blockIdx.y * CHUNK_DIM_Y + wave * BAND_Y; + const size_t chunk_offset_X = blockIdx.x * CHUNK_DIM_X; + if (band_offset_Y >= rows) { + return; + } + tdm::load_2d(in_sh[wave], input, rows, cols, cols, band_offset_Y, chunk_offset_X); + const float scale_inv = *scale_inv_ptr; + tdm::wait<0>(); + const int lane_offset_X = (lane % LANES_PER_ROW) * ELEMS; +#pragma unroll + for (int step = 0; step < BAND_Y / ROWS_PER_STEP; step++) { + const int y = step * ROWS_PER_STEP + lane / LANES_PER_ROW; + Vec in; + Vec out; + in.load_from(&in_sh[wave][y][lane_offset_X]); +#pragma unroll + for (int i = 0; i < ELEMS; i++) { + out.data.elt[i] = cvt_from_float(static_cast(in.data.elt[i]) * scale_inv); + } + out.store_to(&out_sh[wave][y][lane_offset_X]); + } + tdm::store_2d(output, out_sh[wave], rows, cols, cols, band_offset_Y, chunk_offset_X); + tdm::wait<0>(); +#endif // __gfx1250__ +} + +// FP8 2D block scaling of chunk (blockIdx.y, blockIdx.x) with the scale from the chunk's amax into +// output, if given, and, with TRANSPOSE, into output_t. The scale inverse goes where the generic +// block_scaled_cast_transpose_kernel_notaligned puts it. +template +__device__ __forceinline__ void quantize_block2d_chunk( + const IType *input, OType *output, OType *output_t, float *scale_inv, float *scale_inv_t, + const size_t rows, const size_t cols, const size_t scale_stride_x, const size_t scale_stride_y, + const size_t scale_t_stride_x, const size_t scale_t_stride_y, const float epsilon, + const bool pow_2_scaling) { +#ifdef __gfx1250__ + __shared__ alignas(128) band_t in_sh[WAVES]; + __shared__ alignas(128) band_t out_sh[WAVES]; + __shared__ alignas(128) OType t_sh[TRANSPOSE ? WAVES : 1][BAND_Y][CHUNK_DIM_Y]; + + const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / WAVE_SIZE); + const int lane = threadIdx.x % WAVE_SIZE; + const int lane_offset_X = lane * COLS_PER_LANE; + const size_t chunk_offset_Y = blockIdx.y * CHUNK_DIM_Y; + const size_t chunk_offset_X = blockIdx.x * CHUNK_DIM_X; + const size_t band_offset_Y = chunk_offset_Y + wave * BAND_Y; + const bool has_band = band_offset_Y < rows; + auto &in = in_sh[wave]; + auto &out = out_sh[wave]; + + // Out-of-bounds elements are zero in LDS and leave the amax unchanged. + float amax = 0; + if (has_band) { + tdm::load_2d(in, input, rows, cols, cols, band_offset_Y, chunk_offset_X); + tdm::wait<0>(); +#pragma unroll + for (int y = 0; y < BAND_Y; y++) { + Vec x; + x.load_from(&in[y][lane_offset_X]); +#pragma unroll + for (int j = 0; j < COLS_PER_LANE; j++) { + __builtin_assume(amax >= 0); + amax = fmaxf(fabsf(static_cast(x.data.elt[j])), amax); + } + } + } + // The waves exchange their amaxes through their not yet written out_sh bands. + amax = rocm_subwarp_allreduce(amax, rocm_op::max{}); + asm volatile("" ::: "memory"); // out_sh is used as float + if (lane == 0) { + *reinterpret_cast(out) = amax; + } + __syncthreads(); + float block_amax = 0; +#pragma unroll + for (int w = 0; w < WAVES; w++) { + block_amax = fmaxf(block_amax, *reinterpret_cast(out_sh[w])); + } + __syncthreads(); + asm volatile("" ::: "memory"); + const float scale = compute_scale_from_types(block_amax, epsilon, pow_2_scaling); + if (threadIdx.x == 0) { + const float s_inv = 1.0f / scale; + if (scale_inv != nullptr) { + scale_inv[blockIdx.y * scale_stride_y + blockIdx.x * scale_stride_x] = s_inv; + } + if constexpr (TRANSPOSE) { + scale_inv_t[blockIdx.x * scale_t_stride_y + blockIdx.y * scale_t_stride_x] = s_inv; + } + } + + if (has_band) { + // As in quantize_chunk, the lanes start at different quads to spread the transposed writes. + const int first_quad = TRANSPOSE ? lane % QUADS : 0; +#pragma unroll + for (int step = 0; step < QUADS; step++) { + const int quad = (first_quad + step) % QUADS; + uint32_t packed[4]; +#pragma unroll + for (int i = 0; i < 4; i++) { + const int y = quad * 4 + i; + Vec x; + x.load_from(&in[y][lane_offset_X]); + packed[i] = cvt_4x_to_fp8( + static_cast(x.data.elt[0]) * scale, static_cast(x.data.elt[1]) * scale, + static_cast(x.data.elt[2]) * scale, static_cast(x.data.elt[3]) * scale); + if (output != nullptr) { + *reinterpret_cast(&out[y][lane_offset_X]) = packed[i]; + } + } + if constexpr (TRANSPOSE) { + store_transposed_quad(t_sh, packed, lane_offset_X, wave, quad); + } + } + if (output != nullptr) { + tdm::store_2d(output, out, rows, cols, cols, band_offset_Y, chunk_offset_X); + } + } + if constexpr (TRANSPOSE) { + __syncthreads(); + // Rows past the tensor's are garbage in t_sh but outside the transpose's columns. + const size_t t_offset_Y = chunk_offset_X + wave * BAND_Y; + if (t_offset_Y < cols) { + tdm::store_2d(output_t, t_sh[wave], cols, rows, rows, t_offset_Y, chunk_offset_Y); + } + } + tdm::wait<0>(); +#endif // __gfx1250__ +} + +// FP8 1D block scaling of chunk (blockIdx.y, blockIdx.x): each chunk row with its own scale into +// output, if given, and, with TRANSPOSE, each chunk column with its own scale into output_t. The +// scale inverses go where the generic block_scaled_1d_cast_transpose_kernel puts them. +template +__device__ __forceinline__ void quantize_block1d_chunk( + const IType *input, OType *output, OType *output_t, float *scale_inv, float *scale_inv_t, + const size_t rows, const size_t cols, const size_t scale_stride_x, const size_t scale_stride_y, + const size_t scale_t_stride_x, const size_t scale_t_stride_y, const float epsilon, + const bool pow_2_scaling) { +#ifdef __gfx1250__ + constexpr int ELEMS = 16; + constexpr int LANES_PER_ROW = CHUNK_DIM_X / ELEMS; // 8 + constexpr int ROWS_PER_STEP = WAVE_SIZE / LANES_PER_ROW; // 4 + __shared__ alignas(128) band_t in_sh[WAVES]; + __shared__ alignas(128) band_t out_sh[WAVES]; + __shared__ alignas(128) OType t_sh[TRANSPOSE ? WAVES : 1][BAND_Y][CHUNK_DIM_Y]; + + const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / WAVE_SIZE); + const int lane = threadIdx.x % WAVE_SIZE; + const size_t chunk_offset_Y = blockIdx.y * CHUNK_DIM_Y; + const size_t chunk_offset_X = blockIdx.x * CHUNK_DIM_X; + const size_t band_offset_Y = chunk_offset_Y + wave * BAND_Y; + const bool has_band = band_offset_Y < rows; + auto &in = in_sh[wave]; + auto &out = out_sh[wave]; + + if (has_band) { + tdm::load_2d(in, input, rows, cols, cols, band_offset_Y, chunk_offset_X); + tdm::wait<0>(); + } + + // For the transpose, a lane owns COLS_PER_LANE columns. Out-of-bounds elements are zero in LDS + // and leave the amaxes unchanged. + const int col_offset_X = lane * COLS_PER_LANE; + float col_scale[COLS_PER_LANE]; + if constexpr (TRANSPOSE) { + Vec col_amax; + col_amax.clear(); + if (has_band) { +#pragma unroll + for (int y = 0; y < BAND_Y; y++) { + Vec x; + x.load_from(&in[y][col_offset_X]); +#pragma unroll + for (int j = 0; j < COLS_PER_LANE; j++) { + __builtin_assume(col_amax.data.elt[j] >= 0); + col_amax.data.elt[j] = + fmaxf(col_amax.data.elt[j], fabsf(static_cast(x.data.elt[j]))); + } + } + } + // The waves exchange their band's column amaxes through their not yet written out_sh bands. + asm volatile("" ::: "memory"); // out_sh is used as float + col_amax.store_to(reinterpret_cast(out) + col_offset_X); + __syncthreads(); +#pragma unroll + for (int w = 1; w < WAVES; w++) { + Vec other; + other.load_from(reinterpret_cast(out_sh[(wave + w) % WAVES]) + col_offset_X); +#pragma unroll + for (int j = 0; j < COLS_PER_LANE; j++) { + col_amax.data.elt[j] = fmaxf(col_amax.data.elt[j], other.data.elt[j]); + } + } + __syncthreads(); + asm volatile("" ::: "memory"); +#pragma unroll + for (int j = 0; j < COLS_PER_LANE; j++) { + col_scale[j] = compute_scale_from_types(col_amax.data.elt[j], epsilon, + pow_2_scaling); + const size_t col = chunk_offset_X + col_offset_X + j; + if (wave == 0 && col < cols) { + scale_inv_t[col * scale_t_stride_y + blockIdx.y * scale_t_stride_x] = 1.0 / col_scale[j]; + } + } + } + + if (has_band && output != nullptr) { + // Per step, 8 lanes cover a row with 16 elements each. + const int lane_offset_X = (lane % LANES_PER_ROW) * ELEMS; +#pragma unroll + for (int step = 0; step < BAND_Y / ROWS_PER_STEP; step++) { + const int y = step * ROWS_PER_STEP + lane / LANES_PER_ROW; + Vec x; + x.load_from(&in[y][lane_offset_X]); + float amax = 0; +#pragma unroll + for (int i = 0; i < ELEMS; i++) { + __builtin_assume(amax >= 0); + amax = fmaxf(amax, fabsf(static_cast(x.data.elt[i]))); + } + amax = rocm_subwarp_allreduce(amax, rocm_op::max{}); + const float scale = compute_scale_from_types(amax, epsilon, pow_2_scaling); + const size_t row = band_offset_Y + y; + if (lane % LANES_PER_ROW == 0 && row < rows) { + scale_inv[row * scale_stride_y + blockIdx.x * scale_stride_x] = 1.0 / scale; + } + uint32_t packed[ELEMS / 4]; +#pragma unroll + for (int i = 0; i < ELEMS / 4; i++) { + packed[i] = cvt_4x_to_fp8(static_cast(x.data.elt[4 * i]) * scale, + static_cast(x.data.elt[4 * i + 1]) * scale, + static_cast(x.data.elt[4 * i + 2]) * scale, + static_cast(x.data.elt[4 * i + 3]) * scale); + } + *reinterpret_cast(&out[y][lane_offset_X]) = + make_uint4(packed[0], packed[1], packed[2], packed[3]); + } + tdm::store_2d(output, out, rows, cols, cols, band_offset_Y, chunk_offset_X); + } + + if constexpr (TRANSPOSE) { + if (has_band) { + // As in quantize_chunk, the lanes start at different quads to spread the transposed writes. + const int first_quad = lane % QUADS; +#pragma unroll + for (int step = 0; step < QUADS; step++) { + const int quad = (first_quad + step) % QUADS; + uint32_t packed[4]; +#pragma unroll + for (int i = 0; i < 4; i++) { + Vec x; + x.load_from(&in[quad * 4 + i][col_offset_X]); + packed[i] = cvt_4x_to_fp8(static_cast(x.data.elt[0]) * col_scale[0], + static_cast(x.data.elt[1]) * col_scale[1], + static_cast(x.data.elt[2]) * col_scale[2], + static_cast(x.data.elt[3]) * col_scale[3]); + } + store_transposed_quad(t_sh, packed, col_offset_X, wave, quad); + } + } + __syncthreads(); + // Rows past the tensor's are garbage in t_sh but outside the transpose's columns. + const size_t t_offset_Y = chunk_offset_X + wave * BAND_Y; + if (t_offset_Y < cols) { + tdm::store_2d(output_t, t_sh[wave], cols, rows, rows, t_offset_Y, chunk_offset_Y); + } + } + tdm::wait<0>(); +#endif // __gfx1250__ +} + +// FP8 block scaling of chunk (blockIdx.y, blockIdx.x) with BLOCK_DIM 1 (1x128 and 128x1 tiles) or 2 +// (128x128 tiles). +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + quantize_block_fp8_tdm_kernel(const IType *input, OType *output, OType *output_t, + float *scale_inv, float *scale_inv_t, const size_t rows, + const size_t cols, const size_t scale_stride_x, + const size_t scale_stride_y, const size_t scale_t_stride_x, + const size_t scale_t_stride_y, const float epsilon, + const bool pow_2_scaling, const float *noop) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + if constexpr (BLOCK_DIM == 1) { + quantize_block1d_chunk(input, output, output_t, scale_inv, scale_inv_t, rows, cols, + scale_stride_x, scale_stride_y, scale_t_stride_x, + scale_t_stride_y, epsilon, pow_2_scaling); + } else { + quantize_block2d_chunk(input, output, output_t, scale_inv, scale_inv_t, rows, cols, + scale_stride_x, scale_stride_y, scale_t_stride_x, + scale_t_stride_y, epsilon, pow_2_scaling); + } +} + } // namespace tdm_kernel // FP8 tensor scaling of input, after an activation or its derivative, into output's rowwise data @@ -299,10 +816,14 @@ bool quantize_tdm(const Tensor &input, const Tensor *act_input, const Tensor *no const bool transpose = output->has_columnwise_data(); const auto [rows, cols] = input.flat_2d_dims(); const bool has_noop = noop != nullptr && noop->data.dptr != nullptr; + // Up to this size, the amax finalize launch costs a plain cast-transpose more than TDM saves. + constexpr size_t MAX_GENERIC_AMAX_CT_ELEMS = size_t{1} << 25; + const bool small_amax_ct = !FUSED && transpose && output->amax.dptr != nullptr && + rows * cols <= MAX_GENERIC_AMAX_CT_ELEMS; // The transpose's rows are rows long. if (!output->has_data() || !is_tensor_scaling(output->scaling_mode) || is_fp8_dtype(input.dtype()) || !is_fp8_dtype(output->dtype()) || rows == 0 || - (transpose && rows % 16 != 0) || (FUSED && has_noop) || + (transpose && rows % 16 != 0) || (FUSED && has_noop) || small_amax_ct || (IS_DACT && (act_input == nullptr || !act_input->has_data())) || (IS_DBIAS && (dbias == nullptr || workspace == nullptr)) || !use_tdm(rows, cols, @@ -382,9 +903,11 @@ bool quantize_tdm(const Tensor &input, const Tensor *act_input, const Tensor *no // FP8 cast and transpose of each input into its output, whose rows past the input's (with // valid_num_rows, valid_num_rows[i] of them) are zero padding. TDM only if every tensor qualifies. -inline bool multi_cast_transpose_tdm(const std::vector &input_list, - const std::vector &output_list, - const int *valid_num_rows, hipStream_t stream) { +// A template, like dequantize_tdm, so that only the files calling it compile its kernels. +template +bool multi_cast_transpose_tdm(const std::vector &input_list, + const std::vector &output_list, const int *valid_num_rows, + hipStream_t stream) { using namespace tdm_kernel; const size_t num_tensors = input_list.size(); for (size_t i = 0; i < num_tensors; i++) { @@ -435,6 +958,115 @@ inline bool multi_cast_transpose_tdm(const std::vector &input_list, return true; } +// FP8 tensor scaling of a gated activation of gated_input or, with IS_BWD, of its gradient for grad +// into output's rowwise data and, if allocated, its transposed columnwise data. +template +bool quantize_gated_tdm(const Tensor &gated_input, const Tensor &grad, Tensor *output, + const ParamOP &p, hipStream_t stream) { + using namespace tdm_kernel; + const bool transpose = output->has_columnwise_data(); + const size_t rows = gated_input.flat_first_dim(); + const size_t cols = gated_input.flat_last_dim() / 2; + // The transpose's rows are rows long. + if (!output->has_data() || !is_fp8_dtype(output->dtype()) || + is_fp8_dtype(gated_input.dtype()) || rows == 0 || (transpose && rows % 16 != 0) || + !use_tdm(rows, cols, + {gated_input.data.dptr, IS_BWD ? grad.data.dptr : nullptr, output->data.dptr, + output->columnwise_data.dptr})) { + return false; + } + if (transpose) { + NVTE_CHECK(output->data.dtype == output->columnwise_data.dtype, + "Cast and transposed output types must match."); + } + float *const amax = reinterpret_cast(output->amax.dptr); + const dim3 grid(DIVUP(cols, CHUNK_DIM_X), DIVUP(rows, CHUNK_DIM_Y)); + float *const amax_slots = amax != nullptr ? rocm_amax_slots(grid.x * grid.y, stream) : nullptr; + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + gated_input.dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output->dtype(), OType, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + transpose, TRANSPOSE, + quantize_gated_fp8_tdm_kernel + <<>>( + IS_BWD ? reinterpret_cast(grad.data.dptr) : nullptr, + reinterpret_cast(gated_input.data.dptr), + reinterpret_cast(output->data.dptr), + reinterpret_cast(output->columnwise_data.dptr), + reinterpret_cast(output->scale.dptr), amax, amax_slots, + reinterpret_cast(output->scale_inv.dptr), rows, cols, p);))); // NOLINT(*) + NVTE_CHECK_CUDA(hipGetLastError()); + if (amax_slots != nullptr) { + rocm_amax_finalize(amax_slots, amax, nullptr, stream); + } + return true; +} + +// FP8 tensor-scaling dequantize of input into output. +template +bool dequantize_tdm(const Tensor &input, Tensor *output, hipStream_t stream) { + using namespace tdm_kernel; + const auto [rows, cols] = input.flat_2d_dims(); + if (!use_tdm(rows, cols, {input.data.dptr, output->data.dptr})) { + return false; + } + const dim3 grid(DIVUP(cols, CHUNK_DIM_X), DIVUP(rows, CHUNK_DIM_Y)); + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + input.data.dtype, IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + output->data.dtype, OType, + dequantize_fp8_tdm_kernel<<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(output->data.dptr), + reinterpret_cast(input.scale_inv.dptr), rows, cols);)); // NOLINT(*) + NVTE_CHECK_CUDA(hipGetLastError()); + return true; +} + +// FP8 block scaling of the rows x cols input with BLOCK_DIM 1 (1x128 and 128x1 tiles) or 2 (128x128 +// tiles) into output, with rowwise, and into its transpose output_t, with return_transpose. Takes +// the scale-inverse strides of quantize_transpose_{vector,square}_blockwise. +template +bool quantize_blockwise_tdm(const SimpleTensor &input, const SimpleTensor &output, + const SimpleTensor &output_t, const SimpleTensor &scale_inv, + const SimpleTensor &scale_inv_t, const size_t rows, const size_t cols, + const bool rowwise, const bool return_transpose, + const size_t scale_stride_x, const size_t scale_stride_y, + const size_t scale_t_stride_x, const size_t scale_t_stride_y, + const float epsilon, const bool pow_2_scale, const float *noop, + hipStream_t stream) { + using namespace tdm_kernel; + // The transpose's rows are rows long. + if (is_fp8_dtype(input.dtype) || rows == 0 || cols == 0 || + (return_transpose && rows % 16 != 0) || + !use_tdm(rows, cols, + {input.dptr, rowwise ? output.dptr : nullptr, + return_transpose ? output_t.dptr : nullptr})) { + return false; + } + const DType out_dtype = rowwise ? output.dtype : output_t.dtype; + const dim3 grid(DIVUP(cols, CHUNK_DIM_X), DIVUP(rows, CHUNK_DIM_Y)); + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input.dtype, IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + out_dtype, OType, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + return_transpose, TRANSPOSE, + quantize_block_fp8_tdm_kernel + <<>>( + reinterpret_cast(input.dptr), + rowwise ? reinterpret_cast(output.dptr) : nullptr, + reinterpret_cast(output_t.dptr), + reinterpret_cast(scale_inv.dptr), + reinterpret_cast(scale_inv_t.dptr), rows, cols, scale_stride_x, + scale_stride_y, scale_t_stride_x, scale_t_stride_y, epsilon, pow_2_scale, + noop);))); // NOLINT(*) + NVTE_CHECK_CUDA(hipGetLastError()); + return true; +} + } // namespace fp8 } // namespace dispatch } // namespace transformer_engine diff --git a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh index 2ebfc9570..2da642954 100644 --- a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh @@ -36,6 +36,7 @@ #include "../../common.h" #include "../../util/ptx.cuh" #include "../../utils.cuh" +#include "./dequantize_mxfp8.cuh" #include "./rocm_vectorized_2d.cuh" #endif // #ifndef __HIP_PLATFORM_AMD__ @@ -500,12 +501,13 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output, // Row strides must be a whole number of vectors; VECTOR_WIDTH is 16 elements. constexpr size_t VECTOR_ELEMS = 16; + std::vector last_dims; if (last_dims_ptr == nullptr) { NVTE_CHECK(last_logical_dim % VECTOR_ELEMS == 0, "Grouped dequantize requires the last dimension to be divisible by ", VECTOR_ELEMS, "; got ", last_logical_dim, "."); } else { - const std::vector last_dims = read_dims_to_host(last_dims_ptr); + last_dims = read_dims_to_host(last_dims_ptr); for (size_t t = 0; t < num_tensors; t++) { NVTE_CHECK(static_cast(last_dims[t]) % VECTOR_ELEMS == 0, "Grouped dequantize requires each per-tensor last dimension to be divisible by ", @@ -513,6 +515,10 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output, } } + if (dequantize_kernel::group_dequantize_tdm(*input, output, last_dims, stream)) { + return; + } + const size_t scale_dim_X_rowwise = use_rowwise_scaling ? 32 : 1; const size_t scale_dim_Y_colwise = use_colwise_scaling ? 32 : 1; diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 82ea039e7..cfb91d4be 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -1193,8 +1193,9 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations } } - if (group_quantize_tdm(input, activations, noop, output, - dbias, workspace, stream)) { + const bool use_2d_quantization = quant_config != nullptr && quant_config->mxfp8_2d_quantization; + if (group_quantize_tdm( + input, activations, noop, output, dbias, workspace, use_2d_quantization, stream)) { return; } diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index d4524a8c8..4072d9bbb 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -694,9 +694,9 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, const auto [rows, cols] = input.flat_2d_dims(); #ifdef __HIP_PLATFORM_AMD__ - if (!use_2d_quantization && - quantize_tdm(input, act_input, noop, output, dbias, - workspace, stream)) { + if (quantize_tdm(input, act_input, noop, output, dbias, + workspace, use_2d_quantization, + stream)) { return; } diff --git a/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh b/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh index 959ee97ae..a781e2e8c 100644 --- a/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh +++ b/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8_tdm.cuh @@ -15,13 +15,14 @@ constexpr size_t TDM_WAVE_SIZE = 32; constexpr size_t TDM_WAVES = THREADS_PER_CHUNK / TDM_WAVE_SIZE; // 4 constexpr size_t TDM_BAND_Y = CHUNK_DIM_Y / TDM_WAVES; // 32 +// Dequantizes chunk (chunk_Y, chunk_X) of a rows x cols tensor. template -__global__ void __launch_bounds__(THREADS_PER_CHUNK) - dequantize_mxfp8_tdm_kernel(const IType *input_ptr, OType *output_ptr, - const e8m0_t *const scales_ptr, const size_t rows, - const size_t cols, const size_t scales_stride, - const size_t mx_swizzle_padded_dim) { +__device__ __forceinline__ void dequantize_chunk(const IType *input_ptr, OType *output_ptr, + const e8m0_t *const scales_ptr, const size_t rows, + const size_t cols, const size_t scales_stride, + const size_t mx_swizzle_padded_dim, + const size_t chunk_Y, const size_t chunk_X) { #ifdef __gfx1250__ constexpr bool USE_ROWWISE_SCALING = SCALE_DIM_X > 1; @@ -30,8 +31,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / TDM_WAVE_SIZE); const int lane = threadIdx.x % TDM_WAVE_SIZE; - const size_t band_offset_Y = blockIdx.y * CHUNK_DIM_Y + wave * TDM_BAND_Y; - const size_t chunk_offset_X = blockIdx.x * CHUNK_DIM_X; + const size_t band_offset_Y = chunk_Y * CHUNK_DIM_Y + wave * TDM_BAND_Y; + const size_t chunk_offset_X = chunk_X * CHUNK_DIM_X; if (band_offset_Y >= rows) { return; } @@ -114,6 +115,40 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) #endif // __gfx1250__ } +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + dequantize_mxfp8_tdm_kernel(const IType *input_ptr, OType *output_ptr, + const e8m0_t *const scales_ptr, const size_t rows, + const size_t cols, const size_t scales_stride, + const size_t mx_swizzle_padded_dim) { + dequantize_chunk( + input_ptr, output_ptr, scales_ptr, rows, cols, scales_stride, mx_swizzle_padded_dim, + blockIdx.y, blockIdx.x); +} + +// A block per chunk of find_group_chunk, of tensors whose compact scales stack in tensor order. +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + grouped_dequantize_mxfp8_tdm_kernel(const IType *input_ptr, OType *output_ptr, + const e8m0_t *const scales_ptr, + const size_t first_logical_dim, const size_t num_tensors, + const int64_t *offsets_ptr, const int64_t *first_dims_ptr, + const int64_t *last_dims_ptr) { + constexpr bool USE_ROWWISE_SCALING = SCALE_DIM_X > 1; + const auto chunk = tdm_kernel::find_group_chunk( + blockIdx.x, num_tensors, first_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr); + if (chunk.rows == 0) { + return; + } + // Both dimensions are divisible by the scale blocks. + const size_t scales_offset = chunk.offset / (USE_ROWWISE_SCALING ? SCALE_DIM_X : SCALE_DIM_Y); + dequantize_chunk( + input_ptr + chunk.offset, output_ptr + chunk.offset, scales_ptr + scales_offset, chunk.rows, + chunk.cols, USE_ROWWISE_SCALING ? chunk.cols / SCALE_DIM_X : chunk.cols, 0, chunk.chunk_Y, + chunk.chunk_X); +} + // Returns false, leaving the work to the generic kernel, where TDM does not apply. The scale // arguments are the ones of the generic kernel. inline bool dequantize_tdm(const Tensor &input, const SimpleTensor &input_data, Tensor *output, @@ -143,3 +178,58 @@ inline bool dequantize_tdm(const Tensor &input, const SimpleTensor &input_data, NVTE_CHECK_CUDA(hipGetLastError()); return true; } + +// Grouped dequantize, whose compact scales stack in tensor order. With a common last dimension, the +// group is one tensor. Otherwise, host_last_dims has the last dimensions, and the chunks of +// find_group_chunk need them and the first dimensions divisible by 128, as upstream requires of +// varying first dimensions, which only the device has. +inline bool group_dequantize_tdm(const GroupedTensor &input, GroupedTensor *output, + const std::vector &host_last_dims, hipStream_t stream) { + const bool rowwise = input.has_data(); + const SimpleTensor &input_data = rowwise ? input.data : input.columnwise_data; + const size_t first_logical_dim = input.logical_shape.data[0]; + const size_t last_logical_dim = input.logical_shape.data[1]; + const bool same_last_dim = input.all_same_last_dim(); + const auto chunk_aligned = [](const int64_t dim) { return dim % CHUNK_DIM_X == 0; }; + if (!tdm_kernel::use_tdm(first_logical_dim, last_logical_dim, + {input_data.dptr, output->data.dptr}) || + (!same_last_dim && + (input.tensor_offsets.dptr == nullptr || + (input.all_same_first_dim() && first_logical_dim % CHUNK_DIM_Y != 0) || + !std::all_of(host_last_dims.begin(), host_last_dims.end(), chunk_aligned)))) { + return false; + } + const auto *const scales_ptr = reinterpret_cast( + rowwise ? input.scale_inv.dptr : input.columnwise_scale_inv.dptr); + const auto *const offsets_ptr = reinterpret_cast(input.tensor_offsets.dptr); + const auto *const first_dims_ptr = reinterpret_cast(input.first_dims.dptr); + const auto *const last_dims_ptr = reinterpret_cast(input.last_dims.dptr); + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + (rowwise ? 1 : 32), SCALE_DIM_Y, + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + (rowwise ? 32 : 1), SCALE_DIM_X, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + input.dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + output->dtype(), OType, + const auto *const input_ptr = reinterpret_cast(input_data.dptr); + auto *const output_ptr = reinterpret_cast(output->data.dptr); + if (same_last_dim) { + const dim3 grid(DIVUP(last_logical_dim, CHUNK_DIM_X), + DIVUP(first_logical_dim, CHUNK_DIM_Y)); + dequantize_mxfp8_tdm_kernel + <<>>( + input_ptr, output_ptr, scales_ptr, first_logical_dim, + last_logical_dim, + rowwise ? DIVUP(last_logical_dim, SCALE_DIM_X) : last_logical_dim, 0); + } else { + const dim3 grid( + DIVUP(first_logical_dim * last_logical_dim, CHUNK_DIM_Y * CHUNK_DIM_X)); + grouped_dequantize_mxfp8_tdm_kernel + <<>>( + input_ptr, output_ptr, scales_ptr, first_logical_dim, + input.num_tensors, offsets_ptr, first_dims_ptr, last_dims_ptr); + })))); // NOLINT(*) + NVTE_CHECK_CUDA(hipGetLastError()); + return true; +} diff --git a/transformer_engine/common/cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh b/transformer_engine/common/cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh index 53b6684d9..f1c5d7724 100644 --- a/transformer_engine/common/cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh +++ b/transformer_engine/common/cast/mxfp8/rocm_quantize_mxfp8_tdm.cuh @@ -21,6 +21,8 @@ namespace dispatch { namespace mxfp8 { namespace tdm_kernel { +using common::find_group_chunk; +using common::find_tensor; using common::reduce_dbias_tdm; using common::use_tdm; @@ -31,6 +33,7 @@ constexpr size_t WAVE_SIZE = 32; constexpr size_t WAVES = THREADS_PER_CHUNK / WAVE_SIZE; // 4 constexpr size_t BAND_Y = CHUNK_DIM_Y / WAVES; // 32 constexpr size_t COLS_PER_LANE = CHUNK_DIM_X / WAVE_SIZE; // 4 +constexpr size_t TILE_2D = 32; // 2D scaling shares a scale per TILE_2D x TILE_2D tile constexpr size_t MULTI_MAX_TENSORS = 1024; static_assert(CHUNK_DIM_Y == common::tdm_dbias_kernel::CHUNK_DIM_Y, "A dbias partial row covers a chunk"); @@ -39,15 +42,7 @@ template using band_t = T[BAND_Y][CHUNK_DIM_X]; #ifdef __gfx1250__ -// hip_bfloat16 rounds with per-element branches; __bf16 converts with one instruction. -template -__device__ __forceinline__ T cvt_from_float(const float x) { - if constexpr (std::is_same_v) { - return __builtin_bit_cast(T, static_cast<__bf16>(x)); - } else { - return static_cast(x); - } -} +using common::cvt_from_float; // The e8m0 scale rounds up, so scaled finite values stay within the FP8 range and need no clamp: // this matches static_cast bit for bit, including Inf and NaN, at two values per instruction. @@ -65,15 +60,22 @@ __device__ __forceinline__ uint32_t cvt_4x_to_fp8(const float x0, const float x1 return static_cast(packed); } +// With 2D scaling, byte t of tile_exponents is the biased exponent of the band's tile t; this +// returns the one of band column x. +__device__ __forceinline__ e8m0_t tile_exponent(const uint32_t tile_exponents, const int x) { + return static_cast(tile_exponents >> (x / static_cast(TILE_2D) * 8)); +} + // Rowwise MXFP8 of a band in LDS: per step, 8 lanes cover a row with 16 elements each and a lane -// pair shares a scale. Writes the scales of the first valid_rows rows and valid_scales scale -// columns, and returns the lane's amax. -template +// pair shares a scale, or with IS_2D the one of its tile. Writes the scales of the first valid_rows +// rows and valid_scales scale columns, and returns the lane's amax, which IS_2D leaves at 0. +template __device__ __forceinline__ float quantize_band_rowwise(const band_t &in, band_t &out, e8m0_t *const scales, const size_t scale_stride, const size_t valid_rows, - const size_t valid_scales) { + const size_t valid_scales, + const uint32_t tile_exponents = 0) { constexpr size_t ELEMS = 16; constexpr size_t LANES_PER_ROW = CHUNK_DIM_X / ELEMS; // 8 constexpr size_t ROWS_PER_STEP = WAVE_SIZE / LANES_PER_ROW; // 4 @@ -88,15 +90,19 @@ __device__ __forceinline__ float quantize_band_rowwise(const band_t &in, const size_t y = step * ROWS_PER_STEP + lane_offset_Y; Vec x; x.load_from(&in[y][lane_offset_X]); - float thread_amax = 0; + e8m0_t biased_exponent; + if constexpr (IS_2D) { + biased_exponent = tile_exponent(tile_exponents, lane_offset_X); + } else { + float thread_amax = 0; #pragma unroll - for (int j = 0; j < ELEMS; j++) { - thread_amax = fmaxf(thread_amax, fabsf(static_cast(x.data.elt[j]))); + for (int j = 0; j < ELEMS; j++) { + thread_amax = fmaxf(thread_amax, fabsf(static_cast(x.data.elt[j]))); + } + amax = fmaxf(amax, thread_amax); + const float scale_amax = rocm_subwarp_allreduce(thread_amax, rocm_op::max{}); + biased_exponent = ptx::float_to_e8m0(scale_amax * Quantized_Limits::max_norm_rcp); } - amax = fmaxf(amax, thread_amax); - const float scale_amax = rocm_subwarp_allreduce(thread_amax, rocm_op::max{}); - const e8m0_t biased_exponent = - ptx::float_to_e8m0(scale_amax * Quantized_Limits::max_norm_rcp); if (lane % LANES_PER_SCALE == 0 && y < valid_rows && scale_offset_X < valid_scales) { scales[y * scale_stride + scale_offset_X] = biased_exponent; } @@ -116,15 +122,11 @@ __device__ __forceinline__ float quantize_band_rowwise(const band_t &in, return amax; } -// Colwise MXFP8 of a band in LDS: each lane owns COLS_PER_LANE adjacent columns, with one pass for -// the amaxes and one to convert. Writes the scales of the first valid_cols columns and returns the -// lane's amax. -template -__device__ __forceinline__ float quantize_band_colwise(const band_t &in, band_t &out, - e8m0_t *const scales, - const size_t valid_cols) { +// Raises col_amax to the amaxes of the lane's COLS_PER_LANE adjacent columns of a band in LDS. +template +__device__ __forceinline__ void band_column_amaxes(const band_t &in, + float (&col_amax)[COLS_PER_LANE]) { const int lane_offset_X = (threadIdx.x % WAVE_SIZE) * COLS_PER_LANE; - float col_amax[COLS_PER_LANE] = {}; #pragma unroll for (int y = 0; y < BAND_Y; y++) { Vec x; @@ -134,13 +136,30 @@ __device__ __forceinline__ float quantize_band_colwise(const band_t &in, col_amax[i] = fmaxf(col_amax[i], fabsf(static_cast(x.data.elt[i]))); } } +} + +// Colwise MXFP8 of a band in LDS: each lane owns COLS_PER_LANE adjacent columns, with one pass for +// the amaxes and one to convert, or with IS_2D only the latter, with the scales of the tiles. +// Writes the scales of the first valid_cols columns and returns the lane's amax, which IS_2D leaves +// at 0. +template +__device__ __forceinline__ float quantize_band_colwise(const band_t &in, band_t &out, + e8m0_t *const scales, + const size_t valid_cols, + const uint32_t tile_exponents = 0) { + const int lane_offset_X = (threadIdx.x % WAVE_SIZE) * COLS_PER_LANE; + float col_amax[COLS_PER_LANE] = {}; + if constexpr (!IS_2D) { + band_column_amaxes(in, col_amax); + } float amax = 0; float block_scale_inverse[COLS_PER_LANE]; #pragma unroll for (int i = 0; i < COLS_PER_LANE; i++) { amax = fmaxf(amax, col_amax[i]); const e8m0_t biased_exponent = - ptx::float_to_e8m0(col_amax[i] * Quantized_Limits::max_norm_rcp); + IS_2D ? tile_exponent(tile_exponents, lane_offset_X) + : ptx::float_to_e8m0(col_amax[i] * Quantized_Limits::max_norm_rcp); if (static_cast(lane_offset_X + i) < valid_cols) { scales[lane_offset_X + i] = biased_exponent; } @@ -162,10 +181,10 @@ __device__ __forceinline__ float quantize_band_colwise(const band_t &in, #endif // __gfx1250__ // Quantizes chunk (chunk_Y, chunk_X) of a rows x cols tensor and writes its dbias partials to row -// dbias_row of dbias_workspace. +// dbias_row of dbias_workspace. IS_2D scales TILE_2D x TILE_2D tiles instead of rows and columns. template + size_t SCALE_DIM_X, bool IS_2D> __device__ __forceinline__ void quantize_chunk( const IType *input_ptr, const IType *act_input_ptr, OType *output_rowwise, OType *output_colwise, e8m0_t *scales_rowwise, e8m0_t *scales_colwise, @@ -182,6 +201,7 @@ __device__ __forceinline__ void quantize_chunk( constexpr size_t AUX_BYTES = std::max({OUT_BYTES, ACT_BYTES, size_t{1}}); static_assert(!USE_COLWISE_SCALING || SCALE_DIM_Y == BAND_Y, "A band must be one colwise scale block"); + static_assert(!IS_2D || TILE_2D == BAND_Y, "A band must be one row of tiles"); __shared__ alignas(128) band_t in_sh[WAVES]; // Stages the activation input of DACT, then holds the outputs once in_sh has the activations. @@ -244,22 +264,45 @@ __device__ __forceinline__ void quantize_chunk( } } + // The amax of a tile scales all its elements; in the colwise layout, its lanes are adjacent. + uint32_t tile_exponents = 0; + if constexpr (IS_2D) { + constexpr int TILE_LANES = TILE_2D / COLS_PER_LANE; // 8 + float col_amax[COLS_PER_LANE] = {}; + band_column_amaxes(in, col_amax); + float lane_amax = 0; +#pragma unroll + for (int i = 0; i < COLS_PER_LANE; i++) { + lane_amax = fmaxf(lane_amax, col_amax[i]); + } + amax = fmaxf(amax, lane_amax); + const int biased_exponent = + ptx::float_to_e8m0(rocm_subwarp_allreduce(lane_amax, rocm_op::max{}) * + Quantized_Limits::max_norm_rcp); +#pragma unroll + for (int t = 0; t < CHUNK_DIM_X / TILE_2D; t++) { + tile_exponents |= + static_cast(__builtin_amdgcn_readlane(biased_exponent, t * TILE_LANES)) + << (8 * t); + } + } + if constexpr (USE_ROWWISE_SCALING) { const size_t scale_offset_X = chunk_offset_X / SCALE_DIM_X; // Like the 128-wide chunks of the generic kernel, this also fills the scale padding. - amax = fmaxf(amax, quantize_band_rowwise( + amax = fmaxf(amax, quantize_band_rowwise( in, out_rowwise, scales_rowwise + band_offset_Y * scale_stride_rowwise + scale_offset_X, scale_stride_rowwise, rows - band_offset_Y, - scale_stride_rowwise - scale_offset_X)); + scale_stride_rowwise - scale_offset_X, tile_exponents)); tdm::store_2d(output_rowwise, out_rowwise, rows, cols, cols, band_offset_Y, chunk_offset_X); } if constexpr (USE_COLWISE_SCALING) { - amax = fmaxf(amax, quantize_band_colwise( + amax = fmaxf(amax, quantize_band_colwise( in, out_colwise, scales_colwise + band_offset_Y / SCALE_DIM_Y * scale_stride_colwise + chunk_offset_X, - cols - chunk_offset_X)); + cols - chunk_offset_X, tile_exponents)); tdm::store_2d(output_colwise, out_colwise, rows, cols, cols, band_offset_Y, chunk_offset_X); } } @@ -298,7 +341,7 @@ __device__ __forceinline__ void quantize_chunk( template + size_t SCALE_DIM_X, bool IS_2D> __global__ void __launch_bounds__(THREADS_PER_CHUNK) quantize_mxfp8_tdm_kernel(const IType *input, const IType *act_input, OType *output_rowwise, OType *output_colwise, e8m0_t *scales_rowwise, @@ -311,69 +354,83 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) return; } } - quantize_chunk( - input, act_input, output_rowwise, output_colwise, scales_rowwise, scales_colwise, - dbias_workspace, amax, rows, cols, scale_stride_rowwise, scale_stride_colwise, blockIdx.y, - blockIdx.x, blockIdx.y); + quantize_chunk(input, act_input, output_rowwise, output_colwise, scales_rowwise, + scales_colwise, dbias_workspace, amax, rows, cols, scale_stride_rowwise, + scale_stride_colwise, blockIdx.y, blockIdx.x, blockIdx.y); } -// The tensors share the last dimension and each has a multiple of CHUNK_DIM_Y rows, so no chunk -// straddles two of them. The dbias partials of chunk row blockIdx.y go to that workspace row. +// With a common last dimension, the tensors each have a multiple of CHUNK_DIM_Y rows, so no chunk +// of the 2D grid straddles two of them, and the dbias partials of chunk row blockIdx.y go to that +// workspace row. Otherwise, the 1D grid has a block per chunk of find_group_chunk and no dbias. template + size_t SCALE_DIM_X, bool IS_2D> __global__ void __launch_bounds__(THREADS_PER_CHUNK) grouped_quantize_mxfp8_tdm_kernel(const IType *input, const IType *act_input, OType *output_rowwise, OType *output_colwise, e8m0_t *scales_rowwise, e8m0_t *scales_colwise, const float *noop, float *dbias_workspace, float *amax, - const size_t first_logical_dim, const size_t cols, + const size_t first_logical_dim, const size_t last_logical_dim, const size_t num_tensors, const int64_t *offsets_ptr, - const int64_t *first_dims_ptr) { + const int64_t *first_dims_ptr, + const int64_t *last_dims_ptr) { if constexpr (!IS_DBIAS && !IS_DACT && !IS_ACT) { if (noop != nullptr && noop[0] == 1.0f) { return; } } - const size_t chunk_offset_Y = static_cast(blockIdx.y) * CHUNK_DIM_Y; - size_t tensor_id; size_t tensor_offset; - if (offsets_ptr == nullptr) { - const size_t rows_per_tensor = first_logical_dim / num_tensors; - tensor_id = chunk_offset_Y / rows_per_tensor; - tensor_offset = tensor_id * rows_per_tensor * cols; - } else { - // The last tensor starting at or before the chunk - size_t lo = 1; - size_t hi = num_tensors; - while (lo < hi) { - const size_t mid = lo + (hi - lo) / 2; - if (static_cast(offsets_ptr[mid]) <= chunk_offset_Y * cols) { - lo = mid + 1; - } else { - hi = mid; - } + size_t rows; + size_t cols; + size_t chunk_Y; + size_t chunk_X; + size_t scales_offset_rowwise; + size_t scales_offset_colwise; + if (last_dims_ptr == nullptr) { + cols = last_logical_dim; + const size_t chunk_offset_Y = static_cast(blockIdx.y) * CHUNK_DIM_Y; + size_t tensor_id; + if (offsets_ptr == nullptr) { + const size_t rows_per_tensor = first_logical_dim / num_tensors; + tensor_id = chunk_offset_Y / rows_per_tensor; + tensor_offset = tensor_id * rows_per_tensor * cols; + } else { + tensor_id = find_tensor(offsets_ptr, num_tensors, chunk_offset_Y * cols); + tensor_offset = static_cast(offsets_ptr[tensor_id]); } - tensor_id = lo - 1; - tensor_offset = static_cast(offsets_ptr[tensor_id]); + rows = first_dims_ptr != nullptr ? static_cast(first_dims_ptr[tensor_id]) + : first_logical_dim / num_tensors; + const size_t tensor_offset_Y = tensor_offset / cols; + chunk_Y = blockIdx.y - tensor_offset_Y / CHUNK_DIM_Y; + chunk_X = blockIdx.x; + scales_offset_rowwise = tensor_offset_Y * DIVUP(cols, SCALE_DIM_X); + scales_offset_colwise = tensor_offset_Y / SCALE_DIM_Y * cols; + } else { + const auto chunk = find_group_chunk( + blockIdx.x, num_tensors, first_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr); + tensor_offset = chunk.offset; + rows = chunk.rows; + cols = chunk.cols; + chunk_Y = chunk.chunk_Y; + chunk_X = chunk.chunk_X; + // Both dimensions are divisible by the scale blocks. + scales_offset_rowwise = tensor_offset / SCALE_DIM_X; + scales_offset_colwise = tensor_offset / SCALE_DIM_Y; } - const size_t rows = first_dims_ptr != nullptr ? static_cast(first_dims_ptr[tensor_id]) - : first_logical_dim / num_tensors; - const size_t tensor_offset_Y = tensor_offset / cols; - const size_t chunk_Y = blockIdx.y - tensor_offset_Y / CHUNK_DIM_Y; if (chunk_Y * CHUNK_DIM_Y >= rows) { return; } constexpr bool ROWWISE = SCALE_DIM_X > 1; constexpr bool COLWISE = SCALE_DIM_Y > 1; - const size_t scale_stride_rowwise = DIVUP(cols, SCALE_DIM_X); - quantize_chunk( - input + tensor_offset, IS_DACT ? act_input + tensor_offset : nullptr, - ROWWISE ? output_rowwise + tensor_offset : nullptr, - COLWISE ? output_colwise + tensor_offset : nullptr, - ROWWISE ? scales_rowwise + tensor_offset_Y * scale_stride_rowwise : nullptr, - COLWISE ? scales_colwise + tensor_offset_Y / SCALE_DIM_Y * cols : nullptr, dbias_workspace, - amax, rows, cols, scale_stride_rowwise, cols, chunk_Y, blockIdx.x, blockIdx.y); + quantize_chunk(input + tensor_offset, IS_DACT ? act_input + tensor_offset : nullptr, + ROWWISE ? output_rowwise + tensor_offset : nullptr, + COLWISE ? output_colwise + tensor_offset : nullptr, + ROWWISE ? scales_rowwise + scales_offset_rowwise : nullptr, + COLWISE ? scales_colwise + scales_offset_colwise : nullptr, + dbias_workspace, amax, rows, cols, DIVUP(cols, SCALE_DIM_X), cols, chunk_Y, + chunk_X, blockIdx.y); } struct MultiQuantizeArgs { @@ -409,7 +466,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) if (chunk_Y * CHUNK_DIM_Y >= rows || blockIdx.x * CHUNK_DIM_X >= cols) { return; } - quantize_chunk( + quantize_chunk( static_cast(args.input[t]), nullptr, static_cast(args.output_rowwise[t]), static_cast(args.output_colwise[t]), static_cast(args.scales_rowwise[t]), static_cast(args.scales_colwise[t]), @@ -421,7 +479,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) template bool quantize_tdm(const Tensor &input, const Tensor *act_input, const Tensor *noop, Tensor *output, - Tensor *dbias, Tensor *workspace, hipStream_t stream) { + Tensor *dbias, Tensor *workspace, const bool use_2d_quantization, + hipStream_t stream) { using namespace tdm_kernel; const bool rowwise = output->has_data(); const bool colwise = output->has_columnwise_data(); @@ -448,19 +507,22 @@ bool quantize_tdm(const Tensor &input, const Tensor *act_input, const Tensor *no (colwise ? 32 : 1), SCALE_DIM_Y, TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( (rowwise ? 32 : 1), SCALE_DIM_X, - quantize_mxfp8_tdm_kernel - <<>>( - reinterpret_cast(input.data.dptr), - IS_DACT ? reinterpret_cast(act_input->data.dptr) : nullptr, - reinterpret_cast(output->data.dptr), - reinterpret_cast(output->columnwise_data.dptr), - reinterpret_cast(output->scale_inv.dptr), - reinterpret_cast(output->columnwise_scale_inv.dptr), - reinterpret_cast(noop->data.dptr), workspace_ptr, - reinterpret_cast(output->amax.dptr), rows, cols, - rowwise ? output->scale_inv.shape[1] : 1, - colwise ? output->columnwise_scale_inv.shape[1] : 1);)))); // NOLINT(*) + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_2d_quantization, IS_2D, + quantize_mxfp8_tdm_kernel + <<>>( + reinterpret_cast(input.data.dptr), + IS_DACT ? reinterpret_cast(act_input->data.dptr) + : nullptr, + reinterpret_cast(output->data.dptr), + reinterpret_cast(output->columnwise_data.dptr), + reinterpret_cast(output->scale_inv.dptr), + reinterpret_cast(output->columnwise_scale_inv.dptr), + reinterpret_cast(noop->data.dptr), workspace_ptr, + reinterpret_cast(output->amax.dptr), rows, cols, + rowwise ? output->scale_inv.shape[1] : 1, + colwise ? output->columnwise_scale_inv.shape[1] : 1);))))); // NOLINT(*) NVTE_CHECK_CUDA(hipGetLastError()); if constexpr (IS_DBIAS) { TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( @@ -476,28 +538,35 @@ template bool group_quantize_tdm(const GroupedTensor *input, const GroupedTensor *activations, const Tensor *noop, GroupedTensor *output, GroupedTensor *dbias, - Tensor *workspace, hipStream_t stream) { + Tensor *workspace, const bool use_2d_quantization, hipStream_t stream) { using namespace tdm_kernel; const bool rowwise = output->has_data(); const bool colwise = output->has_columnwise_data(); const size_t first_logical_dim = input->logical_shape.data[0]; - const size_t cols = input->logical_shape.data[1]; + const size_t last_logical_dim = input->logical_shape.data[1]; const size_t num_tensors = input->num_tensors; const bool same_shape = output->all_same_shape(); - // Per-tensor rows of a varying first dim are device-side and, as upstream requires, 128-aligned. + const bool same_last_dim = output->all_same_last_dim(); + const int64_t *const offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr); + // Varying dimensions are device-side and, as upstream requires, divisible by 128. Upstream also + // has no dbias for varying last dimensions. const size_t tensor_rows = same_shape ? first_logical_dim / num_tensors : first_logical_dim; - if (!output->all_same_last_dim() || tensor_rows % CHUNK_DIM_Y != 0 || - !use_tdm(first_logical_dim, cols, + if ((!output->varying_both_dims() && tensor_rows % CHUNK_DIM_Y != 0) || + (!same_last_dim && (IS_DBIAS || offsets_ptr == nullptr)) || + !use_tdm(first_logical_dim, last_logical_dim, {input->data.dptr, IS_DACT ? activations->data.dptr : nullptr, rowwise ? output->data.dptr : nullptr, colwise ? output->columnwise_data.dptr : nullptr, IS_DBIAS ? workspace->data.dptr : nullptr})) { return false; } - const int64_t *const offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr); const int64_t *const first_dims_ptr = reinterpret_cast(output->first_dims.dptr); + const int64_t *const last_dims_ptr = reinterpret_cast(output->last_dims.dptr); float *const workspace_ptr = IS_DBIAS ? reinterpret_cast(workspace->data.dptr) : nullptr; - const dim3 grid(DIVUP(cols, CHUNK_DIM_X), DIVUP(first_logical_dim, CHUNK_DIM_Y)); + const dim3 grid = same_last_dim ? dim3(DIVUP(last_logical_dim, CHUNK_DIM_X), + DIVUP(first_logical_dim, CHUNK_DIM_Y)) + : dim3(DIVUP(first_logical_dim * last_logical_dim, + CHUNK_DIM_Y * CHUNK_DIM_X)); TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input->dtype(), IType, TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( @@ -506,24 +575,30 @@ bool group_quantize_tdm(const GroupedTensor *input, const GroupedTensor *activat (colwise ? 32 : 1), SCALE_DIM_Y, TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( (rowwise ? 32 : 1), SCALE_DIM_X, - grouped_quantize_mxfp8_tdm_kernel - <<>>( - reinterpret_cast(input->data.dptr), - IS_DACT ? reinterpret_cast(activations->data.dptr) : nullptr, - reinterpret_cast(output->data.dptr), - reinterpret_cast(output->columnwise_data.dptr), - reinterpret_cast(output->scale_inv.dptr), - reinterpret_cast(output->columnwise_scale_inv.dptr), - reinterpret_cast(noop->data.dptr), workspace_ptr, - reinterpret_cast(output->amax.dptr), first_logical_dim, cols, - num_tensors, offsets_ptr, first_dims_ptr);)))); // NOLINT(*) + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_2d_quantization, IS_2D, + grouped_quantize_mxfp8_tdm_kernel + <<>>( + reinterpret_cast(input->data.dptr), + IS_DACT ? reinterpret_cast(activations->data.dptr) + : nullptr, + reinterpret_cast(output->data.dptr), + reinterpret_cast(output->columnwise_data.dptr), + reinterpret_cast(output->scale_inv.dptr), + reinterpret_cast(output->columnwise_scale_inv.dptr), + reinterpret_cast(noop->data.dptr), workspace_ptr, + reinterpret_cast(output->amax.dptr), first_logical_dim, + last_logical_dim, num_tensors, offsets_ptr, first_dims_ptr, + last_dims_ptr);))))); // NOLINT(*) NVTE_CHECK_CUDA(hipGetLastError()); if constexpr (IS_DBIAS) { TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input->dtype(), IType, - reduce_dbias_tdm(reinterpret_cast(dbias->data.dptr), workspace_ptr, cols, - num_tensors, grid.y, same_shape ? nullptr : offsets_ptr, + reduce_dbias_tdm(reinterpret_cast(dbias->data.dptr), workspace_ptr, + last_logical_dim, num_tensors, grid.y, + same_shape ? nullptr : offsets_ptr, same_shape ? nullptr : first_dims_ptr, stream);); } return true; diff --git a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu index de5571df5..816ac8cda 100644 --- a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu +++ b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu @@ -21,6 +21,8 @@ #include "common/utils.cuh" #ifdef __HIP_PLATFORM_AMD__ #include "common/util/rocm_device_utils.cuh" +#include "common/util/system.h" +#include "common/cast/fp8/rocm_quantize_fp8_tdm.cuh" #endif #if (!defined(__CUDA_MINIMUM_ARCH__) && __CUDA_ARCH__ >= 900) || \ @@ -570,6 +572,17 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor const auto out_dtype = return_identity ? output.dtype : output_t.dtype; +#ifdef __HIP_PLATFORM_AMD__ + static_assert(dispatch::fp8::tdm_kernel::CHUNK_DIM_X == BLOCK_TILE_DIM && + dispatch::fp8::tdm_kernel::CHUNK_DIM_Y == BLOCK_TILE_DIM); + if (dispatch::fp8::quantize_blockwise_tdm<2>( + input, output, output_t, scale_inv, scale_inv_t, num_rows, row_length, return_identity, + return_transpose, scale_stride_x, scale_stride_y, scale_t_stride_x, scale_t_stride_y, + epsilon, pow_2_scale, noop_ptr, stream)) { + return; + } +#endif + const size_t num_blocks_x = DIVUP(row_length, BLOCK_TILE_DIM); const size_t num_blocks_y = DIVUP(num_rows, BLOCK_TILE_DIM); diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu index e224757d4..ab776480f 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu @@ -21,6 +21,10 @@ #include "common/transpose/cast_transpose.h" #include "common/util/cuda_runtime.h" #include "common/utils.cuh" +#ifdef __HIP_PLATFORM_AMD__ +#include "common/util/system.h" +#include "common/cast/fp8/rocm_quantize_fp8_tdm.cuh" +#endif namespace transformer_engine { namespace { @@ -954,6 +958,20 @@ void quantize_transpose_vector_blockwise(const SimpleTensor& input, SimpleTensor const float* noop_ptr = reinterpret_cast(noop_tensor.dptr); +#ifdef __HIP_PLATFORM_AMD__ + static_assert(dispatch::fp8::tdm_kernel::CHUNK_DIM_X == static_cast(kTileDim) && + dispatch::fp8::tdm_kernel::CHUNK_DIM_Y == static_cast(kTileDim)); + if (columnwise_option != FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT && + dispatch::fp8::quantize_blockwise_tdm<1>( + input, output, output_t, scale_inv, scale_inv_t, num_rows, row_length, + rowwise_option != FP8BlockwiseRowwiseOption::NONE, + columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY, + scale_stride_x, scale_stride_y, scale_t_stride_x, scale_t_stride_y, epsilon, pow2_scale, + noop_ptr, stream)) { + return; + } +#endif + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( input.dtype, InputType, From 8c85cb918f2367be2404e1cf29c0ec7a259dd250 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Thu, 29 Oct 2026 15:58:07 +0000 Subject: [PATCH 7/8] add mxfp4 --- .../common/cast/mxfp4/quantize_mxfp4.cuh | 14 + .../cast/mxfp4/rocm_quantize_mxfp4_tdm.cuh | 330 ++++++++++++++++++ 2 files changed, 344 insertions(+) create mode 100644 transformer_engine/common/cast/mxfp4/rocm_quantize_mxfp4_tdm.cuh diff --git a/transformer_engine/common/cast/mxfp4/quantize_mxfp4.cuh b/transformer_engine/common/cast/mxfp4/quantize_mxfp4.cuh index 21878ce97..c79b12a22 100644 --- a/transformer_engine/common/cast/mxfp4/quantize_mxfp4.cuh +++ b/transformer_engine/common/cast/mxfp4/quantize_mxfp4.cuh @@ -12,7 +12,10 @@ #define TRANSFORMER_ENGINE_QUANTIZE_MXFP4_CUH_ #include "../../common.h" +#include "../../util/cuda_runtime.h" +#include "../../util/system.h" #include "cast_transpose_mxfp4_shuffled.cuh" +#include "rocm_quantize_mxfp4_tdm.cuh" namespace transformer_engine { namespace dispatch { @@ -87,6 +90,17 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, colwise_scale_N_pad = static_cast(output->columnwise_scale_inv.shape[1]); } + const tdm_kernel::ScaleLayout rowwise_scales{ + reinterpret_cast(output->scale_inv.dptr), static_cast(rowwise_scale_stride), + rowwise_scale_N, rowwise_scale_M_pad, rowwise_scale_N_pad, scale_shuffle}; + const tdm_kernel::ScaleLayout colwise_scales{ + reinterpret_cast(output->columnwise_scale_inv.dptr), + static_cast(colwise_scale_stride), colwise_scale_N, colwise_scale_M_pad, + colwise_scale_N_pad, scale_shuffle}; + if (quantize_tdm(input, output, use_hadamard, rowwise_scales, colwise_scales, stream)) { + return; + } + nvte_cast_transpose_mxfp4_fused_shuffle( input.data.dptr, use_rowwise ? output->data.dptr : nullptr, diff --git a/transformer_engine/common/cast/mxfp4/rocm_quantize_mxfp4_tdm.cuh b/transformer_engine/common/cast/mxfp4/rocm_quantize_mxfp4_tdm.cuh new file mode 100644 index 000000000..06eef4816 --- /dev/null +++ b/transformer_engine/common/cast/mxfp4/rocm_quantize_mxfp4_tdm.cuh @@ -0,0 +1,330 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ +#pragma once +//#include "hip/hip_runtime.h" //dummy include to prevent hipification adding this header + +// gfx1250 MXFP4 quantize with the Tensor Data Mover: like the MXFP8 TDM path, each wave moves its +// own 32-row band of a 128x128 chunk global->LDS->global. It computes the scales and the Hadamard +// transform of cast_transpose_mxfp4_shuffled, whose header the includer must include first, and +// rounds to nearest even with the hardware conversion, as that kernel does on gfx950. + +#include "../../util/rocm_device_utils.cuh" +#include "../../util/rocm_tdm.cuh" +#include "../core/rocm_tdm_common.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp4 { +namespace tdm_kernel { + +using common::use_tdm; + +constexpr size_t CHUNK_DIM_Y = 128; +constexpr size_t CHUNK_DIM_X = 128; +constexpr size_t THREADS_PER_CHUNK = 128; +constexpr size_t WAVE_SIZE = 32; +constexpr size_t WAVES = THREADS_PER_CHUNK / WAVE_SIZE; // 4 +constexpr size_t BAND_Y = CHUNK_DIM_Y / WAVES; // 32 +constexpr size_t COLS_PER_LANE = CHUNK_DIM_X / WAVE_SIZE; // 4 +constexpr size_t BLOCK_SIZE = 32; +constexpr size_t BAND_BYTES = BAND_Y * CHUNK_DIM_X / 2; // packed FP4 of a band +constexpr size_t CHUNK_BYTES = CHUNK_DIM_Y * CHUNK_DIM_X / 2; +// The shuffled FP4 layouts store each 16-row group contiguously, in 16-byte column slices. +constexpr size_t SHUFFLE_ROWS = 16; +constexpr size_t SHUFFLE_SLICE = 16; +static_assert(BAND_Y == BLOCK_SIZE, "A band must be one colwise scale block"); + +// Where the E8M0 scales of one direction go: scale (row, col) for col < valid_cols, at +// row * stride + col or, with shuffle, at its compute_scale_shuffle_index within the padding. +struct ScaleLayout { + uint8_t *ptr; + size_t stride; + int valid_cols; + int pad_rows; + int pad_cols; + bool shuffle; +}; + +#ifdef __gfx1250__ +typedef float f32x8 __attribute__((ext_vector_type(8))); + +__device__ __forceinline__ void write_scale(const ScaleLayout &layout, const int row, const int col, + const uint8_t e8m0) { + if (col >= layout.valid_cols) { + return; + } + if (!layout.shuffle) { + layout.ptr[row * layout.stride + col] = e8m0; + } else if (row < layout.pad_rows && col < layout.pad_cols) { + layout.ptr[te_mxfp4::compute_scale_shuffle_index(row, col, layout.pad_cols)] = e8m0; + } +} + +// hadamard16_inplace of the four quad lanes whose values x[4q..4q+3] one lane holds. +__device__ __forceinline__ void hadamard16(float *x) { +#pragma unroll + for (int q = 0; q < 4; q++) { + const float a0 = x[4 * q] + x[4 * q + 1]; + const float a1 = x[4 * q] - x[4 * q + 1]; + const float a2 = x[4 * q + 2] + x[4 * q + 3]; + const float a3 = x[4 * q + 2] - x[4 * q + 3]; + x[4 * q] = a0 + a2; + x[4 * q + 2] = a0 - a2; + x[4 * q + 1] = a1 + a3; + x[4 * q + 3] = a1 - a3; + } +#pragma unroll + for (int d = 1; d <= 2; d *= 2) { + float y[16]; +#pragma unroll + for (int q = 0; q < 4; q++) { +#pragma unroll + for (int i = 0; i < 4; i++) { + const float p = x[4 * (q ^ d) + i]; + y[4 * q + i] = (q & d) ? p - x[4 * q + i] : p + x[4 * q + i]; + } + } +#pragma unroll + for (int i = 0; i < 16; i++) { + x[i] = y[i]; + } + } +#pragma unroll + for (int i = 0; i < 16; i++) { + x[i] *= 0.25f; + } +} + +// Scales a 32-element block for FP4 (x holds n of its values; a lane group of 32 / n shares it) +// and returns its E8M0 scale. As fmaxf skips NaN, only an all-NaN block has a NaN amax. +template +__device__ __forceinline__ uint8_t block_scale(const float *x, const int n, float &scale) { + float amax = fabsf(x[0]); + for (int i = 1; i < n; i++) { + amax = fmaxf(amax, fabsf(x[i])); + } + amax = rocm_subwarp_allreduce(amax, rocm_op::max{}); + return te_mxfp4::compute_e8m0_scale(amax, scale); +} + +// x / scale for 8 values as FP4, the first in the low nibble. +__device__ __forceinline__ uint32_t cvt_8x_to_fp4(const float *x, const float scale) { + const f32x8 v = {x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7]}; + return __builtin_amdgcn_cvt_scalef32_pk8_fp4_f32(v, scale); +} + +__device__ __forceinline__ float bf16_bits_to_float(const uint32_t bits) { + return __uint_as_float(bits << 16); +} +#endif // __gfx1250__ + +// Quantizes chunk (blockIdx.y, blockIdx.x) of the rows x cols BF16 input to the rowwise FP4 data, +// rows x cols / 2 bytes, and the columnwise FP4 data of its transpose, cols x rows / 2 bytes. +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + quantize_mxfp4_tdm_kernel(const uint16_t *input, uint8_t *rowwise_fp4, uint8_t *colwise_fp4, + const ScaleLayout rowwise_scales, const ScaleLayout colwise_scales, + const uint32_t rows, const uint32_t cols, const bool shuffle_rowwise, + const bool shuffle_colwise) { +#ifdef __gfx1250__ + __shared__ alignas(128) uint16_t in_sh[WAVES][BAND_Y][CHUNK_DIM_X]; + __shared__ alignas(128) uint8_t rowwise_sh[ROWWISE ? WAVES : 1][BAND_BYTES]; + // The transpose of the chunk, so that its rows go out in CHUNK_DIM_Y / 2-byte segments. + __shared__ alignas(128) uint8_t colwise_sh[COLWISE ? CHUNK_BYTES : 1]; + __shared__ alignas(4) uint8_t colwise_scale_sh[COLWISE ? CHUNK_DIM_X : 1][WAVES]; + + const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / WAVE_SIZE); + const int lane = threadIdx.x % WAVE_SIZE; + const uint32_t band_offset_Y = blockIdx.y * CHUNK_DIM_Y + wave * BAND_Y; + const uint32_t chunk_offset_X = blockIdx.x * CHUNK_DIM_X; + // The columnwise stores need every wave at the barrier. + const bool active = band_offset_Y < rows; + if (!COLWISE && !active) { + return; + } + auto &in = in_sh[wave]; + if (active) { + tdm::load_2d(in, input, rows, cols, cols, band_offset_Y, chunk_offset_X); + tdm::wait<0>(); + } + + if (ROWWISE && active) { + // Per step, 8 lanes cover a row with 16 elements, one Hadamard half, each. + constexpr int ELEMS = 16; + constexpr int LANES_PER_ROW = CHUNK_DIM_X / ELEMS; // 8 + constexpr int ROWS_PER_STEP = WAVE_SIZE / LANES_PER_ROW; // 4 + const int lane_offset_Y = lane / LANES_PER_ROW; + const int lane_offset_X = (lane % LANES_PER_ROW) * ELEMS; + const int packed_X = lane_offset_X / 2; + uint8_t *const out = rowwise_sh[wave]; +#pragma unroll + for (int step = 0; step < BAND_Y / ROWS_PER_STEP; step++) { + const int y = step * ROWS_PER_STEP + lane_offset_Y; + uint4 bits[2]; + bits[0] = *reinterpret_cast(&in[y][lane_offset_X]); + bits[1] = *reinterpret_cast(&in[y][lane_offset_X + 8]); + const uint32_t *const words = reinterpret_cast(bits); + float x[ELEMS]; +#pragma unroll + for (int i = 0; i < ELEMS / 2; i++) { + x[2 * i] = bf16_bits_to_float(words[i] & 0xFFFF); + x[2 * i + 1] = bf16_bits_to_float(words[i] >> 16); + } + if constexpr (HADAMARD) { + hadamard16(x); + } + float scale; + const uint8_t e8m0 = block_scale(x, ELEMS, scale); + const uint2 packed = make_uint2(cvt_8x_to_fp4(x, scale), cvt_8x_to_fp4(x + 8, scale)); + const int offset = shuffle_rowwise + ? y / SHUFFLE_ROWS * (SHUFFLE_ROWS * CHUNK_DIM_X / 2) + + packed_X / SHUFFLE_SLICE * (SHUFFLE_ROWS * SHUFFLE_SLICE) + + y % SHUFFLE_ROWS * SHUFFLE_SLICE + packed_X % SHUFFLE_SLICE + : y * (CHUNK_DIM_X / 2) + packed_X; + *reinterpret_cast(out + offset) = packed; + if (lane % 2 == 0 && band_offset_Y + y < rows) { + write_scale(rowwise_scales, band_offset_Y + y, (chunk_offset_X + lane_offset_X) / BLOCK_SIZE, + e8m0); + } + } + const uint32_t packed_cols = cols / 2; + if (shuffle_rowwise) { + tdm::store_2d(rowwise_fp4, + *reinterpret_cast(out), + rows / SHUFFLE_ROWS, packed_cols * SHUFFLE_ROWS, packed_cols * SHUFFLE_ROWS, + band_offset_Y / SHUFFLE_ROWS, chunk_offset_X / 2 * SHUFFLE_ROWS); + } else { + tdm::store_2d(rowwise_fp4, *reinterpret_cast(out), rows, + packed_cols, packed_cols, band_offset_Y, chunk_offset_X / 2); + } + } + + if constexpr (COLWISE) { + // A lane quantizes its COLS_PER_LANE columns one at a time; a band is one block of each. Its + // 16 packed bytes are the band's segment of a row of the transpose: (c, wave) goes to + // c * CHUNK_DIM_Y / 2 + wave * 16 or, shuffled, to its 16-column group's slice wave. + if (active) { + const int lane_offset_X = lane * COLS_PER_LANE; + uint2 bits[BAND_Y]; +#pragma unroll + for (int y = 0; y < BAND_Y; y++) { + bits[y] = *reinterpret_cast(&in[y][lane_offset_X]); + } +#pragma unroll + for (int i = 0; i < COLS_PER_LANE; i++) { + float x[BAND_Y]; +#pragma unroll + for (int y = 0; y < BAND_Y; y++) { + const uint32_t word = i < 2 ? bits[y].x : bits[y].y; + x[y] = bf16_bits_to_float(i % 2 == 0 ? word & 0xFFFF : word >> 16); + } + if constexpr (HADAMARD) { + hadamard16(x); + hadamard16(x + 16); + } + float scale; + const uint8_t e8m0 = block_scale<1>(x, BAND_Y, scale); + const uint4 packed = make_uint4(cvt_8x_to_fp4(x, scale), cvt_8x_to_fp4(x + 8, scale), + cvt_8x_to_fp4(x + 16, scale), cvt_8x_to_fp4(x + 24, scale)); + const int c = lane_offset_X + i; + const int offset = shuffle_colwise + ? c / SHUFFLE_ROWS * (SHUFFLE_ROWS * CHUNK_DIM_Y / 2) + + wave * (SHUFFLE_ROWS * SHUFFLE_SLICE) + + c % SHUFFLE_ROWS * SHUFFLE_SLICE + : c * (CHUNK_DIM_Y / 2) + wave * (BAND_Y / 2); + *reinterpret_cast(colwise_sh + offset) = packed; + colwise_scale_sh[c][wave] = e8m0; + } + } + __syncthreads(); + + // A thread per column writes its scales, at once where they are contiguous and aligned. + const uint32_t col = chunk_offset_X + threadIdx.x; + const int first_scale = blockIdx.y * WAVES; + if (col < cols) { + const ScaleLayout &layout = colwise_scales; + uint8_t *const dst = layout.ptr + col * layout.stride + first_scale; + if (!layout.shuffle && first_scale + static_cast(WAVES) <= layout.valid_cols && + reinterpret_cast(dst) % 4 == 0) { + *reinterpret_cast(dst) = + *reinterpret_cast(colwise_scale_sh[threadIdx.x]); + } else { +#pragma unroll + for (int w = 0; w < WAVES; w++) { + write_scale(layout, col, first_scale + w, colwise_scale_sh[threadIdx.x][w]); + } + } + } + if (wave == 0) { + const uint32_t chunk_offset_Y = blockIdx.y * CHUNK_DIM_Y; + const uint32_t packed_rows = rows / 2; + if (shuffle_colwise) { + tdm::store_2d(colwise_fp4, + *reinterpret_cast(colwise_sh), + cols / SHUFFLE_ROWS, packed_rows * SHUFFLE_ROWS, packed_rows * SHUFFLE_ROWS, + chunk_offset_X / SHUFFLE_ROWS, chunk_offset_Y / 2 * SHUFFLE_ROWS); + } else { + tdm::store_2d(colwise_fp4, + *reinterpret_cast(colwise_sh), + cols, packed_rows, packed_rows, chunk_offset_X, chunk_offset_Y / 2); + } + } + } + tdm::wait<0>(); +#endif // __gfx1250__ +} + +} // namespace tdm_kernel + +// MXFP4 quantize of input into output, with the scale layouts of quantize. Returns false, leaving +// the work to cast_transpose_mxfp4_shuffled, where TDM does not apply. A template, so that only the +// files calling it compile its kernels. +template +bool quantize_tdm(const Tensor &input, Tensor *output, const bool use_hadamard, + const tdm_kernel::ScaleLayout &rowwise_scales, + const tdm_kernel::ScaleLayout &colwise_scales, hipStream_t stream) { + using namespace tdm_kernel; + const size_t rows = input.flat_first_dim(); + const size_t cols = input.flat_last_dim(); + const bool rowwise = output->has_data(); + const bool colwise = output->has_columnwise_data(); + const bool shuffle_rowwise = rowwise && output->mxfp4_shuffle_rowwise_data; + const bool shuffle_colwise = colwise && output->mxfp4_shuffle_columnwise_data; + // Every block and, in the shuffled layouts, 16-row group of 32 packed bytes is complete. + if (!(rowwise || colwise) || cols % BLOCK_SIZE != 0 || (colwise && rows % BLOCK_SIZE != 0) || + (shuffle_rowwise && (rows % SHUFFLE_ROWS != 0 || cols % (4 * SHUFFLE_SLICE) != 0)) || + (shuffle_colwise && (cols % SHUFFLE_ROWS != 0 || rows % (4 * SHUFFLE_SLICE) != 0)) || + !use_tdm(rows, cols, + {input.data.dptr, rowwise ? output->data.dptr : nullptr, + colwise ? output->columnwise_data.dptr : nullptr})) { + return false; + } + const dim3 grid(DIVUP(cols, CHUNK_DIM_X), DIVUP(rows, CHUNK_DIM_Y)); + const auto launch = [&](auto kernel) { + kernel<<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(output->data.dptr), + reinterpret_cast(output->columnwise_data.dptr), rowwise_scales, colwise_scales, + rows, cols, shuffle_rowwise, shuffle_colwise); + }; + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_hadamard, HADAMARD, + if (rowwise && colwise) { + launch(quantize_mxfp4_tdm_kernel); + } else if (rowwise) { + launch(quantize_mxfp4_tdm_kernel); + } else { + launch(quantize_mxfp4_tdm_kernel); + }); // NOLINT(*) + NVTE_CHECK_CUDA(hipGetLastError()); + return true; +} + +} // namespace mxfp4 +} // namespace dispatch +} // namespace transformer_engine From 2bd9a4b964107112f1dbbca2474abad1c2e7e7a8 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Thu, 29 Oct 2026 18:13:37 +0000 Subject: [PATCH 8/8] add nvfp4 support --- .../common/cast/core/rocm_tdm_common.cuh | 46 ++- .../cast/mxfp4/rocm_quantize_mxfp4_tdm.cuh | 40 ++- ...quantize_transpose_vector_blockwise_fp4.cu | 17 + .../transpose/rocm_quantize_nvfp4_tdm.cuh | 296 ++++++++++++++++++ 4 files changed, 382 insertions(+), 17 deletions(-) create mode 100644 transformer_engine/common/transpose/rocm_quantize_nvfp4_tdm.cuh diff --git a/transformer_engine/common/cast/core/rocm_tdm_common.cuh b/transformer_engine/common/cast/core/rocm_tdm_common.cuh index 55a21c258..84da7b4f8 100644 --- a/transformer_engine/common/cast/core/rocm_tdm_common.cuh +++ b/transformer_engine/common/cast/core/rocm_tdm_common.cuh @@ -7,8 +7,9 @@ //#include "hip/hip_runtime.h" //dummy include to prevent hipification adding this header // gfx1250 Tensor Data Mover pieces shared by the TDM quantize paths: the check whether TDM applies, -// the chunks of grouped tensors, the reduction of per-chunk dbias partials, and amax. This expects -// its includer to have included the common headers, util/system.h and util/cuda_runtime.h. +// conversions and stores, the chunks of grouped tensors, the reduction of per-chunk dbias partials, +// and amax. This expects its includer to have included the common headers, util/system.h and +// util/cuda_runtime.h. #include #include @@ -39,8 +40,49 @@ __device__ __forceinline__ T cvt_from_float(const float x) { return static_cast(x); } } + +typedef float f32x8 __attribute__((ext_vector_type(8))); + +// x / scale for 8 values as FP4, rounded to nearest even, the first in the low nibble. +__device__ __forceinline__ uint32_t cvt_8x_to_fp4(const float *x, const float scale) { + const f32x8 v = {x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7]}; + return __builtin_amdgcn_cvt_scalef32_pk8_fp4_f32(v, scale); +} + +// cvt_8x_to_fp4, rounding stochastically with a random byte of rbits per value. The conversion +// rounds its first value up when bits 15:0 of its random operand plus the 16-bit fraction to the +// next FP4 value carry, its second likewise with bits 31:16, and the other 6 with shifts of the +// same bits, so each conversion takes only 2 values. +__device__ __forceinline__ uint32_t cvt_8x_to_fp4_sr(const float *x, const uint2 rbits, + const float scale) { + uint32_t packed = 0; +#pragma unroll + for (int i = 0; i < 4; i++) { + const uint32_t bytes = ((i < 2 ? rbits.x : rbits.y) >> (16 * (i % 2))) & 0xFFFF; + // The bytes on top of half a byte, so that the probability of rounding up rounds to nearest. + const uint32_t r = (bytes & 0xFF) << 8 | (bytes & 0xFF00) << 16 | 0x00800080; + const f32x8 v = {x[2 * i], x[2 * i + 1], 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + packed |= (__builtin_amdgcn_cvt_scalef32_sr_pk8_fp4_f32(v, r, scale) & 0xFF) << (8 * i); + } + return packed; +} #endif // __gfx1250__ +// Stores the first n of the N bytes of src to dst, at once when that is all of them and dst is +// aligned for it. +template +__device__ __forceinline__ void store_bytes(T *dst, const T (&src)[N], const int n) { + static_assert(sizeof(T) == 1 && (N == 4 || N == 8), "Stores 4 or 8 bytes"); + using Word = std::conditional_t; + if (n >= N && reinterpret_cast(dst) % N == 0) { + *reinterpret_cast(dst) = *reinterpret_cast(src); + } else { + for (int i = 0; i < min(n, N); i++) { + dst[i] = src[i]; + } + } +} + // The last of num_tensors tensors that starts at or before element offset. __device__ __forceinline__ size_t find_tensor(const int64_t *offsets, const size_t num_tensors, const size_t offset) { diff --git a/transformer_engine/common/cast/mxfp4/rocm_quantize_mxfp4_tdm.cuh b/transformer_engine/common/cast/mxfp4/rocm_quantize_mxfp4_tdm.cuh index 06eef4816..25749a028 100644 --- a/transformer_engine/common/cast/mxfp4/rocm_quantize_mxfp4_tdm.cuh +++ b/transformer_engine/common/cast/mxfp4/rocm_quantize_mxfp4_tdm.cuh @@ -20,6 +20,7 @@ namespace dispatch { namespace mxfp4 { namespace tdm_kernel { +using common::store_bytes; using common::use_tdm; constexpr size_t CHUNK_DIM_Y = 128; @@ -30,12 +31,14 @@ constexpr size_t WAVES = THREADS_PER_CHUNK / WAVE_SIZE; // 4 constexpr size_t BAND_Y = CHUNK_DIM_Y / WAVES; // 32 constexpr size_t COLS_PER_LANE = CHUNK_DIM_X / WAVE_SIZE; // 4 constexpr size_t BLOCK_SIZE = 32; +constexpr size_t SCALES_PER_ROW = CHUNK_DIM_X / BLOCK_SIZE; // 4 constexpr size_t BAND_BYTES = BAND_Y * CHUNK_DIM_X / 2; // packed FP4 of a band constexpr size_t CHUNK_BYTES = CHUNK_DIM_Y * CHUNK_DIM_X / 2; // The shuffled FP4 layouts store each 16-row group contiguously, in 16-byte column slices. constexpr size_t SHUFFLE_ROWS = 16; constexpr size_t SHUFFLE_SLICE = 16; static_assert(BAND_Y == BLOCK_SIZE, "A band must be one colwise scale block"); +static_assert(BAND_Y == WAVE_SIZE, "A lane per band row writes the rowwise scales"); // Where the E8M0 scales of one direction go: scale (row, col) for col < valid_cols, at // row * stride + col or, with shuffle, at its compute_scale_shuffle_index within the padding. @@ -49,7 +52,7 @@ struct ScaleLayout { }; #ifdef __gfx1250__ -typedef float f32x8 __attribute__((ext_vector_type(8))); +using common::cvt_8x_to_fp4; __device__ __forceinline__ void write_scale(const ScaleLayout &layout, const int row, const int col, const uint8_t e8m0) { @@ -110,12 +113,6 @@ __device__ __forceinline__ uint8_t block_scale(const float *x, const int n, floa return te_mxfp4::compute_e8m0_scale(amax, scale); } -// x / scale for 8 values as FP4, the first in the low nibble. -__device__ __forceinline__ uint32_t cvt_8x_to_fp4(const float *x, const float scale) { - const f32x8 v = {x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7]}; - return __builtin_amdgcn_cvt_scalef32_pk8_fp4_f32(v, scale); -} - __device__ __forceinline__ float bf16_bits_to_float(const uint32_t bits) { return __uint_as_float(bits << 16); } @@ -132,6 +129,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) #ifdef __gfx1250__ __shared__ alignas(128) uint16_t in_sh[WAVES][BAND_Y][CHUNK_DIM_X]; __shared__ alignas(128) uint8_t rowwise_sh[ROWWISE ? WAVES : 1][BAND_BYTES]; + __shared__ alignas(4) uint8_t rowwise_scale_sh[ROWWISE ? WAVES : 1][BAND_Y][SCALES_PER_ROW]; // The transpose of the chunk, so that its rows go out in CHUNK_DIM_Y / 2-byte segments. __shared__ alignas(128) uint8_t colwise_sh[COLWISE ? CHUNK_BYTES : 1]; __shared__ alignas(4) uint8_t colwise_scale_sh[COLWISE ? CHUNK_DIM_X : 1][WAVES]; @@ -160,6 +158,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const int lane_offset_X = (lane % LANES_PER_ROW) * ELEMS; const int packed_X = lane_offset_X / 2; uint8_t *const out = rowwise_sh[wave]; + auto &scales = rowwise_scale_sh[wave]; #pragma unroll for (int step = 0; step < BAND_Y / ROWS_PER_STEP; step++) { const int y = step * ROWS_PER_STEP + lane_offset_Y; @@ -185,9 +184,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) y % SHUFFLE_ROWS * SHUFFLE_SLICE + packed_X % SHUFFLE_SLICE : y * (CHUNK_DIM_X / 2) + packed_X; *reinterpret_cast(out + offset) = packed; - if (lane % 2 == 0 && band_offset_Y + y < rows) { - write_scale(rowwise_scales, band_offset_Y + y, (chunk_offset_X + lane_offset_X) / BLOCK_SIZE, - e8m0); + if (lane % 2 == 0) { + scales[y][lane_offset_X / BLOCK_SIZE] = e8m0; } } const uint32_t packed_cols = cols / 2; @@ -201,6 +199,20 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) tdm::store_2d(rowwise_fp4, *reinterpret_cast(out), rows, packed_cols, packed_cols, band_offset_Y, chunk_offset_X / 2); } + // A lane per band row writes the row's scales. + const uint32_t row = band_offset_Y + lane; + const int first_scale = chunk_offset_X / BLOCK_SIZE; + if (row < rows) { + if (!rowwise_scales.shuffle) { + store_bytes(rowwise_scales.ptr + row * rowwise_scales.stride + first_scale, scales[lane], + rowwise_scales.valid_cols - first_scale); + } else { +#pragma unroll + for (int s = 0; s < SCALES_PER_ROW; s++) { + write_scale(rowwise_scales, row, first_scale + s, scales[lane][s]); + } + } + } } if constexpr (COLWISE) { @@ -247,11 +259,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const int first_scale = blockIdx.y * WAVES; if (col < cols) { const ScaleLayout &layout = colwise_scales; - uint8_t *const dst = layout.ptr + col * layout.stride + first_scale; - if (!layout.shuffle && first_scale + static_cast(WAVES) <= layout.valid_cols && - reinterpret_cast(dst) % 4 == 0) { - *reinterpret_cast(dst) = - *reinterpret_cast(colwise_scale_sh[threadIdx.x]); + if (!layout.shuffle) { + store_bytes(layout.ptr + col * layout.stride + first_scale, colwise_scale_sh[threadIdx.x], + layout.valid_cols - first_scale); } else { #pragma unroll for (int w = 0; w < WAVES; w++) { diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu index e1e043321..3e0776ca7 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -24,6 +24,9 @@ #include "common/utils.cuh" #ifdef __HIP_PLATFORM_AMD__ #include "common/util/cuda_runtime.h" +#include "common/util/system.h" +// After the util headers it uses. +#include "common/cast/core/rocm_tdm_common.cuh" #endif namespace transformer_engine { @@ -1130,6 +1133,10 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo } } // NOLINT(readability/fn_size) 4over6 grows upstream's kernel past 500 lines +#ifdef __HIP_PLATFORM_AMD__ +#include "common/transpose/rocm_quantize_nvfp4_tdm.cuh" +#endif + } // namespace } // namespace quantize_transpose_nvfp4 #endif // CUDA_VERSION >= 12080 @@ -1222,6 +1229,16 @@ void quantize_transpose_vector_blockwise_fp4( rng_state = reinterpret_cast(rng_state_te_tensor.data.dptr); } +#ifdef __HIP_PLATFORM_AMD__ + if (nvfp4_tdm::quantize(input, global_amax, scale_inv, scale_inv_t, output, output_t, num_rows, + row_length, return_identity, return_transpose, swizzled_scale, + use_stochastic_rounding, rng_state, use_2d_quantization, + row_scaled_nvfp4, nvfp4_4over6_mode != kNVTENVFP44Over6Disabled, + noop_ptr, e4m3_max_for_arch, stream)) { + return; + } +#endif + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( input.dtype, InputType, diff --git a/transformer_engine/common/transpose/rocm_quantize_nvfp4_tdm.cuh b/transformer_engine/common/transpose/rocm_quantize_nvfp4_tdm.cuh new file mode 100644 index 000000000..d6de1dacb --- /dev/null +++ b/transformer_engine/common/transpose/rocm_quantize_nvfp4_tdm.cuh @@ -0,0 +1,296 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ +#pragma once +//#include "hip/hip_runtime.h" //dummy include to prevent hipification adding this header + +// gfx1250 NVFP4 quantize with the Tensor Data Mover, included in quantize_transpose_nvfp4 after the +// scale helpers it shares with block_scaled_1d_cast_transpose_kernel, whose numerics it reproduces. +// As in the MXFP8 and MXFP4 TDM paths, each wave moves its own 32-row band of a 128x128 chunk +// global->LDS->global, and the transpose is assembled in LDS for the whole chunk. + +namespace nvfp4_tdm { + +constexpr int kChunkDim = 128; +constexpr int kThreads = 128; +constexpr int kWaveSize = 32; +constexpr int kWaves = kThreads / kWaveSize; // 4 +constexpr int kBandY = kChunkDim / kWaves; // 32 +constexpr int kColsPerLane = kChunkDim / kWaveSize; // 4 +constexpr int kBlockDim = 16; // values per scale +constexpr int kBlocks = kChunkDim / kBlockDim; // 8 per chunk row or column +constexpr int kBandBlocks = kBandY / kBlockDim; // 2 +static_assert(kBandY == kWaveSize, "A lane per band row writes the rowwise scales"); +static_assert(kThreads == kChunkDim, "A thread per chunk column writes the columnwise scales"); + +using Philox = + transformer_engine::curanddx::detail::philox4x32_native_state; + +#ifdef __gfx1250__ +// src[0..N) as floats, read up to 16 bytes at a time. +template +__device__ __forceinline__ void load_floats(const IType *src, float *x) { + constexpr int kVec = N < 16 / sizeof(IType) ? N : 16 / sizeof(IType); +#pragma unroll + for (int i = 0; i < N; i += kVec) { + Vec v; + v.load_from(src + i); +#pragma unroll + for (int j = 0; j < kVec; j++) { + x[i + j] = static_cast(v.data.elt[j]); + } + } +} + +// The 16 values x with amax as NVFP4, the first in the low nibble of packed; returns the E4M3 +// scale. Stochastic rounding takes a byte of one draw of rng per value, as the generic kernel does. +template +__device__ __forceinline__ uint8_t quantize_block(const float *x, const float amax, + const float encode_multiplier, + const float decode_scale, Philox &rng, + uint2 &packed) { + const fp8e4m3 scale_inv = ComputeDecodeScaleFP4(amax, encode_multiplier); + const float encode_scale = ComputeEncodeScaleFP4(scale_inv, decode_scale); + float y[kBlockDim]; +#pragma unroll + for (int i = 0; i < kBlockDim; i++) { + y[i] = x[i] * encode_scale; + } + if constexpr (kSR) { + const uint4 rbits = rng.generate4(); + packed.x = dispatch::common::cvt_8x_to_fp4_sr(y, make_uint2(rbits.x, rbits.y), 1.0f); + packed.y = dispatch::common::cvt_8x_to_fp4_sr(y + 8, make_uint2(rbits.z, rbits.w), 1.0f); + } else { + packed.x = dispatch::common::cvt_8x_to_fp4(y, 1.0f); + packed.y = dispatch::common::cvt_8x_to_fp4(y + 8, 1.0f); + } + return __builtin_bit_cast(uint8_t, scale_inv); +} +#endif // __gfx1250__ + +// Quantizes chunk (blockIdx.y, blockIdx.x) of the rows x cols input to the rowwise NVFP4 data, +// rows x cols / 2 bytes, and the columnwise data of its transpose, cols x rows / 2 bytes. +template +__global__ void __launch_bounds__(kThreads) + quantize_kernel(const IType *input, const float *global_amax, uint8_t *output_c, + uint8_t *output_t, uint8_t *scales_c, uint8_t *scales_t, const uint32_t rows, + const uint32_t cols, const size_t scale_stride, const size_t scale_t_stride, + const bool row_scaled, const size_t *rng_state, const float *noop, + const float e4m3_max) { +#ifdef __gfx1250__ + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + __shared__ alignas(128) IType in_sh[kWaves][kBandY][kChunkDim]; + __shared__ alignas(128) uint8_t rowwise_sh[kRowwise ? kWaves : 1][kBandY][kChunkDim / 2]; + __shared__ alignas(8) uint8_t rowwise_scale_sh[kRowwise ? kWaves : 1][kBandY][kBlocks]; + // The transpose of the chunk, so that its rows go out in kChunkDim / 2-byte segments. + __shared__ alignas(128) uint8_t colwise_sh[kColwise ? kChunkDim : 1][kChunkDim / 2]; + __shared__ alignas(8) uint8_t colwise_scale_sh[kColwise ? kChunkDim : 1][kBlocks]; + __shared__ float block_amax_sh[k2D ? kWaves : 1][kBandBlocks][kBlocks]; + + const int wave = __builtin_amdgcn_readfirstlane(threadIdx.x / kWaveSize); + const int lane = threadIdx.x % kWaveSize; + const uint32_t chunk_offset_Y = blockIdx.y * kChunkDim; + const uint32_t chunk_offset_X = blockIdx.x * kChunkDim; + const uint32_t band_offset_Y = chunk_offset_Y + wave * kBandY; + // The columnwise stores need every wave at the barrier. + const bool active = band_offset_Y < rows; + if (!kColwise && !active) { + return; + } + auto &in = in_sh[wave]; + if (active) { + tdm::load_2d(in, input, rows, cols, cols, band_offset_Y, chunk_offset_X); + } + + constexpr float fp4_max_inv = 1.0f / TypeExtrema::max; + const float global_encode_scale = + global_amax == nullptr ? 1.0f : ComputeGlobalEncodeScaleFP4(global_amax[0], e4m3_max); + const float global_encode_multiplier = global_encode_scale * fp4_max_inv; + const float global_decode_scale = 1.0 / global_encode_scale; + Philox rng; + if constexpr (kSR) { + const size_t sequence = + threadIdx.x + (blockIdx.x + static_cast(blockIdx.y) * gridDim.x) * kThreads; + rng.init(rng_state != nullptr ? rng_state[0] : 0, sequence, + rng_state != nullptr ? rng_state[1] : 0); + } + if (active) { + tdm::wait<0>(); + } + + if constexpr (k2D) { + // The amaxes of the band's 16x16 blocks, each from 4 lanes. + constexpr int kLanesPerBlock = kBlockDim / kColsPerLane; + if (active) { +#pragma unroll + for (int b = 0; b < kBandBlocks; b++) { + float amax = 0.0f; +#pragma unroll + for (int y = 0; y < kBlockDim; y++) { + float x[kColsPerLane]; + load_floats(&in[b * kBlockDim + y][lane * kColsPerLane], x); +#pragma unroll + for (int i = 0; i < kColsPerLane; i++) { + amax = fmaxf(amax, fabsf(x[i])); + } + } + amax = rocm_subwarp_allreduce(amax, rocm_op::max{}); + if (lane % kLanesPerBlock == 0) { + block_amax_sh[wave][b][lane / kLanesPerBlock] = amax; + } + } + __builtin_amdgcn_wave_barrier(); + } + } + + if (kRowwise && active) { + // Per step, a lane quantizes a block of a row, kBlocks lanes the whole row. + constexpr int kRowsPerStep = kWaveSize / kBlocks; // 4 + const int block = lane % kBlocks; + auto &out = rowwise_sh[wave]; + auto &scales = rowwise_scale_sh[wave]; +#pragma unroll + for (int step = 0; step < kBandY / kRowsPerStep; step++) { + const int y = step * kRowsPerStep + lane / kBlocks; + const uint32_t row = band_offset_Y + y; + float x[kBlockDim]; + load_floats(&in[y][block * kBlockDim], x); + float amax = 0.0f; + if constexpr (k2D) { + amax = block_amax_sh[wave][y / kBlockDim][block]; + } else { +#pragma unroll + for (int i = 0; i < kBlockDim; i++) { + amax = fmaxf(amax, fabsf(x[i])); + } + } + float encode_multiplier = global_encode_multiplier; + float decode_scale = global_decode_scale; + if (row_scaled) { + const float row_encode_scale = + row < rows ? ComputeGlobalEncodeScaleFP4(global_amax[row], e4m3_max) : 1.0f; + encode_multiplier = row_encode_scale * fp4_max_inv; + decode_scale = 1.0f / row_encode_scale; + } + uint2 packed; + scales[y][block] = quantize_block(x, amax, encode_multiplier, decode_scale, rng, packed); + *reinterpret_cast(&out[y][block * kBlockDim / 2]) = packed; + } + tdm::store_2d(output_c, out, rows, cols / 2, cols / 2, band_offset_Y, chunk_offset_X / 2); + // A lane per band row writes the row's scales. + const uint32_t row = band_offset_Y + lane; + const uint32_t first_scale = chunk_offset_X / kBlockDim; + if (row < rows) { + dispatch::common::store_bytes(scales_c + row * scale_stride + first_scale, scales[lane], + DIVUP(cols, uint32_t{kBlockDim}) - first_scale); + } + } + + if constexpr (kColwise) { + // A lane quantizes the blocks of its kColsPerLane columns; block b of column c is segment + // wave * kBandBlocks + b of row c of the transpose. + if (active) { + const int lane_offset_X = lane * kColsPerLane; +#pragma unroll + for (int b = 0; b < kBandBlocks; b++) { + float x[kColsPerLane][kBlockDim]; +#pragma unroll + for (int y = 0; y < kBlockDim; y++) { + float v[kColsPerLane]; + load_floats(&in[b * kBlockDim + y][lane_offset_X], v); +#pragma unroll + for (int i = 0; i < kColsPerLane; i++) { + x[i][y] = v[i]; + } + } +#pragma unroll + for (int i = 0; i < kColsPerLane; i++) { + const int c = lane_offset_X + i; + const int segment = wave * kBandBlocks + b; + float amax = 0.0f; + if constexpr (k2D) { + amax = block_amax_sh[wave][b][c / kBlockDim]; + } else { +#pragma unroll + for (int y = 0; y < kBlockDim; y++) { + amax = fmaxf(amax, fabsf(x[i][y])); + } + } + uint2 packed; + colwise_scale_sh[c][segment] = quantize_block( + x[i], amax, global_encode_multiplier, global_decode_scale, rng, packed); + *reinterpret_cast(&colwise_sh[c][segment * kBlockDim / 2]) = packed; + } + } + } + __syncthreads(); + if (wave == 0) { + tdm::store_2d(output_t, colwise_sh, cols, rows / 2, rows / 2, chunk_offset_X, + chunk_offset_Y / 2); + } + // A thread per column writes its scales. + const uint32_t col = chunk_offset_X + threadIdx.x; + const uint32_t first_scale = chunk_offset_Y / kBlockDim; + if (col < cols) { + dispatch::common::store_bytes(scales_t + col * scale_t_stride + first_scale, + colwise_scale_sh[threadIdx.x], + DIVUP(rows, uint32_t{kBlockDim}) - first_scale); + } + } + tdm::wait<0>(); +#endif // __gfx1250__ +} + +// NVFP4 quantize as quantize_transpose_vector_blockwise_fp4 does it. Returns false, leaving the +// work to block_scaled_1d_cast_transpose_kernel, where TDM does not apply; that includes swizzled +// scales, which only GPUs without TDM request. Stochastic rounding draws its own random bits, so +// it matches that kernel in distribution rather than bitwise. +inline bool quantize(const SimpleTensor &input, const SimpleTensor &global_amax, + const SimpleTensor &scale_inv, const SimpleTensor &scale_inv_t, + const SimpleTensor &output, const SimpleTensor &output_t, const size_t rows, + const size_t cols, const bool rowwise, const bool colwise, + const bool swizzled, const bool use_sr, const size_t *rng_state, + const bool use_2d, const bool row_scaled, const bool use_4over6, + const float *noop, const float e4m3_max, cudaStream_t stream) { + // Packed rows must be whole 16-byte TDM units. + if (use_4over6 || swizzled || is_fp8_dtype(input.dtype) || (rowwise && cols % 32 != 0) || + (colwise && rows % 32 != 0) || + !dispatch::common::use_tdm(rows, cols, {input.dptr, rowwise ? output.dptr : nullptr, + colwise ? output_t.dptr : nullptr})) { + return false; + } + const dim3 grid(DIVUP(cols, static_cast(kChunkDim)), + DIVUP(rows, static_cast(kChunkDim))); + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input.dtype, IType, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_2d, k2D, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_sr, kSR, + const auto launch = [&](auto kernel) { + kernel<<>>( + reinterpret_cast(input.dptr), + reinterpret_cast(global_amax.dptr), + reinterpret_cast(output.dptr), + reinterpret_cast(output_t.dptr), + reinterpret_cast(scale_inv.dptr), + reinterpret_cast(scale_inv_t.dptr), rows, cols, + rowwise ? scale_inv.shape[1] : 0, colwise ? scale_inv_t.shape[1] : 0, + row_scaled, rng_state, noop, e4m3_max); + }; + if (rowwise && colwise) { + launch(quantize_kernel); + } else if (rowwise) { + launch(quantize_kernel); + } else { + launch(quantize_kernel); + }););); // NOLINT(*) + NVTE_CHECK_CUDA(cudaGetLastError()); + return true; +} + +} // namespace nvfp4_tdm