diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index 5d7feeb6ce..d0a54f55a5 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -196,9 +196,14 @@ def fill(): return fill(), fill() -def run_quantize(method, act, x, ain, rowwise, columnwise, fp8_dtype, swizzled): +def run_quantize(method, act, x, ain, rowwise, columnwise, fp8_dtype, swizzled, use_2d_quant): """Quantize via the public dispatch; returns (mxfp8_tensor, dbias_or_None).""" - q = MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=rowwise, columnwise=columnwise) + q = MXFP8Quantizer( + fp8_dtype=fp8_dtype, + rowwise=rowwise, + columnwise=columnwise, + with_2d_quantization=use_2d_quant, + ) # Emit scales in the GEMM-swizzled layout (MXFP8QuantConfig::swizzled). q.optimize_for_gemm = swizzled if method == "CAST_ONLY": @@ -216,7 +221,7 @@ def run_quantize(method, act, x, ain, rowwise, columnwise, fp8_dtype, swizzled): raise ValueError(f"unknown method {method!r}") -def get_cfg_key(method, act, in_dtype, fp8_dtype, rowwise, colwise, swizzled): +def get_cfg_key(method, act, in_dtype, fp8_dtype, rowwise, colwise, swizzled, use_2d_quant): """Mirror of MXFP8QuantConfig::to_key (quantize_mxfp8_cutedsl.cuh): the name the CuTeDSL backend registers its compiled kernel under for this config. Used to check if the CuTeDSL implementation is registered """ @@ -229,8 +234,7 @@ def get_cfg_key(method, act, in_dtype, fp8_dtype, rowwise, colwise, swizzled): elif with_dact: desc = f"d{act['desc']}" # with_amax is hardcoded to False for now because currently MXFP8 quantizers never allocate an amax - # The last False is for use_2d_quantization; these cases never request 2D block scaling - flags = (rowwise, colwise, swizzled, False, with_dbias, with_dact, with_act, False) + flags = (rowwise, colwise, swizzled, False, with_dbias, with_dact, with_act, use_2d_quant) major, minor = device_compute_capability() return ( f"cutedsl_mxfp8_sm{major * 10 + minor}_" @@ -271,7 +275,15 @@ def extract_quantized_output(out, rowwise, columnwise, swizzled): def run_test_case( - method, act, shape, block_size, in_dtype, fp8_dtype, swizzled=False, input_offset=0 + method, + act, + shape, + block_size, + in_dtype, + fp8_dtype, + swizzled=False, + input_offset=0, + use_2d_quant=False, ): """Assert the CuTeDSL and CUDA backends produce bit-identical outputs for the same input and config. input_offset (in elements) shifts x's data pointer past the @@ -287,14 +299,14 @@ def run_test_case( set_cutedsl_backend(False) out_cuda, dbias_cuda = run_quantize( - method, act, x, act_input, rowwise, columnwise, fp8_dtype, swizzled + method, act, x, act_input, rowwise, columnwise, fp8_dtype, swizzled, use_2d_quant ) cuda_output = extract_quantized_output(out_cuda, rowwise, columnwise, swizzled) set_cutedsl_backend(True) try: out_cutedsl, dbias_cutedsl = run_quantize( - method, act, x, act_input, rowwise, columnwise, fp8_dtype, swizzled + method, act, x, act_input, rowwise, columnwise, fp8_dtype, swizzled, use_2d_quant ) cutedsl_output = extract_quantized_output(out_cutedsl, rowwise, columnwise, swizzled) finally: @@ -304,7 +316,7 @@ def run_test_case( # CuTeDSL backend supports, so its kernel must have been registered under the # config key. If not, the backend rejected or missed the config and the # comparison above was CUDA vs CUDA. - key = get_cfg_key(method, act, in_dtype, fp8_dtype, rowwise, columnwise, swizzled) + key = get_cfg_key(method, act, in_dtype, fp8_dtype, rowwise, columnwise, swizzled, use_2d_quant) assert tvm_ffi.get_global_func(key, allow_missing=True) is not None, ( f"CuTeDSL kernel not registered for {key}; the CuTeDSL backend fell back " "to CUDA and this case compared CUDA against itself" @@ -314,6 +326,8 @@ def run_test_case( tag = ( f"{method}/{act['name']}/{M}x{N}/{DTYPE_TO_STR[in_dtype]}/{FP8_TO_STR[fp8_dtype]}/{layout}" ) + if use_2d_quant: + tag += "/2d" for name, cuda_bytes in cuda_output.items(): assert torch.equal( cutedsl_output[name], cuda_bytes @@ -334,6 +348,18 @@ def test_cast_only(swizzled, fp8_dtype, in_dtype, block_size, shape): run_test_case("CAST_ONLY", IDENTITY, shape, block_size, in_dtype, fp8_dtype, swizzled) +# 2D block scaling (one scale per 32x32 block) is only reachable as a plain cast. +@pytest.mark.parametrize("shape", MATRIX_SIZES, ids=get_shape_id) +@pytest.mark.parametrize("block_size", BLOCK_SIZES, ids=get_block_id) +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id) +@pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) +def test_cast_only_2d(swizzled, fp8_dtype, in_dtype, block_size, shape): + run_test_case( + "CAST_ONLY", IDENTITY, shape, block_size, in_dtype, fp8_dtype, swizzled, use_2d_quant=True + ) + + def test_cast_only_rowwise_specialized_grid_y_overflow(): """Fall back when the specialized rowwise launch would exceed CUDA's grid.y limit.""" # The specialized kernel covers four rows per CTA, so this shape would require diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_2d.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_2d.py new file mode 100644 index 0000000000..fdca0825d4 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_2d.py @@ -0,0 +1,431 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MXFP8 2D block-scaling quantization kernel implemented in CuTeDSL. + +Ports the kIs2DBlockScaling path of quantize_mxfp8_kernel (quantize_mxfp8.cuh): every 32x32 +block of the input shares one E8M0 scale, which the rowwise scales repeat down the block's 32 +rows and the colwise scales repeat across its 32 columns. + +The TMA pipeline is the general G2S kernel's; only the per-tile math differs, so this subclasses +it and swaps in its own tile function. A tile is 32 rows by TILE_COLS, i.e. TILE_COLS // 32 +blocks side by side, and the colwise pass maps warp w onto columns [32w, 32w + 32), which is +exactly block w, so a warp reduction of the per-column amaxes yields the block amax. The rowwise +pass reads it back from SMEM, or, when there is no colwise pass, reduces its per-row amaxes over +the block's 32 rows itself. + +Only a plain cast is implemented: the C APIs that fuse dbias or an activation pass no +quantization config, so only plain casts (nvte_quantize_v2, nvte_multi_tensor_quantize) can +request 2D scaling. +""" + +from typing import Optional, Type + +import cutlass +from cutlass import cute +from cutlass import Float32, Int32, Int64, Uint16, Uint32, Uint8 + +from transformer_engine.common.CuTeDSL.utils import ( + device_compute_capability, + is_packed16, + abs_max_x2_bf16, + abs_max_x2_f16, + abs_max_scalar_bf16, + abs_max_scalar_f16, + to_f32_bf16, + to_f32_f16, + x2_lo_to_f32_bf16, + x2_lo_to_f32_f16, + x2_hi_to_f32_bf16, + x2_hi_to_f32_f16, + fabs_f32, + exp2f_rcp, + pack_f32x2, + unpack_i64_to_i32x2, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + get_cvt_f32x2_to_fp8x2_func, + cvt_f32_to_fp8e8m0fnu, + mul_f32x2_cvt_f32x4_to_fp8x4, + mul_f32x2_cvt_packed16x4_to_fp8x4, +) +from .quantize_mxfp8_common import ( + MXFP8QuantizeConfig, + MXFP8_BLOCK_SCALING_SIZE, + THREADS_PER_WARP, +) +from .quantize_mxfp8_g2s_general import MXFP8QuantizeGeneralG2SKernel + + +@cute.jit +def quantize_colwise_mxfp8_2d( + sX_tile: cute.Tensor, # (TILE_Y, TILE_X) input smem view, post-TMA + sO_col_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (colwise FP8 output) + mS_col_stage: cute.Tensor, # colwise scale tensor (1D swizzled, or 2D linear) + sBlockAmax: cute.Tensor, # receives the amax of block w at index w + MAX_NORM_RCP: cutlass.Constexpr[float], + tile_row_start: Int32, + tile_col_start: Int32, + M: Int32, + N: Int32, + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + SWIZZLE: cutlass.Constexpr[bool], + TILE_X: cutlass.Constexpr[int], + TILE_Y: cutlass.Constexpr[int], + SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool], +): + """Quantize one SMEM tile colwise with one scale per 32x32 block; returns the block amax.""" + tidx, _, _ = cute.arch.thread_idx() + + # Each thread handles one vertical MXFP8 block, so we have 1 thread along the Y direction since one tile only has + # one block in the Y (column) direction + CTA_THREADS_Y = MXFP8_BLOCK_SCALING_SIZE // TILE_Y # 1 thread along the Y direction + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_THREADS_Y, TILE_X), stride=(TILE_X, 1)), + val_layout=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE, 1), stride=(1, 1)), + ) + sX_thread = cute.composition(sX_tile, tv_layout)[tidx, None] + sO_thread = cute.composition(sO_col_tile, tv_layout)[tidx, None] + + # Gather this thread's 32 values from the MXFP8 block and its own amax. + # Note: TMA zero-fills past M and N, so the padding values never win when comparing amax. + in_c = [None] * MXFP8_BLOCK_SCALING_SIZE + if cutlass.const_expr(is_packed16(DTYPE)): + abs_max_scalar = abs_max_scalar_f16 if DTYPE is cutlass.Float16 else abs_max_scalar_bf16 + to_f32 = to_f32_f16 if DTYPE is cutlass.Float16 else to_f32_bf16 + in_c[0] = sX_thread[0] + amax_value = in_c[0] + for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE): + in_c[i] = sX_thread[i] + amax_value = abs_max_scalar(amax_value, in_c[i]) + amax_c = fabs_f32(to_f32(amax_value)) + else: + to_f32 = None + in_c[0] = Float32(sX_thread[0]) + amax_c = fabs_f32(in_c[0]) + for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE): + in_c[i] = Float32(sX_thread[i]) + amax_c = cute.arch.fmax(amax_c, fabs_f32(in_c[i])) + + # Reduce the amaxes of the 32 threads in this warp to get the block amax, and let the leader thread + # of each warp write it to SMEM. + # Use f32 redux if it is available (blackwell family); otherwise use butterfly shuffle for reduction + if cutlass.const_expr(device_compute_capability() in ((10, 0), (10, 3))): + block_amax = cute.arch.warp_redux_sync(amax_c, kind="fmax") + else: + block_amax = amax_c + for k in cutlass.range_constexpr(THREADS_PER_WARP.bit_length() - 1): + block_amax = cute.arch.fmax( + block_amax, cute.arch.shuffle_sync_bfly(block_amax, offset=1 << k) + ) + if tidx % THREADS_PER_WARP == 0: + sBlockAmax[tidx // THREADS_PER_WARP] = block_amax + + biased_exp_c = cvt_f32_to_fp8e8m0fnu(block_amax * MAX_NORM_RCP) + if cutlass.const_expr(SWIZZLE): + scale_coord = (0, tidx % 32, tidx // 32) + else: + scale_coord = (0, tidx) + if cutlass.const_expr(SKIP_SCALE_BOUNDS): + mS_col_stage[scale_coord] = biased_exp_c + elif tile_row_start < M and tile_col_start + tidx < N: + mS_col_stage[scale_coord] = biased_exp_c + + inv_scale_c = exp2f_rcp(biased_exp_c) + cvt_x2_func = get_cvt_f32x2_to_fp8x2_func(FP8_DTYPE) + sO_thread_fp8 = cute.make_tensor( + cute.recast_ptr(sO_thread.iterator, dtype=FP8_DTYPE), sO_thread.layout + ) + for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): + lo, hi = 2 * j, 2 * j + 1 + x_lo, x_hi = in_c[lo], in_c[hi] + if cutlass.const_expr(to_f32 is not None): + x_lo, x_hi = to_f32(x_lo), to_f32(x_hi) + pair = cvt_x2_func(x_hi * inv_scale_c, x_lo * inv_scale_c) + sO_thread_fp8[lo] = Uint8(pair & Uint16(0xFF)).bitcast(FP8_DTYPE) + sO_thread_fp8[hi] = Uint8(pair >> Uint16(8)).bitcast(FP8_DTYPE) + + return block_amax + + +@cute.jit +def quantize_rowwise_mxfp8_2d( + sX_tile: cute.Tensor, # (TILE_Y, TILE_X) input smem view, post-TMA + sO_row_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (rowwise FP8 output) + mS_row_stage: cute.Tensor, # rowwise scale tensor (1D swizzled, or 2D linear) + sBlockAmax: cute.Tensor, # block amaxes from the colwise pass, or scratch for the reduction + MAX_NORM_RCP: cutlass.Constexpr[float], + tile_row_start: Int32, + tile_col_start: Int32, + M: Int32, + N: Int32, + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + TILE_X: cutlass.Constexpr[int], + TILE_Y: cutlass.Constexpr[int], + SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool], + BLOCK_AMAX_IN_SMEM: cutlass.Constexpr[ + bool + ], # If colwise pass has already left the block amaxes in sBlockAmax +): + """Quantize one SMEM tile rowwise with one scale per 32x32 block; returns the block amax. + + With BLOCK_AMAX_IN_SMEM the colwise pass has already left the block amaxes in sBlockAmax + (and the caller synchronized); otherwise they are reduced here, which syncs the CTA. + """ + tidx, _, _ = cute.arch.thread_idx() + + # Same vectorized, bank-staggered loads as the general kernel's rowwise pass. + PACK_SIZE = 4 + WAVES = MXFP8_BLOCK_SCALING_SIZE // PACK_SIZE + THREADS_PER_BANK = 32 * 4 // MXFP8_BLOCK_SCALING_SIZE + + # Each thread covers one horizontal MXFP8 block, so we have TILE_X // MXFP8_BLOCK_SCALING_SIZE threads + # along X (row direction) + CTA_THREADS_X = TILE_X // MXFP8_BLOCK_SCALING_SIZE + NUM_WARPS = TILE_Y * CTA_THREADS_X // THREADS_PER_WARP + # A CTA consists of NUM_WARPS warps stacked together + thr_layout = cute.blocked_product( + cute.make_layout( + (THREADS_PER_WARP // CTA_THREADS_X, CTA_THREADS_X), stride=(CTA_THREADS_X, 1) + ), # Thread layout in a warp, (16, 2) + cute.make_layout((NUM_WARPS, 1)), # Warp layout in a CTA (2, 1) + ) + + _, tv_layout = cute.make_layout_tv( + thr_layout=thr_layout, + val_layout=cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), + ) + + sX_thread = cute.composition(sX_tile, tv_layout)[tidx, None] + sO_thread = cute.composition(sO_row_tile, tv_layout)[tidx, None] + sO_thread_u32 = cute.make_tensor( + cute.recast_ptr(sO_thread.iterator, dtype=Uint32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,)), + ) + + # Coordinates in thr_layout: (lane_y, lane_x) within the warp, warps stacked along Y. + warp_id = tidx // THREADS_PER_WARP + lane_id = tidx % THREADS_PER_WARP + lane_y = lane_id // CTA_THREADS_X + lane_x = lane_id % CTA_THREADS_X + thread_y = warp_id * (THREADS_PER_WARP // CTA_THREADS_X) + lane_y + thread_x = lane_x + + # Each thread start reading from the specfic bank based on its thread ID so they can do their best to access different banks + # to avoid bank conflict. + bank_group = (tidx % THREADS_PER_WARP) // THREADS_PER_BANK + # The offset this thread should start reading from based on what's its first bank to access. + offset = bank_group * PACK_SIZE + + if cutlass.const_expr(is_packed16(DTYPE)): + sX_thread_rw_i64 = cute.make_tensor( + cute.recast_ptr(sX_thread.iterator, dtype=Int64), + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE // 4), stride=(0, 1)), + ) + in_r = [[None, None] for _ in range(WAVES)] + for w in cutlass.range_constexpr(WAVES): + idx = (w + offset // 4) % (MXFP8_BLOCK_SCALING_SIZE // 4) + in_r[w][0], in_r[w][1] = unpack_i64_to_i32x2(sX_thread_rw_i64[0, idx]) + else: + sX_thread_rw = cute.make_tensor( + sX_thread.iterator, + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), + ) + in_r = [[None] * PACK_SIZE for _ in range(WAVES)] + for w in cutlass.range_constexpr(WAVES): + start = (w * PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE + for i in cutlass.range_constexpr(PACK_SIZE): + in_r[w][i] = Float32(sX_thread_rw[0, start + i]) + + if cutlass.const_expr(BLOCK_AMAX_IN_SMEM): + # If colwise pass has already writes the block amaxes into sBlockAmax, just read it back. + block_amax = sBlockAmax[thread_x] + else: + # Otherwise we need to compute the block amaxes in this rowwise pass + if cutlass.const_expr(is_packed16(DTYPE)): + abs_max_x2 = abs_max_x2_f16 if DTYPE is cutlass.Float16 else abs_max_x2_bf16 + x2_lo_to_f32 = x2_lo_to_f32_f16 if DTYPE is cutlass.Float16 else x2_lo_to_f32_bf16 + x2_hi_to_f32 = x2_hi_to_f32_f16 if DTYPE is cutlass.Float16 else x2_hi_to_f32_bf16 + # Load the fisst wave's two values into amax_2x, then fold in the rest of the waves. + amax_2x = abs_max_x2(in_r[0][0], in_r[0][1]) + for w in cutlass.range_constexpr(1, WAVES): + amax_2x = abs_max_x2(amax_2x, in_r[w][0]) + amax_2x = abs_max_x2(amax_2x, in_r[w][1]) + amax_r = cute.arch.fmax( + fabs_f32(x2_lo_to_f32(amax_2x)), fabs_f32(x2_hi_to_f32(amax_2x)) + ) + else: + vals = [in_r[w][i] for w in range(WAVES) for i in range(PACK_SIZE)] + amax_r = fabs_f32(vals[0]) + for j in cutlass.range_constexpr(1, len(vals)): + amax_r = cute.arch.fmax(amax_r, fabs_f32(vals[j])) + + # Use butterfly shuffles to reduce the 32 threads' local amax in this warp to get 2 block amax + # for the left half and the right half of the warp (each warp with shape (16, 1)) + for k in cutlass.range_constexpr((THREADS_PER_WARP // CTA_THREADS_X).bit_length() - 1): + amax_r = cute.arch.fmax( + amax_r, cute.arch.shuffle_sync_bfly(amax_r, offset=CTA_THREADS_X << k) + ) + + # Write the half-warp block amaxes to SMEM, then let the first warp reduce them to get the final block amax. + sWarpBlockAmax = cute.make_tensor( + sBlockAmax.iterator, + cute.make_layout((NUM_WARPS, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)), + ) + # Let the threads in the first row in each half warp write its partial warp amax to SMEM + # Warp shape is (16, 2) so each warp writes 2 partial amax for a (16, 32) block which will be reduced to produce + # the (32, 32) block amax + if lane_y == 0: + sWarpBlockAmax[(warp_id, lane_x)] = amax_r + cute.arch.sync_threads() + + # Reduce 2 warps' amax to the whole (32, 32) block's amax + block_amax = sWarpBlockAmax[(0, thread_x)] + for w in cutlass.range_constexpr(1, NUM_WARPS): + block_amax = cute.arch.fmax(block_amax, sWarpBlockAmax[(w, thread_x)]) + + biased_exp_r = cvt_f32_to_fp8e8m0fnu(block_amax * MAX_NORM_RCP) + if cutlass.const_expr(SKIP_SCALE_BOUNDS): + mS_row_stage[(thread_y, thread_x)] = biased_exp_r + elif tile_row_start + thread_y < M and tile_col_start + thread_x * MXFP8_BLOCK_SCALING_SIZE < N: + mS_row_stage[(thread_y, thread_x)] = biased_exp_r + + inv_scale_r = exp2f_rcp(biased_exp_r) + scale_2x = pack_f32x2(inv_scale_r, inv_scale_r) + if cutlass.const_expr(is_packed16(DTYPE)): + mul_cvt_x4_func = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, FP8_DTYPE) + else: + mul_cvt_x4_func = mul_f32x2_cvt_f32x4_to_fp8x4(FP8_DTYPE) + for w in cutlass.range_constexpr(WAVES): + idx = ((w * PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE) // 4 + if cutlass.const_expr(is_packed16(DTYPE)): + sO_thread_u32[idx] = mul_cvt_x4_func(in_r[w][0], in_r[w][1], scale_2x) + else: + sO_thread_u32[idx] = mul_cvt_x4_func( + in_r[w][0], in_r[w][1], in_r[w][2], in_r[w][3], scale_2x + ) + + return block_amax + + +@cute.jit +def quantize_mxfp8_smem_tile_2d( # pylint: disable=unused-argument + sX_tile: cute.Tensor, + sActInput_tile: Optional[cute.Tensor], + sO_row_tile: Optional[cute.Tensor], + mS_row_stage: Optional[cute.Tensor], + sO_col_tile: Optional[cute.Tensor], + mS_col_stage: Optional[cute.Tensor], + rowwise_dbias_acc: Optional[cute.Tensor], + tile_row_start: Int32, + tile_col_start: Int32, + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + WITH_AMAX: cutlass.Constexpr[bool], + MAX_NORM_RCP: cutlass.Constexpr[float], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], + TILE_X: cutlass.Constexpr[int], + TILE_Y: cutlass.Constexpr[int], + SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool], + smem: cutlass.memory.SmemAllocator, + **_fusion_args, # The general tile function's fusion knobs, all off for a plain cast. +): + """2D quantize a SMEM tile; returns (amax, None).""" + NUM_BLOCKS_X = TILE_X // MXFP8_BLOCK_SCALING_SIZE + NUM_WARPS = TILE_Y * NUM_BLOCKS_X // THREADS_PER_WARP + + # One amax per (warp, block) pair, for the rowwise-only cross-warp reduction. + @cute.struct + class BlockAmaxStorage: # pylint: disable=missing-class-docstring + sBlockAmax: cute.struct.MemRange[Float32, NUM_WARPS * NUM_BLOCKS_X] + + sBlockAmax = smem.allocate(BlockAmaxStorage).sBlockAmax.get_tensor( + cute.make_layout(NUM_WARPS * NUM_BLOCKS_X) + ) + + block_amax = Float32(0.0) + if cutlass.const_expr(COLWISE): + block_amax = quantize_colwise_mxfp8_2d( + sX_tile, + sO_col_tile, + mS_col_stage, + sBlockAmax, + MAX_NORM_RCP, + tile_row_start, + tile_col_start, + M, + N, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + SWIZZLE=WITH_GEMM_SWIZZLED_SCALES, + TILE_X=TILE_X, + TILE_Y=TILE_Y, + SKIP_SCALE_BOUNDS=SKIP_SCALE_BOUNDS, + ) + if cutlass.const_expr(ROWWISE): + # Publish the block amaxes to the rowwise pass. + cute.arch.sync_threads() + if cutlass.const_expr(ROWWISE): + block_amax = quantize_rowwise_mxfp8_2d( + sX_tile, + sO_row_tile, + mS_row_stage, + sBlockAmax, + MAX_NORM_RCP, + tile_row_start, + tile_col_start, + M, + N, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + TILE_X=TILE_X, + TILE_Y=TILE_Y, + SKIP_SCALE_BOUNDS=SKIP_SCALE_BOUNDS, + BLOCK_AMAX_IN_SMEM=COLWISE, + ) + amax = None + if cutlass.const_expr(WITH_AMAX): + amax = block_amax + return amax, None + + +class MXFP8Quantize2DKernel(MXFP8QuantizeGeneralG2SKernel): + """MXFP8 quantization with one scale per 32x32 block (plain cast only). + + Reuses the general G2S kernel's pipeline with its own tile function and tile configuration. + """ + + _quantize_smem_tile = staticmethod(quantize_mxfp8_smem_tile_2d) + + def __init__( + self, + cfg: MXFP8QuantizeConfig, + SKIP_INPUT_MASKING: bool = False, + SKIP_SCALE_BOUNDS: bool = False, + ): + assert cfg.USE_2D_QUANTIZATION + # 2D quantization is only supported for plain casts and no fusion should be requested + assert not (cfg.WITH_DBIAS or cfg.WITH_ACT or cfg.WITH_DACT) + + super().__init__(cfg, SKIP_INPUT_MASKING, SKIP_SCALE_BOUNDS) + + # Overrides the general kernel's tile configuration with the 2D kernel's + self._NUM_TILES = 2 + self._USE_STATIC_NUM_TILES = True + self._MIN_BLOCKS_PER_MP = 0 + + self._NUM_WARPS = 2 + self._THREADS_PER_CTA = self._NUM_WARPS * THREADS_PER_WARP + + # In rowwise pass, + self._TILE_ROWS = MXFP8_BLOCK_SCALING_SIZE + # In colwise pass, each thread covers one vertical MXFP8 block, and since THREADS_PER_WARP == MXFP8_BLOCK_SCALING_SIZE, + # two warps cover 2 * MXFP8_BLOCK_SCALING_SIZE columns + self._TILE_COLS = self._NUM_WARPS * MXFP8_BLOCK_SCALING_SIZE diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py index 8a66b55e48..6d72d50a88 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -139,8 +139,6 @@ def __init__( use_2d_quantization: bool = False, activation: Optional[str] = None, ): - if use_2d_quantization: - raise ValueError("2D block scaling is not implemented by the CuTeDSL MXFP8 kernels") if dtype is None or dtype not in ("Float32", "Float16", "BFloat16"): raise ValueError(f"unknown input dtype {dtype!r}; expected Float32|Float16|BFloat16") self.DTYPE = str_to_cutlass_dtype(dtype) @@ -193,6 +191,11 @@ def __init__( # uses a cross-thread smem reduction over THREADS_Y. Both mirror the CUDA # kernel's COLWISE_SCALING / rowwise dbias branches. self.WITH_DBIAS = with_dbias + # The C APIs that fuse dbias or an activation pass no quantization config, so 2D block + # scaling only reaches plain casts and fusion is not supported. + if use_2d_quantization and (with_dbias or with_dact or with_act): + raise ValueError("2D block scaling is only implemented for a plain cast") + self.USE_2D_QUANTIZATION = use_2d_quantization self.MAX_NORM_RCP = ( FP8E4M3_MAX_NORM_RCP if fp8_dtype == "Float8E4M3" else FP8E5M2_MAX_NORM_RCP ) @@ -203,7 +206,8 @@ def __str__(self): f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, " f"swizzled={self.WITH_GEMM_SWIZZLED_SCALES}, with_amax={self.WITH_AMAX}, " f"with_dbias={self.WITH_DBIAS}, with_dact={self.WITH_DACT}, " - f"with_act={self.WITH_ACT}, activation={self.ACTIVATION})" + f"with_act={self.WITH_ACT}, use_2d_quantization={self.USE_2D_QUANTIZATION}, " + f"activation={self.ACTIVATION})" ) __repr__ = __str__ @@ -212,6 +216,14 @@ def is_plain_cast(self) -> bool: """Whether cfg is a plain cast: no amax, dbias or (d)activation fused in.""" return not (self.WITH_AMAX or self.WITH_DBIAS or self.WITH_DACT or self.WITH_ACT) + def act_needs_masking(self) -> bool: + """Whether TMA's zero padding must be masked after the activation, i.e. act(0) != 0. + + All activations supported by MXFP8 satisfy act(0) == 0, so this is a precaution for + activations added in the future. + """ + return self.WITH_ACT and self.ACTIVATION not in ("relu", "gelu", "silu", "qgelu", "srelu") + class MXFP8QuantizeKernelBase(abc.ABC): """Base class for MXFP8 quantize kernels.""" @@ -268,34 +280,20 @@ def __init__(self, cfg: MXFP8QuantizeConfig): from .quantize_mxfp8_g2s_bidim import MXFP8QuantizeBidimensionalG2SKernel from .quantize_mxfp8_g2r_rowwise_2lane import MXFP8QuantizeRowwise2LaneG2RKernel from .quantize_mxfp8_g2r_bidim import MXFP8QuantizeBidimensionalG2RKernel + from .quantize_mxfp8_2d import MXFP8Quantize2DKernel # pylint: enable=import-outside-toplevel self.cfg = cfg - # These activation functions satisfy f(0) = 0 so we don't need to mask with OOB regions - # (zeros filled by TMA are still zeros without applying activation to them) - # Note that as the time of writing, all activations supported by MXFP8 satisfy this property, so !ACT_NEED_MASKING is never reachable. - # This is only a precaution for future use cases if someone wants to add a new activation function. - self.ACT_NEED_MASKING = cfg.WITH_ACT and cfg.ACTIVATION not in ( - "relu", - "gelu", - "silu", - "qgelu", - "srelu", - ) - # Instantiate all possible kernels at compile time, - # and we will pick the right one at runtime based on the input shape and config. - self.general_skip_input_masking_kernel = MXFP8QuantizeGeneralG2SKernel( - cfg, SKIP_INPUT_MASKING=True - ) - self.general_skip_input_masking_and_scale_bounds_kernel = MXFP8QuantizeGeneralG2SKernel( - cfg, SKIP_INPUT_MASKING=True, SKIP_SCALE_BOUNDS=True + # 2D block scaling only swaps the general kernel's per-tile math. + general_cls = ( + MXFP8Quantize2DKernel if cfg.USE_2D_QUANTIZATION else MXFP8QuantizeGeneralG2SKernel ) - # We only need to mask when WITH_ACT is enabled because with activations applied zeros filled by TMA affect block statistics - self.general_kernel = MXFP8QuantizeGeneralG2SKernel(cfg) if self.ACT_NEED_MASKING else None - - # Specialized kernels in a specific dispatch priority order so we prefer to dispatch to kernels that come first in the list. - self.specialized_kernels = [ + # Candidate kernels in dispatch priority order; the first one whose runtime + # eligible_to_dispatch check passes is launched. The specialized kernels come first, + # then the general kernel without and with boundary handling. The last one is always + # eligible. + self.preferred_kernels = [ kernel_cls(cfg) for kernel_cls in ( MXFP8QuantizeRowwise2LaneG2RKernel, @@ -304,6 +302,9 @@ def __init__(self, cfg: MXFP8QuantizeConfig): MXFP8QuantizeBidimensionalG2SKernel, ) if kernel_cls.supports(cfg) + ] + [ + general_cls(cfg, SKIP_INPUT_MASKING=True, SKIP_SCALE_BOUNDS=True), + general_cls(cfg, SKIP_INPUT_MASKING=not cfg.act_needs_masking()), ] @cute.jit @@ -320,13 +321,13 @@ def __call__( mWorkspace: Optional[cute.Tensor], stream: CUstream, ): - dispatched_to_specialized = False - for kernel in self.specialized_kernels: - if not dispatched_to_specialized: + dispatched = False + for kernel in self.preferred_kernels: + if not dispatched: if kernel.eligible_to_dispatch( self.cfg, mX, mO_row, mS_row, mO_col, mS_col, mNoop, mDActInput ): - dispatched_to_specialized = True + dispatched = True kernel( mX, mO_row, @@ -339,69 +340,6 @@ def __call__( mWorkspace, stream, ) - # If not using a specialized kernel, fall back to the general kernel - if not dispatched_to_specialized: - # With _USE_STATIC_NUM_TILES, each CTA always processes _NUM_TILES tiles, including OOB ones, - # so the scale-bounds check can only be dropped if M divides TILE_ROWS * _NUM_TILES. - num_tiles_per_cta = ( - self.general_skip_input_masking_kernel._NUM_TILES - if self.general_skip_input_masking_kernel._USE_STATIC_NUM_TILES - else 1 - ) - # If the input shape can be perfectly tiled by the general kernel's tile size, we can skip some boundary check because - # we know we will not touch any out-of-bounds region. - shape_is_divisible = ( - mX.shape[0] - % (self.general_skip_input_masking_kernel._TILE_ROWS * num_tiles_per_cta) - == 0 - and mX.shape[1] % self.general_skip_input_masking_kernel._TILE_COLS == 0 - ) - # If shape is divisible, then we won't access OOB regions regardless whatever - if shape_is_divisible: - self.general_skip_input_masking_and_scale_bounds_kernel( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - else: - # Otherwise we can't skip scale bound check, but we may still skip masking activation inputs with zero - if cutlass.const_expr(self.ACT_NEED_MASKING): - # Masking and the scale guards now share one condition, so this is a plain - # two-way choice: either the shape tiles exactly (drop both) or it does not. - self.general_kernel( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - else: - # We still need to check the scale bounds, but we can skip some activations masking because their output is 0 - # for OOB regions where TMA fills with zeros - self.general_skip_input_masking_kernel( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) def compile_cutedsl_function_from_cfg(cfg): diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 61598a7d37..3cb611e775 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -609,7 +609,8 @@ class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. return ( - not cfg.WITH_AMAX + not cfg.USE_2D_QUANTIZATION + and not cfg.WITH_AMAX and not cfg.WITH_DBIAS and cfg.DTYPE is cutlass.BFloat16 and cfg.COLWISE diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py index 0c287f02ab..fb325f289b 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py @@ -219,7 +219,13 @@ class MXFP8QuantizeRowwise1LaneG2RKernel(MXFP8QuantizeKernelBase): @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: - return cfg.is_plain_cast() and is_packed16(cfg.DTYPE) and cfg.ROWWISE and not cfg.COLWISE + return ( + not cfg.USE_2D_QUANTIZATION + and cfg.is_plain_cast() + and is_packed16(cfg.DTYPE) + and cfg.ROWWISE + and not cfg.COLWISE + ) @staticmethod @cute.jit diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py index 03d5d510ab..19a879ce72 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -262,7 +262,8 @@ class MXFP8QuantizeRowwise2LaneG2RKernel(MXFP8QuantizeKernelBase): def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. return ( - not cfg.WITH_AMAX + not cfg.USE_2D_QUANTIZATION + and not cfg.WITH_AMAX and not cfg.WITH_DBIAS and cfg.DTYPE is cutlass.BFloat16 and cfg.ROWWISE diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py index 5f0b3d2aed..5e769cc84c 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py @@ -255,7 +255,8 @@ class MXFP8QuantizeBidimensionalG2SKernel(MXFP8QuantizeKernelBase): def supports(cfg: MXFP8QuantizeConfig) -> bool: # This kernel supports swizzled scales but we don't dispatch to it for now. return ( - cfg.is_plain_cast() + not cfg.USE_2D_QUANTIZATION + and cfg.is_plain_cast() and is_packed16(cfg.DTYPE) and cfg.ROWWISE and cfg.COLWISE diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py index 81148416fe..da73da2d77 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py @@ -29,7 +29,7 @@ import cutlass from cutlass import cute from cutlass import pipeline -from cutlass import Float32, Int32, Int64, Uint16, Uint32, Uint8 +from cutlass import Boolean, Float32, Int32, Int64, Uint16, Uint32, Uint8 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( @@ -576,6 +576,7 @@ def quantize_mxfp8_smem_tile_general( DBIAS_REDUCTION_IN_COLWISE: cutlass.Constexpr[bool], COLWISE_DBIAS_REDUCTION_ONLY: cutlass.Constexpr[bool], DBIAS_REDUCTION_IN_ROWWISE: cutlass.Constexpr[bool], + smem: Optional[cutlass.memory.SmemAllocator] = None, # pylint: disable=unused-argument ): """Quantize one SMEM tile in the directions enabled by ROWWISE / COLWISE; returns (amax, dbias). @@ -665,6 +666,9 @@ class MXFP8QuantizeGeneralG2SKernel(MXFP8QuantizeKernelBase): """ _NUM_STAGES = 2 # The pipeline depth is always 2 + _quantize_smem_tile = staticmethod( + quantize_mxfp8_smem_tile_general + ) # Function to quantize one SMEM tile def __init__( self, @@ -726,6 +730,38 @@ def __init__( and cfg.ACTIVATION != "relu" ) + @cute.jit + def eligible_to_dispatch( # pylint: disable=arguments-differ,unused-argument + self, + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + ) -> Boolean: + """Runtime check: dropping the scale bounds checks, or the input masking an activation + with act(0) != 0 needs, is only safe when the CTAs tile the input exactly.""" + eligible = True + # If we skip scale OOB check, or we skip input masking with an activation that doesn't satisfy act(0)=0, + # then it is required to have the input shape perfectly divisible + if cutlass.const_expr( + self.SKIP_SCALE_BOUNDS or (self.SKIP_INPUT_MASKING and cfg.act_needs_masking()) + ): + M, N = mX.shape + # With _USE_STATIC_NUM_TILES, each CTA always processes _NUM_TILES tiles, including OOB ones, + # so the scale-bounds check can only be dropped if M divides TILE_ROWS * _NUM_TILES. + num_tiles_per_cta = self._NUM_TILES if self._USE_STATIC_NUM_TILES else 1 + # If the input shape can be perfectly tiled by the general kernel's tile size, we can skip some boundary check because + # we know we will not touch any out-of-bounds region. + shape_is_divisible = ( + M % (self._TILE_ROWS * num_tiles_per_cta) == 0 and N % self._TILE_COLS == 0 + ) + eligible = shape_is_divisible + return eligible + @cute.jit def __call__( self, @@ -1217,7 +1253,7 @@ class DactStorage: sO_col_tile = sO_col[(None, stage_idx)] mS_col_stage = cute.flatten(mS_col[(None, (tile_idx_y, tile_idx_x))]) - amax, dbias = quantize_mxfp8_smem_tile_general( + amax, dbias = self._quantize_smem_tile( sX_tile, sActInput_tile, sO_row_tile, @@ -1247,6 +1283,7 @@ class DactStorage: DBIAS_REDUCTION_IN_COLWISE=self.DBIAS_REDUCTION_IN_COLWISE, COLWISE_DBIAS_REDUCTION_ONLY=self.COLWISE_DBIAS_REDUCTION_ONLY, DBIAS_REDUCTION_IN_ROWWISE=self.DBIAS_REDUCTION_IN_ROWWISE, + smem=smem, ) if cutlass.const_expr(cfg.WITH_AMAX): per_thread_amax = cute.arch.fmax(per_thread_amax, amax) diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh index 6380dda21d..4356b365e9 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh @@ -231,11 +231,6 @@ bool mxfp8_quantize_cutedsl(const Tensor *input_tensor, const Tensor *act_input_ if (!tvm_ffi_bridge::TVMFFICentral::getInstance().get_cutedsl_backend_enabled()) { return false; } - // TODO(kainingz): port 2D quantization to CuTeDSL - if (use_2d_quantization) { - // report_cutedsl_not_chosen("2D quantization is not supported."); - return false; - } checkCuDriverContext(stream); const MXFP8QuantConfig config{/*dtype=*/input_tensor->dtype(), /*fp8_dtype=*/output_tensor->dtype(),