From f90b9f3af8764be7c7e9d3bd45e68555ff0abd05 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 28 Aug 2026 19:32:33 +0000 Subject: [PATCH 1/2] [Common] Let MXFP8 quantization always reduce dbias in colwise Signed-off-by: Kaining Zhong --- .../common/cast/mxfp8/quantize_mxfp8.cuh | 74 ++++++++++++++++--- 1 file changed, 62 insertions(+), 12 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 9f312ac3f5..90a92cc211 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -80,7 +80,28 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) constexpr size_t STAGES = CHUNK_DIM_Y / BUFF_DIM_Y; static_assert(STAGES >= 1); - constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS && ROWWISE_SCALING && COLWISE_SCALING; + // If columnwise is quantized and dbias is fused, do dbias reduction in colwise which is easier + // (no cross-thread reduction needed). + constexpr bool DBIAS_REDUCTION_IN_COLWISE = IS_DBIAS && COLWISE_SCALING; + // If columnwise is not quantized, and dbias is fused, we will do a columnwise reduction only pass + // (may also cache dact values) without quantization because doing dbias in rowwise requires a SMEM shuffle + // which is even slower + constexpr bool DBIAS_REDUCTION_COLWISE_ONLY = IS_DBIAS && (!COLWISE_SCALING); + // Never do dbias reduction in rowwise because it is slow. Leave this so in case it's enabled in the future + constexpr bool DBIAS_REDUCTION_IN_ROWWISE = false; + + // Fast path: keep the elements in BF16/FP16 and compute the AMAX with the half-precision + // abs-max instead of upcasting every element to FP32. Only possible without activations. + // dbias does not disable it: the partial sums are accumulated in the scaling loop below, which + // upcasts the elements to FP32 anyway to feed the cvt. + constexpr bool USE_HALF_PRECISION = NO_ACTIVATIONS && (!std::is_same_v); + + // Cache activations in-place in the SMEM input tile so the activation is computed only once, in + // the direction we favor (columnwise), and the rowwise pass reads the cached value back instead + // of recomputing it. The columnwise reduction-only pass computes the same values, so it can + // populate the cache too. + constexpr bool IS_CACHED_ACT_OP = + COMPUTE_ACTIVATIONS && ROWWISE_SCALING && (COLWISE_SCALING || DBIAS_REDUCTION_COLWISE_ONLY); const size_t block_offset_Y = blockIdx.y * CHUNK_DIM_Y; const size_t block_offset_X = blockIdx.x * CHUNK_DIM_X; @@ -148,7 +169,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) float partial_dbias_colwise = 0.0f; float thread_dbias_rowwise[SCALE_DIM_X]; - if constexpr (IS_DBIAS) { + if constexpr (DBIAS_REDUCTION_IN_ROWWISE) { #pragma unroll for (int j = 0; j < SCALE_DIM_X; ++j) { thread_dbias_rowwise[j] = 0.0f; @@ -214,7 +235,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) IType in_colwise_IType[BUFF_DIM_Y]; // 1. Read/Compute elements. Find MXFP8-block AMAX - if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { + if constexpr (USE_HALF_PRECISION) { IType thread_amax_f16 = static_cast(0.0f); #pragma unroll for (int i = 0; i < BUFF_DIM_Y; ++i) { @@ -236,7 +257,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) float act_in_elt = static_cast(act_in_sh[shmem_offset_colwise]); elt *= OP(act_in_elt, {}); } - if constexpr (IS_DBIAS) { + if constexpr (DBIAS_REDUCTION_IN_COLWISE) { + // Accumulate before the truncation below so the partial sums stay full precision partial_dbias_colwise += elt; } // Numerical truncation: Downcast to IType (BF16/FP16), then upcast it back to FP32 @@ -287,11 +309,16 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) #pragma unroll for (int i = 0; i < SCALE_DIM_Y; ++i) { float in; - if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { + if constexpr (USE_HALF_PRECISION) { in = static_cast(in_colwise_IType[i]); } else { in = in_compute_colwise[i]; } + // On the half-precision path the read loop kept the elements in IType, so dbias is + // accumulated here instead, reusing the FP32 value the cvt needs anyway. + if constexpr (DBIAS_REDUCTION_IN_COLWISE && USE_HALF_PRECISION) { + partial_dbias_colwise += in; + } const float scaled_out = in * block_scale_inverse; const size_t shmem_offset_elt = shmem_offset_base_colwise + i * BUFF_DIM_X; @@ -299,6 +326,31 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } } + // If the columnwise direction is not quantized but dbias is fused, run a columnwise + // reduction-only pass (no quantization). When activations are fused, this pass also caches the + // post-activation values so that the rowwise pass below does not recompute them. + if constexpr (DBIAS_REDUCTION_COLWISE_ONLY) { + const size_t shmem_offset_base_colwise = buff * BUFF_DIM + tid_X_colwise; +#pragma unroll + for (int i = 0; i < BUFF_DIM_Y; ++i) { + const size_t shmem_offset_colwise = shmem_offset_base_colwise + i * BUFF_DIM_X; + + float elt = static_cast(in_sh[shmem_offset_colwise]); + if constexpr (IS_ACT) { + elt = OP(elt, {}); + } + if constexpr (IS_DACT) { + const float act_in_elt = static_cast(act_in_sh[shmem_offset_colwise]); + elt *= OP(act_in_elt, {}); + } + partial_dbias_colwise += elt; + // Cache computed activations to avoid computing them again in the rowwise pass + if constexpr (IS_CACHED_ACT_OP) { + cached_act_sh[shmem_offset_colwise] = static_cast(elt); + } + } + } + if constexpr (ROWWISE_SCALING) { const size_t shmem_offset_base_rowwise = buff * BUFF_DIM + thread_offset_Y_rowwise * BUFF_DIM_X; @@ -310,7 +362,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) Vec in_IType[WAVES]; // 1. Read/Compute elements. Find MXFP8-block AMAX - if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { + if constexpr (USE_HALF_PRECISION) { IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; #pragma unroll for (int w = 0; w < WAVES; ++w) { @@ -392,7 +444,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } // If DBIAS was computed in the 1st pass (COLWISE) then no need to compute it again - if constexpr (IS_DBIAS && (!COLWISE_SCALING)) { + if constexpr (DBIAS_REDUCTION_IN_ROWWISE) { thread_dbias_rowwise[j] += elt; } // Numerical truncation: Downcast to IType (BF16/FP16), then upcast it back to FP32 @@ -419,9 +471,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) // 2. Compute E8M0 scaling factor e8m0_t biased_exponent; if constexpr (kIs2DBlockScaling) { - using AMax2DType = - std::conditional_t), - IType, float>; + using AMax2DType = std::conditional_t; __shared__ e8m0_t block_scales_2d[THREADS_X]; __shared__ AMax2DType block_amax_2d[THREADS_X * THREADS_Y]; block_amax_2d[tid_X_rowwise * THREADS_Y + tid_Y_rowwise] = @@ -469,7 +519,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) for (int e = 0; e < PACK_SIZE / 2; ++e) { IType2 in; OType2 &out_pair = reinterpret_cast(out.data.elt[e]); - if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { + if constexpr (USE_HALF_PRECISION) { in = in_IType[w].data.elt[e]; } else if constexpr (IS_CACHED_ACT_OP) { in.x = in_cached[w].data.elt[2 * e]; @@ -523,7 +573,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) if constexpr (IS_DBIAS) { float thread_partial_dbias = 0.0f; - if constexpr (COLWISE_SCALING) { + if constexpr (!DBIAS_REDUCTION_IN_ROWWISE) { thread_partial_dbias = partial_dbias_colwise; } else { ptx::cp_async_bulk_wait_group_read<0>(); From eb3f010d8bedb8e77083bc8fdbfeffd48a0d015d Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 28 Aug 2026 23:48:00 +0000 Subject: [PATCH 2/2] fix Signed-off-by: Kaining Zhong --- .../common/cast/mxfp8/quantize_mxfp8.cuh | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 90a92cc211..2b8826f93d 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -84,11 +84,13 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) // (no cross-thread reduction needed). constexpr bool DBIAS_REDUCTION_IN_COLWISE = IS_DBIAS && COLWISE_SCALING; // If columnwise is not quantized, and dbias is fused, we will do a columnwise reduction only pass - // (may also cache dact values) without quantization because doing dbias in rowwise requires a SMEM shuffle - // which is even slower - constexpr bool DBIAS_REDUCTION_COLWISE_ONLY = IS_DBIAS && (!COLWISE_SCALING); - // Never do dbias reduction in rowwise because it is slow. Leave this so in case it's enabled in the future - constexpr bool DBIAS_REDUCTION_IN_ROWWISE = false; + // (may also cache dact values) without quantization because doing dbias in rowwise is slower + // With fp32 input and DBIAS+DACT in which this approach is slower so we exclude that case specifically. + constexpr bool DBIAS_REDUCTION_COLWISE_ONLY = + IS_DBIAS && (!COLWISE_SCALING) && !(COMPUTE_ACTIVATIONS && std::is_same_v); + // Rowwise reduction is usually slower unless under the exception mentioned above. + constexpr bool DBIAS_REDUCTION_IN_ROWWISE = + IS_DBIAS && (!COLWISE_SCALING) && !DBIAS_REDUCTION_COLWISE_ONLY; // Fast path: keep the elements in BF16/FP16 and compute the AMAX with the half-precision // abs-max instead of upcasting every element to FP32. Only possible without activations.