Skip to content

[Common] Let MXFP8 quantization always reduce dbias in colwise #3439

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kainzhong wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[Common] Let MXFP8 quantization always reduce dbias in colwise #3439

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f90b9f3
[Common] Let MXFP8 quantization always reduce dbias in colwise
kainzhong Aug 28, 2026
eb3f010
fix
kainzhong Aug 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 64 additions & 12 deletions transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,30 @@ __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 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<IType, float>);
// 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.
// 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<IType, float>);

// 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;
Expand Down Expand Up @@ -148,7 +171,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;
Expand Down Expand Up @@ -214,7 +237,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<IType, float>)) {
if constexpr (USE_HALF_PRECISION) {
IType thread_amax_f16 = static_cast<IType>(0.0f);
#pragma unroll
for (int i = 0; i < BUFF_DIM_Y; ++i) {
Expand All @@ -236,7 +259,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
float act_in_elt = static_cast<float>(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
Expand Down Expand Up @@ -287,18 +311,48 @@ __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<IType, float>)) {
if constexpr (USE_HALF_PRECISION) {
in = static_cast<float>(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;
out_colwise_data_sh[shmem_offset_elt] = static_cast<OType>(scaled_out);
}
}

// 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<float>(in_sh[shmem_offset_colwise]);
if constexpr (IS_ACT) {
elt = OP(elt, {});
}
if constexpr (IS_DACT) {
const float act_in_elt = static_cast<float>(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<IType>(elt);
}
}
}

if constexpr (ROWWISE_SCALING) {
const size_t shmem_offset_base_rowwise =
buff * BUFF_DIM + thread_offset_Y_rowwise * BUFF_DIM_X;
Expand All @@ -310,7 +364,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
Vec<IType2, PACK_SIZE / 2> in_IType[WAVES];

// 1. Read/Compute elements. Find MXFP8-block AMAX
if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v<IType, float>)) {
if constexpr (USE_HALF_PRECISION) {
IType2 thread_amax_2x = {static_cast<IType>(0.0f), static_cast<IType>(0.0f)};
#pragma unroll
for (int w = 0; w < WAVES; ++w) {
Expand Down Expand Up @@ -392,7 +446,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
Expand All @@ -419,9 +473,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<NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v<IType, float>),
IType, float>;
using AMax2DType = std::conditional_t<USE_HALF_PRECISION, IType, float>;
__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] =
Expand Down Expand Up @@ -469,7 +521,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
for (int e = 0; e < PACK_SIZE / 2; ++e) {
IType2 in;
OType2 &out_pair = reinterpret_cast<OType2 &>(out.data.elt[e]);
if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v<IType, float>)) {
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];
Expand Down Expand Up @@ -523,7 +575,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>();
Expand Down
Loading