From 953f13d988f83ffb4945549fa73f9db7aeebdc6f Mon Sep 17 00:00:00 2001 From: Xiaokang Shang Date: Mon, 15 Jun 2026 14:40:12 +0800 Subject: [PATCH 1/3] blockwise fp8 fused multi-tensor kernel --- .../include/transformer_engine/transpose.h | 14 + .../common/transpose/cast_transpose.h | 6 + .../quantize_transpose_vector_blockwise.cu | 386 ++++++++++++++++-- .../pytorch/csrc/extensions/cast.cpp | 56 +++ 4 files changed, 438 insertions(+), 24 deletions(-) diff --git a/transformer_engine/common/include/transformer_engine/transpose.h b/transformer_engine/common/include/transformer_engine/transpose.h index cc069ee3ec..161ae49165 100644 --- a/transformer_engine/common/include/transformer_engine/transpose.h +++ b/transformer_engine/common/include/transformer_engine/transpose.h @@ -93,6 +93,20 @@ void nvte_fp8_transpose_dbias(const NVTETensor input, NVTETensor transposed_outp void nvte_multi_cast_transpose(size_t num_tensors, const NVTETensor* input_list, NVTETensor* output_list, cudaStream_t stream); +/*! \brief Cast multiple input tensors to 1D block-scaled FP8 tensors. + * + * \param[in] num_tensors Number of tensors. + * \param[in] input_list List of input tensors. + * \param[in,out] output_list List of 1D block-scaled FP8 output tensors. + * \param[in] quant_config Quantization configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_multi_quantize_transpose_vector_blockwise(size_t num_tensors, + const NVTETensor* input_list, + NVTETensor* output_list, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream); + /*! \brief Compute backward of GeLU operation on the input, then cast and transpose. * Additionally, reduce the result of the GeLU backward along the first dimension. * diff --git a/transformer_engine/common/transpose/cast_transpose.h b/transformer_engine/common/transpose/cast_transpose.h index 89266f4bbc..5d9ab9dfeb 100644 --- a/transformer_engine/common/transpose/cast_transpose.h +++ b/transformer_engine/common/transpose/cast_transpose.h @@ -63,6 +63,12 @@ void quantize_transpose_vector_blockwise(const SimpleTensor &input, SimpleTensor const bool pow_2_scale, const SimpleTensor &noop_tensor, cudaStream_t stream); +void multi_quantize_transpose_vector_blockwise( + const std::vector &input_list, std::vector &output_list, + const float epsilon, FP8BlockwiseRowwiseOption rowwise_option, + FP8BlockwiseColumnwiseOption columnwise_option, const bool pow_2_scale, + const SimpleTensor &noop_tensor, cudaStream_t stream); + void quantize_transpose_vector_blockwise_fp4( const SimpleTensor &input, const SimpleTensor &global_amax, SimpleTensor &scale_inv, SimpleTensor &scale_inv_t, SimpleTensor &output, SimpleTensor &output_t, const float epsilon, diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu index d38bf79963..88d1c99c3d 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu @@ -12,12 +12,14 @@ #include #include #include +#include #include #include "common/common.h" #include "common/recipe/recipe_common.cuh" #include "common/transpose/cast_transpose.h" #include "common/utils.cuh" +#include "transformer_engine/transpose.h" namespace transformer_engine { namespace { @@ -164,20 +166,33 @@ constexpr int kNumThreadsStore = kTileDim / kNVecOut; static_assert(kNumThreadsLoad <= kThreadsPerWarp, "kNumThreadsLoad must be <= kThreadsPerWarp"); static_assert(kNumThreadsStore <= kThreadsPerWarp, "kNumThreadsStore must be <= kThreadsPerWarp"); constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp; +constexpr int kMaxTensorsPerBlockwiseKernel = 32; // Keep kernel args comfortably under 4 KB. + +struct MultiBlockwiseQuantizeArgs { + void* input_list[kMaxTensorsPerBlockwiseKernel]; + void* output_c_list[kMaxTensorsPerBlockwiseKernel]; + void* output_t_list[kMaxTensorsPerBlockwiseKernel]; + void* scale_inv_c_list[kMaxTensorsPerBlockwiseKernel]; + void* scale_inv_t_list[kMaxTensorsPerBlockwiseKernel]; + int row_length_list[kMaxTensorsPerBlockwiseKernel]; + int num_rows_list[kMaxTensorsPerBlockwiseKernel]; + int scale_stride_x_list[kMaxTensorsPerBlockwiseKernel]; + int scale_stride_y_list[kMaxTensorsPerBlockwiseKernel]; + int scale_t_stride_x_list[kMaxTensorsPerBlockwiseKernel]; + int scale_t_stride_y_list[kMaxTensorsPerBlockwiseKernel]; + int block_range[kMaxTensorsPerBlockwiseKernel + 1]; + int num_tensors; +}; template -__global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpose_kernel( +__device__ __forceinline__ void block_scaled_1d_cast_transpose_impl( const IType* const input, OType* const output_c, OType* const output_t, CType* const tile_scales_inv_c, CType* const tile_scales_inv_t, const size_t row_length, const size_t num_rows, 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, FP8BlockwiseRowwiseOption rowwise_option, FP8BlockwiseColumnwiseOption columnwise_option, - const bool pow_2_scaling, const float* noop_ptr) { - // skip execution if noop - if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) { - return; - } - + const bool pow_2_scaling, const size_t tile_idx_x, const size_t tile_idx_y, + Vec* const smem) { bool return_rowwise = rowwise_option != FP8BlockwiseRowwiseOption::NONE; bool return_columnwise_gemm_ready = columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY; @@ -191,9 +206,6 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo Vec smem_type; }; - extern __shared__ char smem_base[]; - SMemVec* smem = reinterpret_cast(&smem_base[0]); - // Step 1: Load input to shared memory { constexpr int r_stride = kThreadsPerBlock / kNumThreadsLoad; // stride in rows of shared memory @@ -202,8 +214,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo (threadIdx.x % kNumThreadsLoad) * (kNVecIn / kNVecSMem); // Column in shared memory int r_s = threadIdx.x / kNumThreadsLoad; // Row in shared memory const size_t c_g = - static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem; // Column in global memory - size_t r_g = static_cast(blockIdx.y) * kTileDim + r_s; // Row in global memory + tile_idx_x * kTileDim + c_s * kNVecSMem; // Column in global memory + size_t r_g = tile_idx_y * kTileDim + r_s; // Row in global memory const size_t stride_g = static_cast(r_stride) * row_length; // Stride in global memory const size_t num_ele = c_g < row_length ? min(static_cast(kNVecIn), row_length - c_g) : 0; // For not aligned case @@ -248,8 +260,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo (threadIdx.x % kNumThreadsStore) * (kNVecOut / kNVecSMem); // Column in shared memory int r_s = threadIdx.x / kNumThreadsStore; // Row in shared memory const size_t c_g = - static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem; // Column in global memory - size_t r_g = static_cast(blockIdx.y) * kTileDim + r_s; // Row in global memory + tile_idx_x * kTileDim + c_s * kNVecSMem; // Column in global memory + size_t r_g = tile_idx_y * kTileDim + r_s; // Row in global memory const size_t stride_g = static_cast(r_stride) * row_length; // Stride in global memory const size_t num_ele = c_g < row_length ? min(static_cast(kNVecOut), row_length - c_g) : 0; // For not aligned case @@ -299,8 +311,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo } if (write_scale_inv) { CType scale_inv = 1.0 / scale; - size_t row_idx = static_cast(blockIdx.y) * kTileDim + r_s; - size_t col_idx = static_cast(blockIdx.x); + size_t row_idx = tile_idx_y * kTileDim + r_s; + size_t col_idx = tile_idx_x; tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv; } // Step 2.6: Quantize @@ -338,8 +350,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo const int r_s = (threadIdx.x % kNumThreadsStore) * kNVecOut; // Row in shared memory int c_s = threadIdx.x / kNumThreadsStore; // Column in shared memory size_t r_g = - static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem; // Row in global memory - const size_t c_g = static_cast(blockIdx.y) * kTileDim + r_s; // Column in global memory + tile_idx_x * kTileDim + c_s * kNVecSMem; // Row in global memory + const size_t c_g = tile_idx_y * kTileDim + r_s; // Column in global memory const size_t stride_g = static_cast(c_stride) * kNVecSMem * num_rows; // Stride in global memory const size_t num_ele = c_g < num_rows ? min(static_cast(kNVecOut), num_rows - c_g) @@ -388,8 +400,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo } if (write_scale_inv) { CType scale_inv = 1.0 / scale; - size_t row_idx = static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem + smem_idx; - size_t col_idx = static_cast(blockIdx.y); + size_t row_idx = tile_idx_x * kTileDim + c_s * kNVecSMem + smem_idx; + size_t col_idx = tile_idx_y; tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv; } // Step 3.6: Quantize @@ -435,9 +447,9 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo const int warp_idx = threadIdx.x / kThreadsPerWarp; const int r_s = thr_idx_in_warp * kThreadTileRow; // Row in shared memory int c_s = warp_idx * num_smem_reads; // Column in shared memory - size_t r_g = static_cast(blockIdx.y) * kTileDim + r_s; // Row in global memory + size_t r_g = tile_idx_y * kTileDim + r_s; // Row in global memory const size_t c_g = - static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem; // Column in global memory + tile_idx_x * kTileDim + c_s * kNVecSMem; // Column in global memory const size_t num_ele = c_g < row_length ? min(static_cast(kThreadTileCol), row_length - c_g) : 0; // For not aligned case @@ -485,8 +497,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo } if (write_scale_inv) { CType scale_inv = 1.0 / scale; - size_t row_idx = static_cast(blockIdx.y); - size_t col_idx = static_cast(blockIdx.x) * kTileDim + c_s * kNVecSMem + reg_idx; + size_t row_idx = tile_idx_y; + size_t col_idx = tile_idx_x * kTileDim + c_s * kNVecSMem + reg_idx; tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv; } } @@ -515,6 +527,93 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo } } +template +__global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpose_kernel( + const IType* const input, OType* const output_c, OType* const output_t, + CType* const tile_scales_inv_c, CType* const tile_scales_inv_t, const size_t row_length, + const size_t num_rows, 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, + FP8BlockwiseRowwiseOption rowwise_option, FP8BlockwiseColumnwiseOption columnwise_option, + const bool pow_2_scaling, const float* noop_ptr) { + // skip execution if noop + if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) { + return; + } + + using SMemVec = Vec; + extern __shared__ char smem_base[]; + SMemVec* smem = reinterpret_cast(&smem_base[0]); + + block_scaled_1d_cast_transpose_impl( + input, output_c, output_t, tile_scales_inv_c, tile_scales_inv_t, row_length, num_rows, + scale_stride_x, scale_stride_y, scale_t_stride_x, scale_t_stride_y, epsilon, rowwise_option, + columnwise_option, pow_2_scaling, static_cast(blockIdx.x), + static_cast(blockIdx.y), smem); +} + +template +__global__ void __launch_bounds__(kThreadsPerBlock) multi_block_scaled_1d_cast_transpose_kernel( + MultiBlockwiseQuantizeArgs args, const float epsilon, + FP8BlockwiseRowwiseOption rowwise_option, FP8BlockwiseColumnwiseOption columnwise_option, + const bool pow_2_scaling, const float* noop_ptr) { + // skip execution if noop + if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) { + return; + } + + int tensor_id = 0; + const int bid = blockIdx.x; + while (args.block_range[tensor_id + 1] <= bid) { + ++tensor_id; + } + + const size_t row_length = static_cast(args.row_length_list[tensor_id]); + const size_t num_rows = static_cast(args.num_rows_list[tensor_id]); + const size_t num_tiles_x = DIVUP(row_length, static_cast(kTileDim)); + const int tile_id = bid - args.block_range[tensor_id]; + const size_t tile_idx_x = static_cast(tile_id) % num_tiles_x; + const size_t tile_idx_y = static_cast(tile_id) / num_tiles_x; + + using SMemVec = Vec; + extern __shared__ char smem_base[]; + SMemVec* smem = reinterpret_cast(&smem_base[0]); + + block_scaled_1d_cast_transpose_impl( + reinterpret_cast(args.input_list[tensor_id]), + reinterpret_cast(args.output_c_list[tensor_id]), + reinterpret_cast(args.output_t_list[tensor_id]), + reinterpret_cast(args.scale_inv_c_list[tensor_id]), + reinterpret_cast(args.scale_inv_t_list[tensor_id]), row_length, num_rows, + static_cast(args.scale_stride_x_list[tensor_id]), + static_cast(args.scale_stride_y_list[tensor_id]), + static_cast(args.scale_t_stride_x_list[tensor_id]), + static_cast(args.scale_t_stride_y_list[tensor_id]), epsilon, rowwise_option, + columnwise_option, pow_2_scaling, tile_idx_x, tile_idx_y, smem); +} + +template +void launch_multi_block_scaled_1d_cast_transpose_kernel( + const MultiBlockwiseQuantizeArgs& kernel_args, const float epsilon, + FP8BlockwiseRowwiseOption rowwise_option, FP8BlockwiseColumnwiseOption columnwise_option, + const bool pow_2_scaling, const float* noop_ptr, cudaStream_t stream) { + if (kernel_args.num_tensors == 0) { + return; + } + + const int n_blocks = kernel_args.block_range[kernel_args.num_tensors]; + size_t smem_bytes = kSMemSize * sizeof(InputType); + if (smem_bytes >= 48 * 1024) { + cudaError_t err = cudaFuncSetAttribute( + &multi_block_scaled_1d_cast_transpose_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + NVTE_CHECK(err == cudaSuccess, "Failed to set dynamic shared memory size."); + } + multi_block_scaled_1d_cast_transpose_kernel + <<>>( + kernel_args, epsilon, rowwise_option, columnwise_option, pow_2_scaling, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + } // namespace } // namespace transformer_engine @@ -632,4 +731,243 @@ void quantize_transpose_vector_blockwise(const SimpleTensor& input, SimpleTensor NVTE_CHECK_CUDA(cudaGetLastError()); } +void multi_quantize_transpose_vector_blockwise( + const std::vector& input_list, std::vector& output_list, + const float epsilon, FP8BlockwiseRowwiseOption rowwise_option, + FP8BlockwiseColumnwiseOption columnwise_option, const bool pow2_scale, + const SimpleTensor& noop_tensor, cudaStream_t stream) { + NVTE_API_CALL(multi_quantize_transpose_vector_blockwise); + + NVTE_CHECK(input_list.size() == output_list.size(), + "Number of input and output tensors must match."); + if (input_list.empty()) { + return; + } + if (rowwise_option == FP8BlockwiseRowwiseOption::NONE && + columnwise_option == FP8BlockwiseColumnwiseOption::NONE) { + return; + } + + const DType input_dtype = input_list[0]->data.dtype; + DType output_dtype = DType::kNumTypes; + for (const auto* output : output_list) { + if (rowwise_option != FP8BlockwiseRowwiseOption::NONE && + (output->data.dptr != nullptr || !output->data.shape.empty())) { + output_dtype = output->data.dtype; + break; + } + if (columnwise_option != FP8BlockwiseColumnwiseOption::NONE && + (output->columnwise_data.dptr != nullptr || !output->columnwise_data.shape.empty())) { + output_dtype = output->columnwise_data.dtype; + break; + } + } + NVTE_CHECK(output_dtype != DType::kNumTypes, "Unable to infer output dtype."); + const float* noop_ptr = reinterpret_cast(noop_tensor.dptr); + + auto check_int_range = [](size_t value, const char* name) -> int { + NVTE_CHECK(value <= static_cast(std::numeric_limits::max()), name, + " exceeds int range: ", value); + return static_cast(value); + }; + + auto reset_kernel_args = [](MultiBlockwiseQuantizeArgs& args) { + args.num_tensors = 0; + args.block_range[0] = 0; + }; + + MultiBlockwiseQuantizeArgs kernel_args_aligned, kernel_args_unaligned; + reset_kernel_args(kernel_args_aligned); + reset_kernel_args(kernel_args_unaligned); + + auto launch_kernel_args = [&](MultiBlockwiseQuantizeArgs& args, bool aligned) { + if (args.num_tensors == 0) { + return; + } + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input_dtype, InputType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output_dtype, OutputType, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + aligned, kAligned, + launch_multi_block_scaled_1d_cast_transpose_kernel( + args, epsilon, rowwise_option, columnwise_option, pow2_scale, noop_ptr, + stream);) // kAligned + ) // OutputType + ) // InputType + reset_kernel_args(args); + }; + + for (size_t tensor_id = 0; tensor_id < input_list.size(); ++tensor_id) { + const auto& input = input_list[tensor_id]->data; + auto& output = output_list[tensor_id]->data; + auto& output_t = output_list[tensor_id]->columnwise_data; + auto& scale_inv = output_list[tensor_id]->scale_inv; + auto& scale_inv_t = output_list[tensor_id]->columnwise_scale_inv; + + NVTE_CHECK(output_list[tensor_id]->scaling_mode == NVTE_BLOCK_SCALING_1D, + "Output tensor ", tensor_id, " must use 1D block scaling."); + NVTE_CHECK(input.dtype == input_dtype, "Input tensor types do not match."); + + const size_t row_length = input.shape.size() > 0 ? input.shape.at(input.shape.size() - 1) : 1u; + size_t num_elements = row_length; + size_t num_rows = 1; + for (size_t i = 0; (i < input.shape.size() - 1) && (input.shape.size() > 0); ++i) { + num_rows *= input.shape.at(i); + num_elements *= input.shape.at(i); + } + + if (num_elements == 0) { + continue; + } + + size_t scale_stride_x = 0; + size_t scale_stride_y = 0; + size_t scale_t_stride_x = 0; + size_t scale_t_stride_y = 0; + + if (rowwise_option != FP8BlockwiseRowwiseOption::NONE) { + NVTE_CHECK(rowwise_option == FP8BlockwiseRowwiseOption::ROWWISE_GEMM_READY || + rowwise_option == FP8BlockwiseRowwiseOption::ROWWISE_COMPACT, + "Unexpected rowwise enum value."); + NVTE_CHECK(input.shape == output.shape, "Input and rowwise output shapes must match."); + NVTE_CHECK(output.dtype == output_dtype, "Rowwise output tensor types do not match."); + NVTE_CHECK(scale_inv.shape.size() == 2, "Scale dimension must be 2."); + bool rowwise_compact = rowwise_option == FP8BlockwiseRowwiseOption::ROWWISE_COMPACT; + size_t scale_k = scale_inv.shape[1]; + scale_stride_x = rowwise_compact ? 1 : scale_k; + scale_stride_y = rowwise_compact ? scale_k : 1; + } + + if (columnwise_option != FP8BlockwiseColumnwiseOption::NONE) { + NVTE_CHECK(output_t.shape.size() == input.shape.size(), + "Columnwise output must have same number of dimensions as input."); + if (output_t.shape.size() > 0) { + if (columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY) { + NVTE_CHECK(output_t.shape[0] == row_length, "Wrong dimension 0 of columnwise output."); + for (size_t i = 1; i < output_t.shape.size(); ++i) { + NVTE_CHECK(output_t.shape.at(i) == input.shape.at(i - 1), + "Wrong dimension in columnwise output."); + } + } else { + NVTE_CHECK(columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT, + "Unexpected columnwise enum value."); + NVTE_CHECK(output_t.shape[0] == input.shape[0], + "Wrong dimension 0 of compact columnwise output."); + NVTE_CHECK(input.shape == output_t.shape, + "Input and compact columnwise output shapes must match."); + } + } + NVTE_CHECK(output_t.dtype == output_dtype, "Columnwise output tensor types do not match."); + if (rowwise_option != FP8BlockwiseRowwiseOption::NONE) { + NVTE_CHECK(output.dtype == output_t.dtype, + "Rowwise and columnwise outputs need to have the same dtype."); + } + NVTE_CHECK(scale_inv_t.shape.size() == 2, "Scale_t dimension must be 2."); + bool columnwise_compact = + columnwise_option == FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT; + size_t scale_t_k = scale_inv_t.shape[1]; + scale_t_stride_x = columnwise_compact ? 1 : scale_t_k; + scale_t_stride_y = columnwise_compact ? scale_t_k : 1; + } + + const size_t num_tiles_x = DIVUP(row_length, static_cast(kTileDim)); + const size_t num_tiles_y = DIVUP(num_rows, static_cast(kTileDim)); + const size_t num_tiles = num_tiles_x * num_tiles_y; + if (num_tiles == 0) { + continue; + } + check_int_range(num_tiles, "Number of tiles"); + if (rowwise_option != FP8BlockwiseRowwiseOption::NONE) { + NVTE_CHECK(output.dptr != nullptr, "Rowwise output data pointer must not be null."); + NVTE_CHECK(scale_inv.dptr != nullptr, "Rowwise scale inverse pointer must not be null."); + } + if (columnwise_option != FP8BlockwiseColumnwiseOption::NONE) { + NVTE_CHECK(output_t.dptr != nullptr, "Columnwise output data pointer must not be null."); + NVTE_CHECK(scale_inv_t.dptr != nullptr, "Columnwise scale inverse pointer must not be null."); + } + + const bool aligned = row_length % kTileDim == 0 && num_rows % kTileDim == 0; + auto& kernel_args = aligned ? kernel_args_aligned : kernel_args_unaligned; + if (kernel_args.num_tensors == kMaxTensorsPerBlockwiseKernel) { + launch_kernel_args(kernel_args, aligned); + } + + const int pos = kernel_args.num_tensors; + kernel_args.input_list[pos] = input.dptr; + kernel_args.output_c_list[pos] = + rowwise_option != FP8BlockwiseRowwiseOption::NONE ? output.dptr : nullptr; + kernel_args.output_t_list[pos] = + columnwise_option != FP8BlockwiseColumnwiseOption::NONE ? output_t.dptr : nullptr; + kernel_args.scale_inv_c_list[pos] = + rowwise_option != FP8BlockwiseRowwiseOption::NONE ? scale_inv.dptr : nullptr; + kernel_args.scale_inv_t_list[pos] = + columnwise_option != FP8BlockwiseColumnwiseOption::NONE ? scale_inv_t.dptr : nullptr; + kernel_args.row_length_list[pos] = check_int_range(row_length, "Row length"); + kernel_args.num_rows_list[pos] = check_int_range(num_rows, "Number of rows"); + kernel_args.scale_stride_x_list[pos] = check_int_range(scale_stride_x, "Scale stride x"); + kernel_args.scale_stride_y_list[pos] = check_int_range(scale_stride_y, "Scale stride y"); + kernel_args.scale_t_stride_x_list[pos] = check_int_range(scale_t_stride_x, "Scale_t stride x"); + kernel_args.scale_t_stride_y_list[pos] = check_int_range(scale_t_stride_y, "Scale_t stride y"); + kernel_args.block_range[pos + 1] = + kernel_args.block_range[pos] + check_int_range(num_tiles, "Number of tiles"); + ++kernel_args.num_tensors; + } + + launch_kernel_args(kernel_args_aligned, true); + launch_kernel_args(kernel_args_unaligned, false); +} + } // namespace transformer_engine::detail + +void nvte_multi_quantize_transpose_vector_blockwise(size_t num_tensors, + const NVTETensor* input_list, + NVTETensor* output_list, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream) { + NVTE_API_CALL(nvte_multi_quantize_transpose_vector_blockwise); + using namespace transformer_engine; + + std::vector input_list_, output_list_; + input_list_.reserve(num_tensors); + output_list_.reserve(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + input_list_.push_back(convertNVTETensorCheck(input_list[i])); + output_list_.push_back(convertNVTETensorCheck(output_list[i])); + } + + QuantizationConfig quant_config_cpp; + if (quant_config != nullptr) { + quant_config_cpp = *reinterpret_cast(quant_config); + } + + Tensor dummy_tensor; + Tensor* noop_tensor = &dummy_tensor; + if (quant_config_cpp.noop_tensor != nullptr) { + noop_tensor = convertNVTETensorCheck(quant_config_cpp.noop_tensor); + } + + bool return_rowwise = false; + bool return_columnwise = false; + for (const auto* output : output_list_) { + return_rowwise |= output->has_data(); + return_columnwise |= output->has_columnwise_data(); + } + + const bool compact = quant_config_cpp.float8_block_scale_tensor_format == + Float8BlockScaleTensorFormat::COMPACT; + const auto rowwise_option = + return_rowwise ? (compact ? detail::FP8BlockwiseRowwiseOption::ROWWISE_COMPACT + : detail::FP8BlockwiseRowwiseOption::ROWWISE_GEMM_READY) + : detail::FP8BlockwiseRowwiseOption::NONE; + const auto columnwise_option = + return_columnwise + ? (compact ? detail::FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT + : detail::FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY) + : detail::FP8BlockwiseColumnwiseOption::NONE; + + detail::multi_quantize_transpose_vector_blockwise( + input_list_, output_list_, quant_config_cpp.amax_epsilon, rowwise_option, columnwise_option, + quant_config_cpp.force_pow_2_scales, noop_tensor->data, stream); +} diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 2c1edae4c6..dbb95f96d0 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -141,6 +141,62 @@ void multi_tensor_quantize_impl(const std::vector &input_list, nvte_multi_cast_transpose(nvte_tensor_input_list.size(), nvte_tensor_input_list.data(), nvte_tensor_output_list.data(), at::cuda::getCurrentCUDAStream()); }); + return; + } + + // Fused kernel for 1D FP8 blockwise multi-tensor quantize. + bool with_blockwise_fused_kernel = true; + Float8BlockQuantizer *first_blockwise_quantizer = nullptr; + for (size_t i = 0; i < num_tensors; i++) { + if (!detail::IsFloat8BlockwiseQuantizers(quantizer_py_list[i].ptr())) { + with_blockwise_fused_kernel = false; + break; + } + + auto *blockwise_quantizer = dynamic_cast(quantizer_cpp_list[i].get()); + if (blockwise_quantizer == nullptr || + blockwise_quantizer->get_scaling_mode() != NVTE_BLOCK_SCALING_1D) { + with_blockwise_fused_kernel = false; + break; + } + + if (i == 0) { + first_blockwise_quantizer = blockwise_quantizer; + continue; + } + + if (blockwise_quantizer->dtype != first_blockwise_quantizer->dtype || + blockwise_quantizer->force_pow_2_scales != + first_blockwise_quantizer->force_pow_2_scales || + blockwise_quantizer->amax_epsilon != first_blockwise_quantizer->amax_epsilon || + blockwise_quantizer->all_gather_usage != first_blockwise_quantizer->all_gather_usage || + blockwise_quantizer->rowwise_usage != first_blockwise_quantizer->rowwise_usage || + blockwise_quantizer->columnwise_usage != first_blockwise_quantizer->columnwise_usage) { + with_blockwise_fused_kernel = false; + break; + } + } + + if (with_blockwise_fused_kernel && first_blockwise_quantizer != nullptr) { + std::vector nvte_tensor_input_list; + std::vector nvte_tensor_output_list; + for (size_t i = 0; i < num_tensors; ++i) { + nvte_tensor_input_list.push_back(input_list[i].data()); + nvte_tensor_output_list.push_back(output_list[i].data()); + } + + QuantizationConfigWrapper quant_config; + quant_config.set_force_pow_2_scales(first_blockwise_quantizer->force_pow_2_scales); + quant_config.set_amax_epsilon(first_blockwise_quantizer->amax_epsilon); + if (first_blockwise_quantizer->all_gather_usage) { + quant_config.set_float8_block_scale_tensor_format(Float8BlockScaleTensorFormat::COMPACT); + } + + NVTE_SCOPED_GIL_RELEASE({ + nvte_multi_quantize_transpose_vector_blockwise( + nvte_tensor_input_list.size(), nvte_tensor_input_list.data(), + nvte_tensor_output_list.data(), quant_config, at::cuda::getCurrentCUDAStream()); + }); } else { // Quantize kernels individually for (size_t i = 0; i < num_tensors; ++i) { From fcd4b382753214f9a34c413c6a1eade1ab7d58f6 Mon Sep 17 00:00:00 2001 From: Xiaokang Shang Date: Mon, 15 Jun 2026 16:16:58 +0800 Subject: [PATCH 2/3] debug print --- .../transpose/quantize_transpose_vector_blockwise.cu | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu index 88d1c99c3d..cb61963c31 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu @@ -11,6 +11,7 @@ #include #include +#include #include #include #include @@ -167,6 +168,7 @@ static_assert(kNumThreadsLoad <= kThreadsPerWarp, "kNumThreadsLoad must be <= kT static_assert(kNumThreadsStore <= kThreadsPerWarp, "kNumThreadsStore must be <= kThreadsPerWarp"); constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp; constexpr int kMaxTensorsPerBlockwiseKernel = 32; // Keep kernel args comfortably under 4 KB. +bool g_printed_blockwise_fp8_multi_tensor_kernel = false; struct MultiBlockwiseQuantizeArgs { void* input_list[kMaxTensorsPerBlockwiseKernel]; @@ -608,6 +610,14 @@ void launch_multi_block_scaled_1d_cast_transpose_kernel( cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); NVTE_CHECK(err == cudaSuccess, "Failed to set dynamic shared memory size."); } + if (!g_printed_blockwise_fp8_multi_tensor_kernel) { + g_printed_blockwise_fp8_multi_tensor_kernel = true; + std::fprintf(stderr, + "[TransformerEngine] Using blockwise FP8 multi-tensor quantize kernel " + "(num_tensors=%d, num_blocks=%d, aligned=%d)\n", + kernel_args.num_tensors, n_blocks, kAligned ? 1 : 0); + std::fflush(stderr); + } multi_block_scaled_1d_cast_transpose_kernel <<>>( kernel_args, epsilon, rowwise_option, columnwise_option, pow_2_scaling, noop_ptr); From c437581f18a5c970646e33f7501ea04d6e6f9623 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 10:07:25 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../quantize_transpose_vector_blockwise.cu | 67 +++++++++---------- .../pytorch/csrc/extensions/cast.cpp | 3 +- 2 files changed, 32 insertions(+), 38 deletions(-) diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu index cb61963c31..b56a888d79 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu @@ -213,10 +213,9 @@ __device__ __forceinline__ void block_scaled_1d_cast_transpose_impl( constexpr int r_stride = kThreadsPerBlock / kNumThreadsLoad; // stride in rows of shared memory constexpr int num_iterations = kTileDim / r_stride; const int c_s = - (threadIdx.x % kNumThreadsLoad) * (kNVecIn / kNVecSMem); // Column in shared memory - int r_s = threadIdx.x / kNumThreadsLoad; // Row in shared memory - const size_t c_g = - tile_idx_x * kTileDim + c_s * kNVecSMem; // Column in global memory + (threadIdx.x % kNumThreadsLoad) * (kNVecIn / kNVecSMem); // Column in shared memory + int r_s = threadIdx.x / kNumThreadsLoad; // Row in shared memory + const size_t c_g = tile_idx_x * kTileDim + c_s * kNVecSMem; // Column in global memory size_t r_g = tile_idx_y * kTileDim + r_s; // Row in global memory const size_t stride_g = static_cast(r_stride) * row_length; // Stride in global memory const size_t num_ele = c_g < row_length ? min(static_cast(kNVecIn), row_length - c_g) @@ -259,10 +258,9 @@ __device__ __forceinline__ void block_scaled_1d_cast_transpose_impl( kThreadsPerBlock / kNumThreadsStore; // stride in rows of shared memory constexpr int num_iterations = kTileDim / r_stride; const int c_s = - (threadIdx.x % kNumThreadsStore) * (kNVecOut / kNVecSMem); // Column in shared memory - int r_s = threadIdx.x / kNumThreadsStore; // Row in shared memory - const size_t c_g = - tile_idx_x * kTileDim + c_s * kNVecSMem; // Column in global memory + (threadIdx.x % kNumThreadsStore) * (kNVecOut / kNVecSMem); // Column in shared memory + int r_s = threadIdx.x / kNumThreadsStore; // Row in shared memory + const size_t c_g = tile_idx_x * kTileDim + c_s * kNVecSMem; // Column in global memory size_t r_g = tile_idx_y * kTileDim + r_s; // Row in global memory const size_t stride_g = static_cast(r_stride) * row_length; // Stride in global memory const size_t num_ele = c_g < row_length ? min(static_cast(kNVecOut), row_length - c_g) @@ -351,9 +349,8 @@ __device__ __forceinline__ void block_scaled_1d_cast_transpose_impl( constexpr int num_iterations = kTileDim / (c_stride * kNVecSMem); const int r_s = (threadIdx.x % kNumThreadsStore) * kNVecOut; // Row in shared memory int c_s = threadIdx.x / kNumThreadsStore; // Column in shared memory - size_t r_g = - tile_idx_x * kTileDim + c_s * kNVecSMem; // Row in global memory - const size_t c_g = tile_idx_y * kTileDim + r_s; // Column in global memory + size_t r_g = tile_idx_x * kTileDim + c_s * kNVecSMem; // Row in global memory + const size_t c_g = tile_idx_y * kTileDim + r_s; // Column in global memory const size_t stride_g = static_cast(c_stride) * kNVecSMem * num_rows; // Stride in global memory const size_t num_ele = c_g < num_rows ? min(static_cast(kNVecOut), num_rows - c_g) @@ -447,11 +444,10 @@ __device__ __forceinline__ void block_scaled_1d_cast_transpose_impl( "num_iterations should be 1 for columnwise non-transpose case"); const int thr_idx_in_warp = threadIdx.x % kThreadsPerWarp; const int warp_idx = threadIdx.x / kThreadsPerWarp; - const int r_s = thr_idx_in_warp * kThreadTileRow; // Row in shared memory - int c_s = warp_idx * num_smem_reads; // Column in shared memory - size_t r_g = tile_idx_y * kTileDim + r_s; // Row in global memory - const size_t c_g = - tile_idx_x * kTileDim + c_s * kNVecSMem; // Column in global memory + const int r_s = thr_idx_in_warp * kThreadTileRow; // Row in shared memory + int c_s = warp_idx * num_smem_reads; // Column in shared memory + size_t r_g = tile_idx_y * kTileDim + r_s; // Row in global memory + const size_t c_g = tile_idx_x * kTileDim + c_s * kNVecSMem; // Column in global memory const size_t num_ele = c_g < row_length ? min(static_cast(kThreadTileCol), row_length - c_g) : 0; // For not aligned case @@ -554,10 +550,12 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo } template -__global__ void __launch_bounds__(kThreadsPerBlock) multi_block_scaled_1d_cast_transpose_kernel( - MultiBlockwiseQuantizeArgs args, const float epsilon, - FP8BlockwiseRowwiseOption rowwise_option, FP8BlockwiseColumnwiseOption columnwise_option, - const bool pow_2_scaling, const float* noop_ptr) { +__global__ void __launch_bounds__(kThreadsPerBlock) + multi_block_scaled_1d_cast_transpose_kernel(MultiBlockwiseQuantizeArgs args, + const float epsilon, + FP8BlockwiseRowwiseOption rowwise_option, + FP8BlockwiseColumnwiseOption columnwise_option, + const bool pow_2_scaling, const float* noop_ptr) { // skip execution if noop if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) { return; @@ -742,10 +740,9 @@ void quantize_transpose_vector_blockwise(const SimpleTensor& input, SimpleTensor } void multi_quantize_transpose_vector_blockwise( - const std::vector& input_list, std::vector& output_list, - const float epsilon, FP8BlockwiseRowwiseOption rowwise_option, - FP8BlockwiseColumnwiseOption columnwise_option, const bool pow2_scale, - const SimpleTensor& noop_tensor, cudaStream_t stream) { + const std::vector& input_list, std::vector& output_list, const float epsilon, + FP8BlockwiseRowwiseOption rowwise_option, FP8BlockwiseColumnwiseOption columnwise_option, + const bool pow2_scale, const SimpleTensor& noop_tensor, cudaStream_t stream) { NVTE_API_CALL(multi_quantize_transpose_vector_blockwise); NVTE_CHECK(input_list.size() == output_list.size(), @@ -800,12 +797,11 @@ void multi_quantize_transpose_vector_blockwise( output_dtype, OutputType, TRANSFORMER_ENGINE_SWITCH_CONDITION( aligned, kAligned, - launch_multi_block_scaled_1d_cast_transpose_kernel( + launch_multi_block_scaled_1d_cast_transpose_kernel( args, epsilon, rowwise_option, columnwise_option, pow2_scale, noop_ptr, stream);) // kAligned - ) // OutputType - ) // InputType + ) // OutputType + ) // InputType reset_kernel_args(args); }; @@ -816,8 +812,8 @@ void multi_quantize_transpose_vector_blockwise( auto& scale_inv = output_list[tensor_id]->scale_inv; auto& scale_inv_t = output_list[tensor_id]->columnwise_scale_inv; - NVTE_CHECK(output_list[tensor_id]->scaling_mode == NVTE_BLOCK_SCALING_1D, - "Output tensor ", tensor_id, " must use 1D block scaling."); + NVTE_CHECK(output_list[tensor_id]->scaling_mode == NVTE_BLOCK_SCALING_1D, "Output tensor ", + tensor_id, " must use 1D block scaling."); NVTE_CHECK(input.dtype == input_dtype, "Input tensor types do not match."); const size_t row_length = input.shape.size() > 0 ? input.shape.at(input.shape.size() - 1) : 1u; @@ -965,17 +961,16 @@ void nvte_multi_quantize_transpose_vector_blockwise(size_t num_tensors, return_columnwise |= output->has_columnwise_data(); } - const bool compact = quant_config_cpp.float8_block_scale_tensor_format == - Float8BlockScaleTensorFormat::COMPACT; + const bool compact = + quant_config_cpp.float8_block_scale_tensor_format == Float8BlockScaleTensorFormat::COMPACT; const auto rowwise_option = return_rowwise ? (compact ? detail::FP8BlockwiseRowwiseOption::ROWWISE_COMPACT : detail::FP8BlockwiseRowwiseOption::ROWWISE_GEMM_READY) : detail::FP8BlockwiseRowwiseOption::NONE; const auto columnwise_option = - return_columnwise - ? (compact ? detail::FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT - : detail::FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY) - : detail::FP8BlockwiseColumnwiseOption::NONE; + return_columnwise ? (compact ? detail::FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT + : detail::FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY) + : detail::FP8BlockwiseColumnwiseOption::NONE; detail::multi_quantize_transpose_vector_blockwise( input_list_, output_list_, quant_config_cpp.amax_epsilon, rowwise_option, columnwise_option, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index dbb95f96d0..2e74f897ca 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -166,8 +166,7 @@ void multi_tensor_quantize_impl(const std::vector &input_list, } if (blockwise_quantizer->dtype != first_blockwise_quantizer->dtype || - blockwise_quantizer->force_pow_2_scales != - first_blockwise_quantizer->force_pow_2_scales || + blockwise_quantizer->force_pow_2_scales != first_blockwise_quantizer->force_pow_2_scales || blockwise_quantizer->amax_epsilon != first_blockwise_quantizer->amax_epsilon || blockwise_quantizer->all_gather_usage != first_blockwise_quantizer->all_gather_usage || blockwise_quantizer->rowwise_usage != first_blockwise_quantizer->rowwise_usage ||