From 7c3abf6bdc9442b8645e4d639c603c1b8fc98cdc Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 7 Oct 2026 07:22:07 +0000 Subject: [PATCH 01/29] [Common] More CuTeDSL MXFP8 special kernels Signed-off-by: Kaining Zhong --- transformer_engine/common/CuTeDSL/__init__.py | 2 +- .../common/CuTeDSL/cast/mxfp8/__init__.py | 2 +- .../cast/mxfp8/quantize_mxfp8_common.py | 685 ++++++++ .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 490 ++++++ .../mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py | 265 +++ .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 368 ++++ .../mxfp8/quantize_mxfp8_g2s_bidimensional.py | 712 ++++++++ ...ize_mxfp8.py => quantize_mxfp8_general.py} | 1498 +---------------- transformer_engine/common/CuTeDSL/utils.py | 20 + .../common/CuTeDSL/utils_fp8.py | 43 + 10 files changed, 2603 insertions(+), 1482 deletions(-) create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py rename transformer_engine/common/CuTeDSL/cast/mxfp8/{quantize_mxfp8.py => quantize_mxfp8_general.py} (50%) diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index 50d536852c2..e3920b13486 100644 --- a/transformer_engine/common/CuTeDSL/__init__.py +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -10,7 +10,7 @@ import tvm_ffi -from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8 import ( +from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8_common import ( get_mxfp8_quantization_function, ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py index 6491573690f..444d85234ec 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py @@ -4,4 +4,4 @@ """CuTeDSL MXFP8 quantization kernels.""" -from . import quantize_mxfp8 +from . import quantize_mxfp8_common diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py new file mode 100644 index 00000000000..f281937ffca --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -0,0 +1,685 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared definitions, kernel dispatch and TVM-FFI registration for CuTeDSL MXFP8.""" + +import abc +import logging +import os +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import Boolean, Float32, Int32, Int64, Uint32 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module +import tvm_ffi + +from transformer_engine.common.CuTeDSL.utils import ( + str_to_cutlass_dtype, + abs_max_x2_bf16, + device_compute_capability, + is_packed16, +) +from transformer_engine.common.CuTeDSL.activations import ( + act_relu, + act_gelu, + act_silu, + act_qgelu, + act_srelu, + dact_drelu, + dact_dsrelu, + dact_dsilu, + dact_dqgelu, + dact_dgelu, +) + +CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" + +logger = logging.getLogger("transformer_engine.cutedsl.mxfp8") + +# Number of elements per MXFP8 scale block. They will share the same E8M0 scale factor +MXFP8_BLOCK_SCALING_SIZE = 32 +# How many threads are in one warp +THREADS_PER_WARP = 32 + +# Prefer the current allocator API while retaining compatibility with older DSL releases. +_SmemAllocator = getattr(cutlass, "memory", cutlass.utils).SmemAllocator + +# FP8E4M3 max representable value +FP8E4M3_MAX_NORM = 448.0 +FP8E4M3_MAX_NORM_RCP = 1.0 / FP8E4M3_MAX_NORM +FP8E5M2_MAX_NORM = 57344.0 +FP8E5M2_MAX_NORM_RCP = 1.0 / FP8E5M2_MAX_NORM + +# If N is not divisible by 16, writing quantized FP8 output with N bytes per row will fail because +# TMA requires 16-byte alignment for each row +SYM_N_DIVISIBILITY = 16 + +SUPPORTED_ACTIVATIONS = { + "relu": act_relu, + "gelu": act_gelu, + "silu": act_silu, + "qgelu": act_qgelu, + "srelu": act_srelu, +} + +SUPPORTED_DACTIVATIONS = { + "drelu": dact_drelu, + "dgelu": dact_dgelu, + "dsilu": dact_dsilu, + "dqgelu": dact_dqgelu, + "dsrelu": dact_dsrelu, +} + + +@cute.jit +def derive_swizzled_scale_layout( + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + mS_row: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], +): + """Derive the swizzled layout for the rowwise and colwise scale tensors.""" + num_scale_cols = cute.ceil_div(N, MXFP8_BLOCK_SCALING_SIZE) + num_scale_rows = cute.ceil_div(M, MXFP8_BLOCK_SCALING_SIZE) + + num_tiles_M = cute.ceil_div(M, 128) + num_tiles_SC = cute.ceil_div(num_scale_cols, 4) + num_tiles_SR = cute.ceil_div(num_scale_rows, 4) + num_tiles_N = cute.ceil_div(N, 128) + + if cutlass.const_expr(ROWWISE): + mS_row = cute.make_tensor( + mS_row.iterator, + cute.make_layout( + ((32, 4, num_tiles_M), (4, num_tiles_SC)), + stride=((16, 4, num_tiles_SC * 512), (1, 512)), + ), + ) + if cutlass.const_expr(COLWISE): + mS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout( + ((4, num_tiles_SR), (32, 4, num_tiles_N)), + stride=((1, 512), (16, 4, num_tiles_SR * 512)), + ), + ) + return mS_row, mS_col + + +@cute.jit +def noop_flag_is_set(mNoop: cute.Pointer) -> Boolean: + """Whether the cast_noop flag says this quantization is a no-op and must be skipped. + + mNoop is a pointer rather than a tensor so that one compiled kernel serves both a present and + an absent flag, hence the address is checked before it is dereferenced, exactly like the CUDA + C++ kernel's `noop != nullptr && noop[0] == 1.0f`. The two checks cannot be joined with `and`, + which the DSL lowers to a non-short-circuiting op that would load from the null pointer. + """ + flag_is_set = Boolean(False) + if mNoop.toint() != Int64(0): + flag_is_set = cute.make_tensor(mNoop, cute.make_layout((1,)))[0] == Float32(1.0) + return flag_is_set + + +class MXFP8QuantizeConfig: + """Configs for the compiled CuTeDSL kernel. These will be fixed once the kernel is compiled and + they will behave as const expressions. + """ + + def __init__( + self, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + with_amax: bool, + with_dbias: bool = False, + with_dact: bool = False, + with_act: bool = False, + 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) + self.DTYPE_STR = dtype # readable input-dtype token, for __str__ + if fp8_dtype not in ("Float8E4M3", "Float8E5M2"): + raise ValueError( + f"unknown FP8 dtype {fp8_dtype!r}; expected 'Float8E4M3' or 'Float8E5M2'" + ) + self.FP8_DTYPE = str_to_cutlass_dtype(fp8_dtype) + self.FP8_DTYPE_STR = fp8_dtype # readable token, for __str__ + self.ROWWISE = rowwise + self.COLWISE = colwise + if not (rowwise or colwise): + raise ValueError("at least one of rowwise or colwise must be true") + self.WITH_GEMM_SWIZZLED_SCALES = with_gemm_swizzled_scales + self.WITH_AMAX = with_amax + if not with_dact and not with_act: + if activation == "none": + self.ACTIVATION = None + else: + raise ValueError( + "activation must be none when with_dact and with_act are both False" + ) + else: + if with_dact and with_act: + raise ValueError( + "with_dact and with_act cannot be true at the same time since they are used for" + " different paths (bwd vs fwd)" + ) + if with_dact: + if activation in SUPPORTED_DACTIVATIONS: + self.ACTIVATION = activation + else: + raise ValueError( + f"unknown activation {activation!r} for with_dact=True; expected one of" + f" {sorted(SUPPORTED_DACTIVATIONS)}" + ) + elif with_act: + if activation in SUPPORTED_ACTIVATIONS: + self.ACTIVATION = activation + else: + raise ValueError( + f"unknown activation {activation!r} for with_act=True; expected one of" + f" {sorted(SUPPORTED_ACTIVATIONS)}" + ) + self.WITH_DACT = with_dact + self.WITH_ACT = with_act + # dbias is the column reduction of the (post-act/dact) element. With colwise + # output each thread owns a full column (trivial reduction); rowwise-only + # 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 + self.MAX_NORM_RCP = ( + FP8E4M3_MAX_NORM_RCP if fp8_dtype == "Float8E4M3" else FP8E5M2_MAX_NORM_RCP + ) + + def __str__(self): + return ( + f"MXFP8QuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " + 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})" + ) + + __repr__ = __str__ + + +class MXFP8QuantizeKernelBase(abc.ABC): + """Base class for MXFP8 quantize kernels.""" + + @abc.abstractmethod + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + """ + Compiled kernel entrypoint (decorate with @cute.jit). + All MXFP8 quantize kernels must implement this interface because this is our C++ call site's contract in `quantize_mxfp8_cutedsl.cuh`. + C++ call site will pass the arguments in this exact order via tvm-ffi and the kernel must accept them even if they are not used. + """ + + +@cute.jit +def bf16_pair_magnitude(pair: Int32) -> Uint32: + """Fold both BF16 halves and clear the sign set by max.xorsign.abs.""" + swapped = cute.arch.inline_ptx( + "prmt.b32 {$w0}, {$r0}, {$r0}, 0x1032;", + write_only_types=[Int32], + read_only_args=[pair], + ) + return Uint32(abs_max_x2_bf16(pair, swapped)) & Uint32(0x7FFF) + + +class MXFP8QuantizeEntry(MXFP8QuantizeKernelBase): + """Select the appropriate MXFP8 quantization kernel based on the configuration and runtime shapes.""" + + def __init__(self, cfg: MXFP8QuantizeConfig): + # Kernels import shared definitions from this module. Resolve their + # implementations only when constructing the compile-time entry object. + # pylint: disable=import-outside-toplevel + from .quantize_mxfp8_general import MXFP8QuantizeKernel + from .quantize_mxfp8_g2r_rowwise_1lane import MXFP8QuantizeG2RRowwise1LaneKernel + from .quantize_mxfp8_g2s_bidimensional import MXFP8QuantizeSpecializedBidimensionalKernel + from .quantize_mxfp8_g2r_rowwise_2lane import MXFP8QuantizeG2RRowwise2LaneKernel + from .quantize_mxfp8_g2r_bidimensional import MXFP8QuantizeRegisterBidimensionalKernel + + # 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 = MXFP8QuantizeKernel(cfg, SKIP_INPUT_MASKING=True) + self.general_skip_input_masking_and_scale_bounds_kernel = MXFP8QuantizeKernel( + cfg, SKIP_INPUT_MASKING=True, SKIP_SCALE_BOUNDS=True + ) + # We only need to mask when WITH_ACT is enabled because with activations applied zeros filled by TMA affect block statistics + self.general_kernel = MXFP8QuantizeKernel(cfg) if self.ACT_NEED_MASKING else None + self.g2r_rowwise_1lane = MXFP8QuantizeG2RRowwise1LaneKernel(cfg) if cfg.ROWWISE else None + self.g2s_bidim = ( + MXFP8QuantizeSpecializedBidimensionalKernel(cfg) + if cfg.ROWWISE and cfg.COLWISE + else None + ) + self.g2r_rowwise_2lane = MXFP8QuantizeG2RRowwise2LaneKernel(cfg) if cfg.ROWWISE else None + self.g2r_bidim = ( + MXFP8QuantizeRegisterBidimensionalKernel(cfg) if cfg.ROWWISE and cfg.COLWISE else None + ) + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + M = mX.shape[0] + N = mX.shape[1] + # Select specialized kernels if possible + plain_cast_only = ( + not self.cfg.WITH_AMAX + and not self.cfg.WITH_DBIAS + and not self.cfg.WITH_DACT + and not self.cfg.WITH_ACT + ) + dispatched_to_specialized = False + # The two-lane rowwise and g2r bidimensional kernels use BF16 packed math and 32-byte loads. + # Keep the existing kernels for FP16, fusions, ragged bidimensional tiles, + # misaligned input views and compilers predating 256-bit global loads. + if cutlass.const_expr( + plain_cast_only + and self.cfg.DTYPE is cutlass.BFloat16 + and cutlass.target_version(min_version="12.9") + ): + input_is_aligned = mX.iterator.toint() % 32 == 0 + if cutlass.const_expr(self.cfg.ROWWISE and not self.cfg.COLWISE): + scale_stride_matches = Boolean(True) + if cutlass.const_expr(not self.cfg.WITH_GEMM_SWIZZLED_SCALES): + scale_stride_matches = mS_row.shape[1] == N // 32 + # Column tiles use CUDA's grid.y axis. Inputs this wide always use two tiles. + grid_cols_fit = ( + N + <= 65535 + * self.g2r_rowwise_2lane._CTA_THREADS_X + * self.g2r_rowwise_2lane._ELEMENTS_PER_THREAD + * 2 + ) + if N % 128 == 0 and input_is_aligned and scale_stride_matches and grid_cols_fit: + dispatched_to_specialized = True + self.g2r_rowwise_2lane( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + stream, + ) + if cutlass.const_expr( + self.cfg.ROWWISE and self.cfg.COLWISE and not self.cfg.WITH_GEMM_SWIZZLED_SCALES + ): + columnwise_scales_are_aligned = mS_col.iterator.toint() % 16 == 0 + if ( + M % 32 == 0 + and N % 256 == 0 + and M // 32 <= 65535 + and input_is_aligned + and columnwise_scales_are_aligned + ): + dispatched_to_specialized = True + self.g2r_bidim( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + stream, + ) + # The one-lane rowwise and g2s bidimensional kernels support packed16 types (bf16/fp16). + if cutlass.const_expr(plain_cast_only and is_packed16(self.cfg.DTYPE)): + if cutlass.const_expr(self.cfg.ROWWISE and not self.cfg.COLWISE): + g2r_rowwise_1lane_grid_fits = ( + cute.ceil_div(M, self.g2r_rowwise_1lane._TILE_ROWS) <= 65535 + ) + # The one-lane rowwise kernel requires N divisible by 128 for vectorized stores. + if not dispatched_to_specialized and N % 128 == 0 and g2r_rowwise_1lane_grid_fits: + dispatched_to_specialized = True + self.g2r_rowwise_1lane( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + stream, + ) + # The g2s bidimensional kernel supports swizzled scales but we don't dispatch to it for now + if cutlass.const_expr( + self.cfg.ROWWISE and self.cfg.COLWISE and not self.cfg.WITH_GEMM_SWIZZLED_SCALES + ): + g2s_bidim_grid_fits = ( + cute.ceil_div( + M, + self.g2s_bidim._TILE_ROWS * self.g2s_bidim._NUM_TILES_Y, + ) + <= 65535 + ) + if not dispatched_to_specialized and g2s_bidim_grid_fits: + dispatched_to_specialized = True + self.g2s_bidim( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + stream, + ) + # If not using a specialized kernel, fall back to the general kernel + if not dispatched_to_specialized: + # 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 == 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): + """ + Return the compiled CuTeDSL function object for the given MXFP8 quantization config. + """ + + kernel_obj = MXFP8QuantizeEntry(cfg) + sym_M = cute.sym_int32() + sym_N = cute.sym_int32(divisibility=SYM_N_DIVISIBILITY) + in_shape = out_shape = (sym_M, sym_N) + # TE allocates scale tensors at a padded shape (see + # MXFP8Quantizer::get_scale_shape in transformer_engine/pytorch/csrc): + # rowwise: (roundup(M, 128), roundup(ceildiv(N, 32), 4)) + # columnwise: (roundup(ceildiv(M, 32), 4), roundup(N, 128)) + # These padded extents are NOT M/N (and SymInt has no `//`/`+`), so give the + # scales their own fresh syms carrying the divisibility the padding + # guarantees (rowwise: 128 x 4; colwise: 4 x 128). + scale_rowwise_shape = (cute.sym_int32(divisibility=128), cute.sym_int32(divisibility=4)) + scale_colwise_shape = (cute.sym_int32(divisibility=4), cute.sym_int32(divisibility=128)) + ws_shape = (cute.sym_int32(), sym_N) # (blocks_Y, N); N ties to input N + # Native FP8/E8M0 dtypes at the FFI boundary (matches the DLPack dtype the C++ + # bridge sends). + out_dtype = cfg.FP8_DTYPE + scale_dtype = cutlass.Float8E8M0FNU + + in_fake = cute.runtime.make_fake_compact_tensor( + cfg.DTYPE, in_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=16 + ) + out_row_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + out_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if cfg.ROWWISE + else None + ) + scale_row_fake = ( + cute.runtime.make_fake_compact_tensor( + scale_dtype, + scale_rowwise_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if cfg.ROWWISE + else None + ) + out_col_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + out_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if cfg.COLWISE + else None + ) + scale_col_fake = ( + cute.runtime.make_fake_compact_tensor( + scale_dtype, + scale_colwise_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if cfg.COLWISE + else None + ) + amax_fake = ( + cute.runtime.make_fake_compact_tensor( + Float32, (1,), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4 + ) + if cfg.WITH_AMAX + else None + ) + # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so that + # the same compiled kernel serves both an absent and a present flag (see noop_flag_is_set). + # The null address here is only a compile-time placeholder for the pointer the C++ dispatcher + # passes at launch. + noop_fake = cute.runtime.nullptr(Float32, mem_space=cute.AddressSpace.gmem, assumed_align=4) + act_input_fake = ( + cute.runtime.make_fake_compact_tensor( + cfg.DTYPE, + in_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if cfg.WITH_DACT + else None + ) + workspace_fake = ( + cute.runtime.make_fake_compact_tensor( + Float32, ws_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=4 + ) + if cfg.WITH_DBIAS + else None + ) + + compiled = cute.compile( + kernel_obj, + in_fake, # mX + out_row_fake, + scale_row_fake, # mO_row, mS_row + out_col_fake, + scale_col_fake, # mO_col, mS_col + amax_fake, # mAmax + noop_fake, # mNoop (pointer to the cast_noop flag, may be null at launch) + act_input_fake, # mDActInput (backward slot, unused) + workspace_fake, # mWorkspace(backward slot, unused) + cute.runtime.make_fake_stream(), # stream (compiled as an explicit tvm-ffi + # "handle" arg; C++ passes the CUDA stream + # as void*) + options="--enable-tvm-ffi", + ) + return compiled + + +def get_mxfp8_quantization_function( + fn_name: str, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + with_amax: bool, + with_dbias: bool, + with_dact: bool, + with_act: bool, + use_2d_quantization: bool, + activation: str, +) -> bool: + """Compile the MXFP8 quantize kernel for this config and register it in the TVM-FFI global registry + under EXACTLY `fn_name` (the key the C++ dispatcher built; Python treats it as an opaque name). + Returns True if a kernel is successfully registered under `fn_name` (the C++ side then fetches it with GetGlobal(fn_name)); + False if the config is unsupported, so the caller caches the negative result and falls back to the CUDA C++ kernel. + """ + try: + # Already registered (e.g. by a prior call) -> supported. + if tvm_ffi.get_global_func(fn_name, allow_missing=True) is not None: + return True + + major, minor = device_compute_capability() + if major < 10: + logger.warning( + "CuTeDSL MXFP8 backend requires compute capability >= 10.0 (Blackwell), " + "but detected %d.%d; falling back to the CUDA C++ kernel.", + major, + minor, + ) + return False + + try: + cfg = MXFP8QuantizeConfig( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + with_amax=with_amax, + with_dbias=with_dbias, + with_dact=with_dact, + with_act=with_act, + use_2d_quantization=use_2d_quantization, + activation=activation, + ) + except ValueError as e: + # The exception message states exactly why the config is unsupported + # (unknown dtype/activation, dbias not implemented, ...). Surfacing it as a + # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. + logger.warning( + "CuTeDSL MXFP8 backend does not support this config, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False + + logger.debug("Compiling CuTeDSL MXFP8 quantization kernel for %s", cfg) + compiled = compile_cutedsl_function_from_cfg(cfg) + # The returned compiled object is not neccessarily the compiled function itself. It could be a Python + # wrapper that parses arguments and calls the underlying function. This is needed in the general case + # as TVM-FFI is positional-only, however, for us this is unnecessary launch overhead as we pass the + # arguments in the correct order on the C++ side. + # Relevant issue: https://github.com/NVIDIA/cutlass/issues/3527 + # https://github.com/NVIDIA/cutlass/pull/3589 recommends using an explicit positional only marker (/) + # in the function signature, however it is not merged and it is something easy to miss and silently + # introduce unnecessarry overhead. + native = getattr(compiled, "__tvm_ffi_object__", lambda: None)() + tvm_ffi.register_global_func( + fn_name, native if native is not None else compiled, override=True + ) + return True + except Exception as e: # pylint: disable=broad-exception-caught + logger.error( + "CuTeDSL MXFP8 kernel compilation & registration failed, falling back to the CUDA" + " C++ kernel: %s", + e, + ) + # Unconditionally fallback to CUDA path because we can't tell if this exception is transient or permanent. + return False diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py new file mode 100644 index 00000000000..0713d68f0fd --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -0,0 +1,490 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Register-resident CuTeDSL bidimensional MXFP8 cast, ported from cast_bidim.cuh.""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import Int32, Int64, Uint32, Uint8 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + abs_max_x2_bf16, + create_l2_policy, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 +from .quantize_mxfp8_common import ( + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + _SmemAllocator, + bf16_pair_magnitude, + noop_flag_is_set, +) + + +@cute.jit +def bf16_mx_scale_reciprocal(magnitude: Uint32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: + """Exact BF16 reciprocal scale bits, including E8M0 zero, Inf and NaN cases.""" + rounded = (magnitude + Uint32(31)) & Uint32(0x7F80) + if rounded < Uint32(MAX_EXPONENT << 7): # pylint: disable=consider-using-max-builtin + rounded = Uint32(MAX_EXPONENT << 7) + reciprocal = Uint32((254 + MAX_EXPONENT) << 7) - rounded + if magnitude >= Uint32(0x7F80): + reciprocal = Uint32(0x0040) + if magnitude != Uint32(0x7F80): + reciprocal = Uint32(0x7FFF) + return reciprocal + + +@cute.jit +def bf16_mx_scale_reciprocal_pair(pair: Int32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: + """Compute two reciprocal scales without borrowing across packed BF16 halves.""" + magnitudes = Uint32(pair) & Uint32(0x7FFF7FFF) + rounded = (magnitudes + Uint32(0x001F001F)) & Uint32(0xFF80FF80) + floor = (MAX_EXPONENT << 23) | (MAX_EXPONENT << 7) + # Packed BF16 max accepts registers, so materialize the constant inside PTX. + clamped = cute.arch.inline_ptx( + "{ .reg.b32 floor_pair; " + f"mov.b32 floor_pair, {floor}; " + "max.xorsign.abs.bf16x2 {$w0}, {$r0}, floor_pair; }", + write_only_types=[Int32], + read_only_args=[rounded.bitcast(Int32)], + ) + bias = ((254 + MAX_EXPONENT) << 23) | ((254 + MAX_EXPONENT) << 7) + reciprocal = Uint32(bias) - Uint32(clamped) + if ((magnitudes + Uint32(0x00800080)) & Uint32(0x80008000)) != Uint32(0): + lo = bf16_mx_scale_reciprocal(magnitudes & Uint32(0x7FFF), MAX_EXPONENT) + hi = bf16_mx_scale_reciprocal(magnitudes >> Uint32(16), MAX_EXPONENT) + reciprocal = lo | (hi << Uint32(16)) + return reciprocal + + +@cute.jit +def bf16_mx_scale_bytes_pair(reciprocal: Uint32) -> Uint32: + """Gather the E8M0 bytes of two packed BF16 reciprocals into the low half.""" + bytes_pair = (Uint32(0x01FE01FE) - ((reciprocal >> Uint32(7)) & Uint32(0x00FF00FF))) & Uint32( + 0x00FF00FF + ) + return (bytes_pair & Uint32(0xFF)) | ((bytes_pair >> Uint32(8)) & Uint32(0xFF00)) + + +@cute.jit +def quantize_bidimensional_mxfp8_g2r( + gX_tile: cute.Tensor, + gO_row_tile: cute.Tensor, + gS_row_tile: cute.Tensor, + gO_col_tile: cute.Tensor, + gS_col_tile: cute.Tensor, + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ROWS_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_PER_THREAD: cutlass.Constexpr[int], + cache_policy: Int64, +): + """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile in both directions. + + TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD = 32, so every columnwise + MXFP8 block is contained in the tile. Each warp owns ROWS_PER_THREAD + consecutive rows, and each lane owns ELEMENTS_PER_THREAD adjacent columns. + Load the tile once, then reuse its registers for both quantization passes. + """ + TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD + TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD + LANES_PER_MX_BLOCK = 32 // ELEMENTS_PER_THREAD + COLUMN_PAIRS = TILE_COLS // 2 + COLUMN_PAIRS_PER_THREAD = ELEMENTS_PER_THREAD // 2 + PAIRS_PER_HALF = COLUMN_PAIRS_PER_THREAD // 2 + MAX_EXPONENT = 8 if cfg.FP8_DTYPE is cutlass.Float8E4M3FN else 15 + SCALE_DTYPE = cutlass.Float8E8M0FNU + tidx, _, _ = cute.arch.thread_idx() + warp, lane = tidx // CTA_THREADS_X, tidx % CTA_THREADS_X + + _, tv_payload = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)), + val_layout=cute.make_layout( + (ROWS_PER_THREAD, ELEMENTS_PER_THREAD), stride=(ELEMENTS_PER_THREAD, 1) + ), + ) + # Reshape each thread's values into (columns, rows), with contiguous columns. + fragment_layout = cute.make_layout( + (ELEMENTS_PER_THREAD, ROWS_PER_THREAD), stride=(1, ELEMENTS_PER_THREAD) + ) + gX_thread = cute.composition(cute.composition(gX_tile, tv_payload)[tidx, None], fragment_layout) + gO_row_thread = cute.composition( + cute.composition(gO_row_tile, tv_payload)[tidx, None], fragment_layout + ) + gO_col_thread = cute.composition( + cute.composition(gO_col_tile, tv_payload)[tidx, None], fragment_layout + ) + # Each group of lanes shares one rowwise scale for each of its rows. + tv_row_scale = cute.make_layout( + ( + (LANES_PER_MX_BLOCK, (CTA_THREADS_X // LANES_PER_MX_BLOCK, CTA_THREADS_Y)), + ROWS_PER_THREAD, + ), + stride=((0, (TILE_ROWS, ROWS_PER_THREAD)), 1), + ) + gS_row_thread = cute.composition(gS_row_tile, tv_row_scale)[tidx, None] + # Warp 0 writes one scale per column, in each lane's original column order. + gS_col_thread = cute.local_tile(gS_col_tile, (1, ELEMENTS_PER_THREAD), (0, lane))[0, None] + + # Keep the complete input fragment; reuse one row's output registers. + rX = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, ROWS_PER_THREAD), cfg.DTYPE) + rO = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.FP8_DTYPE) + rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) + rS_col = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), SCALE_DTYPE) + load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cfg.DTYPE, + num_bits_per_copy=ELEMENTS_PER_THREAD * cfg.DTYPE.width, + invariant=True, + ) + store_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + cfg.FP8_DTYPE, + num_bits_per_copy=ELEMENTS_PER_THREAD * cfg.FP8_DTYPE.width, + ) + row_scale_atom = cute.make_copy_atom(cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8) + col_scale_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=ELEMENTS_PER_THREAD * 8 + ) + scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(cfg.FP8_DTYPE) + rO_u32 = cute.recast_tensor(rO, Uint32) + + for row_idx in cutlass.range_constexpr(ROWS_PER_THREAD): + cute.copy(load_atom, gX_thread[None, row_idx], rX[None, row_idx], cache_policy=cache_policy) + + @cute.jit + def extract_row_amax(values_i32: cute.Tensor) -> Uint32: + """Reduce a 32-column MXFP8 block across its two or four lane owners.""" + amax = values_i32[0] + for pair_idx in cutlass.range_constexpr(1, cute.size(values_i32)): + amax = abs_max_x2_bf16(amax, values_i32[pair_idx]) + for stage in cutlass.range_constexpr(LANES_PER_MX_BLOCK.bit_length() - 1): + amax = abs_max_x2_bf16(amax, cute.arch.shuffle_sync_bfly(amax, 1 << stage)) + return bf16_pair_magnitude(amax) + + # Rowwise pass: lane shuffles close each 32-column reduction. + for row_idx in cutlass.range_constexpr(ROWS_PER_THREAD): + rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + reciprocal = bf16_mx_scale_reciprocal(extract_row_amax(rX_row_i32), MAX_EXPONENT) + reciprocal_packed = reciprocal | (reciprocal << Uint32(16)) + for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): + rO_u32[pack_idx] = scale_and_pack( + rX_row_i32[2 * pack_idx], + reciprocal_packed, + rX_row_i32[2 * pack_idx + 1], + reciprocal_packed, + ) + cute.copy(store_atom, rO, gO_row_thread[None, row_idx], cache_policy=cache_policy) + if lane % LANES_PER_MX_BLOCK == 0: + rS_row[0] = Uint8(Uint32(254) - (reciprocal >> Uint32(7))).bitcast(SCALE_DTYPE) + cute.copy( + row_scale_atom, + rS_row, + cute.local_tile(gS_row_thread, (1,), (row_idx,)), + cache_policy=cache_policy, + ) + + # Columnwise pass: first reduce each thread's rows into one maximum per column. + rColumnAmax = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.DTYPE) + rColumnScale = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.DTYPE) + rColumnAmax_i32 = cute.recast_tensor(rColumnAmax, Int32) + rColumnScale_u32 = cute.recast_tensor(rColumnScale, Uint32) + rX_i32 = cute.recast_tensor(rX, Int32) + for pair_idx in cutlass.range_constexpr(COLUMN_PAIRS_PER_THREAD): + amax = rX_i32[pair_idx, 0] + for row_idx in cutlass.range_constexpr(1, ROWS_PER_THREAD): + amax = abs_max_x2_bf16(amax, rX_i32[pair_idx, row_idx]) + rColumnAmax_i32[pair_idx] = amax + + allocator = _SmemAllocator() + sColumnAmax = allocator.allocate_tensor( + cfg.DTYPE, + cute.make_layout((CTA_THREADS_Y, TILE_COLS), stride=(TILE_COLS, 1)), + byte_alignment=16, + ) + sColumnScale = allocator.allocate_tensor( + cfg.DTYPE, cute.make_layout((TILE_COLS,)), byte_alignment=16 + ) + sColumnAmax_i32 = cute.recast_tensor(sColumnAmax, Int32) + sColumnScale_u32 = cute.recast_tensor(sColumnScale, Uint32) + # Group each lane's column pairs into two halves. Placing the first halves + # of all lanes before the second halves preserves CUDA's conflict-free + # vectorized shared-memory accesses, without calculating pointer offsets. + column_pair_layout = cute.make_layout( + ((PAIRS_PER_HALF, 2), CTA_THREADS_X), + stride=((1, COLUMN_PAIRS // 2), PAIRS_PER_HALF), + ) + sColumnAmax_thread = cute.composition(sColumnAmax_i32, (None, column_pair_layout))[ + warp, (None, lane) + ] + sColumnScale_thread = cute.composition(sColumnScale_u32, column_pair_layout)[None, lane] + for half_idx in cutlass.range_constexpr(2): + cute.autovec_copy( + cute.local_tile(rColumnAmax_i32, (PAIRS_PER_HALF,), (half_idx,)), + cute.local_tile(sColumnAmax_thread, (PAIRS_PER_HALF,), (half_idx,)), + ) + cute.arch.sync_threads() + + # One thread reduces each column pair across all warps (all 32 tile rows). + if tidx < COLUMN_PAIRS: + amax = sColumnAmax_i32[0, tidx] + for warp_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): + amax = abs_max_x2_bf16(amax, sColumnAmax_i32[warp_idx, tidx]) + sColumnScale_u32[tidx] = bf16_mx_scale_reciprocal_pair(amax, MAX_EXPONENT) + cute.arch.sync_threads() + + # The same layout returns the reciprocals to each lane's original columns. + for half_idx in cutlass.range_constexpr(2): + cute.autovec_copy( + cute.local_tile(sColumnScale_thread, (PAIRS_PER_HALF,), (half_idx,)), + cute.local_tile(rColumnScale_u32, (PAIRS_PER_HALF,), (half_idx,)), + ) + if warp == 0: + rS_col_u32 = cute.recast_tensor(rS_col, Uint32) + for pack_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): + rS_col_u32[pack_idx] = bf16_mx_scale_bytes_pair(rColumnScale_u32[2 * pack_idx]) | ( + bf16_mx_scale_bytes_pair(rColumnScale_u32[2 * pack_idx + 1]) << Uint32(16) + ) + cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) + + for row_idx in cutlass.range_constexpr(ROWS_PER_THREAD): + rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): + rO_u32[pack_idx] = scale_and_pack( + rX_row_i32[2 * pack_idx], + rColumnScale_u32[2 * pack_idx], + rX_row_i32[2 * pack_idx + 1], + rColumnScale_u32[2 * pack_idx + 1], + ) + cute.copy(store_atom, rO, gO_col_thread[None, row_idx], cache_policy=cache_policy) + + +class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): + """Load one CTA tile into registers and quantize it along both axes. + + TILE_ROWS = _CTA_THREADS_Y * _ROWS_PER_THREAD = 32. + TILE_COLS = _CTA_THREADS_X * ELEMENTS_PER_THREAD = 256 or 512. + The 32-row tile contains complete columnwise MXFP8 blocks; shared memory + combines the partial column maxima from the eight warps. Requires aligned + BF16 input, M divisible by 32, N divisible by 256 and plain scale layouts. + """ + + _CTA_THREADS_Y = 8 + _CTA_THREADS_X = 32 + _ROWS_PER_THREAD = 4 + _ELEMENTS_PER_THREAD = 16 + + def __init__(self, cfg: MXFP8QuantizeConfig): + self.cfg = cfg + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + M, N = mX.shape + TILE_ROWS = self._CTA_THREADS_Y * self._ROWS_PER_THREAD + WIDE_TILE_COLS = self._CTA_THREADS_X * self._ELEMENTS_PER_THREAD + if N % WIDE_TILE_COLS == 0: + if Int64(M // TILE_ROWS) * (N // WIDE_TILE_COLS) >= 4096: + self._launch( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mNoop, + stream, + CTA_THREADS_Y=self._CTA_THREADS_Y, + CTA_THREADS_X=self._CTA_THREADS_X, + ROWS_PER_THREAD=self._ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + CLUSTER_WIDTH=4, + ) + else: + self._launch( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mNoop, + stream, + CTA_THREADS_Y=self._CTA_THREADS_Y, + CTA_THREADS_X=self._CTA_THREADS_X, + ROWS_PER_THREAD=self._ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + CLUSTER_WIDTH=1, + ) + else: + self._launch( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mNoop, + stream, + CTA_THREADS_Y=self._CTA_THREADS_Y, + CTA_THREADS_X=self._CTA_THREADS_X, + ROWS_PER_THREAD=self._ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD // 2, + CLUSTER_WIDTH=8, + ) + + @cute.jit + def _launch( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: cute.Tensor, + mS_col: cute.Tensor, + mNoop: cute.Pointer, + stream: CUstream, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ROWS_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_PER_THREAD: cutlass.Constexpr[int], + CLUSTER_WIDTH: cutlass.Constexpr[int], + ): + assert CTA_THREADS_X == 32 + assert CTA_THREADS_Y * ROWS_PER_THREAD == 32 + assert ELEMENTS_PER_THREAD in (8, 16) + M, N = mX.shape + TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD + TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD + THREADS_PER_CTA = CTA_THREADS_Y * CTA_THREADS_X + grid_cols = N // TILE_COLS + grid_rows = M // TILE_ROWS + cluster_cols = Int32(CLUSTER_WIDTH) + while cluster_cols > 1 and grid_cols % cluster_cols != 0: + cluster_cols -= 1 + # Keep column CTAs on grid.x so clusters read adjacent column tiles. + grid = [grid_cols, grid_rows] + block = [THREADS_PER_CTA] + cluster = [cluster_cols, 1, 1] + # Preserve CUDA's compile-time strides for the two common column counts. + if N == 2048: + self.kernel( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mNoop, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ROWS_PER_THREAD=ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, + N_CONST=2048, + ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) + elif N == 4096: + self.kernel( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mNoop, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ROWS_PER_THREAD=ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, + N_CONST=4096, + ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) + else: + self.kernel( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mNoop, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ROWS_PER_THREAD=ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, + N_CONST=0, + ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: cute.Tensor, + mS_col: cute.Tensor, + mNoop: cute.Pointer, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ROWS_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_PER_THREAD: cutlass.Constexpr[int], + N_CONST: cutlass.Constexpr[int], + ): + """Check noop before constructing or accessing any CTA tile.""" + if not noop_flag_is_set(mNoop): + M, N = mX.shape + if cutlass.const_expr(N_CONST != 0): + N = N_CONST + payload_layout = cute.make_layout((M, N), stride=(N, 1)) + gX = cute.make_tensor(mX.iterator.align(32), payload_layout) + gO_row = cute.make_tensor(mO_row.iterator, payload_layout) + gO_col = cute.make_tensor(mO_col.iterator, payload_layout) + gS_row = cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), N // 32), stride=(N // 32, 1)), + ) + gS_col = cute.make_tensor( + mS_col.iterator.align(16), + cute.make_layout((cute.size(mS_col, mode=[0]), N), stride=(N, 1)), + ) + else: + # The dispatcher checks these alignments before choosing this kernel. + gX = cute.make_tensor(mX.iterator.align(32), mX.layout) + gO_row = mO_row + gO_col = mO_col + gS_row = mS_row + gS_col = cute.make_tensor(mS_col.iterator.align(16), mS_col.layout) + cta_col, cta_row, _ = cute.arch.block_idx() + cta_coord = (cta_row, cta_col) + TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD + TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD + TILER = (TILE_ROWS, TILE_COLS) + ROW_SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) + COL_SCALE_TILER = (TILE_ROWS // 32, TILE_COLS) + gX_tile = cute.local_tile(gX, TILER, cta_coord) + gO_row_tile = cute.local_tile(gO_row, TILER, cta_coord) + gO_col_tile = cute.local_tile(gO_col, TILER, cta_coord) + gS_row_tile = cute.local_tile(gS_row, ROW_SCALE_TILER, cta_coord) + gS_col_tile = cute.local_tile(gS_col, COL_SCALE_TILER, cta_coord) + cache_policy = create_l2_policy(True) + quantize_bidimensional_mxfp8_g2r( + gX_tile, + gO_row_tile, + gS_row_tile, + gO_col_tile, + gS_col_tile, + self.cfg, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ROWS_PER_THREAD=ROWS_PER_THREAD, + ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, + cache_policy=cache_policy, + ) 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 new file mode 100644 index 00000000000..bbdf9a4b510 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py @@ -0,0 +1,265 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""One-lane GMEM-to-RMEM rowwise MXFP8 cast with vectorized global I/O.""" + +# Local @cute.struct classes are shared-memory storage descriptors. +# pylint: disable=missing-class-docstring + +from typing import Optional, Type + +import cutlass +from cutlass import cute +from cutlass import Float8E8M0FNU, Int32, Uint32, Uint8 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + abs_max_x2_bf16, + abs_max_x2_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, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + cvt_f32_to_fp8e8m0fnu, + mul_f32x2_cvt_packed16x4_to_fp8x4, +) +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + MXFP8_BLOCK_SCALING_SIZE, + derive_swizzled_scale_layout, + noop_flag_is_set, +) + + +class MXFP8QuantizeG2RRowwise1LaneKernel(MXFP8QuantizeKernelBase): + """Specialized cast-only ROWWISE-only MXFP8 kernel. + + Requires N % 128 == 0 (full vectorizable column chunks). + + Plain rowwise-only quantize. Each thread owns one 32-element MXFP8 chunk and + uses vectorized global loads/stores (no TMA used).""" + + _TILE_ROWS = 4 + _TILE_COLS = 1024 + _THREADS_PER_CTA = 128 + + def __init__(self, cfg: MXFP8QuantizeConfig): + self.cfg = cfg + # If True, then this kernel will first write each thread's scale byte to a shared memory buffer, + # then utilize vectorized store to flush the buffer to global memory. + self._STASH_SCALE_TO_SMEM = True # Hardcode to true for now + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], # Unused, kept for API compatibility + mAmax: Optional[cute.Tensor], # Unused, kept for API compatibility + mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device + mDActInput: Optional[cute.Tensor], # Unused, kept for API compatibility + mWorkspace: Optional[cute.Tensor], # Unused, kept for API compatibility + stream: CUstream, + ): + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + f"[CuTeDSL] MXFP8QuantizeG2RRowwise1LaneKernel.__call__() with config: {self.cfg}\n" + ) + + M = mX.shape[0] + N = mX.shape[1] + + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row, _ = derive_swizzled_scale_layout(M, N, True, False, mS_row, None) + + grid = [ + cute.ceil_div(Int32(N), self._TILE_COLS), + cute.ceil_div(M, self._TILE_ROWS), + ] + block = [self._THREADS_PER_CTA] + + self.kernel( + mX, + mO_row, + mS_row, + mNoop, + mX.element_type, + ).launch(grid=grid, block=block, stream=stream) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mNoop: cute.Pointer, + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + ): + """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" + if not noop_flag_is_set(mNoop): + self._kernel_main(mX, mO_row, mS_row, DTYPE) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + ): + """Device entry for the specialized rowwise-only cast kernel (vectorized global loads/stores, no TMA).""" + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + M = mX.shape[0] + N = mX.shape[1] + + # Each thread handles one 32-element MXFP8 chunk (= one scale block). + # The 128 threads in the CTA are grouped as (4, 32), so they cover a + # (4, 1024) input tile and the matching (4, 32) scale tile. + CTA_Y = self._TILE_ROWS + CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE + tiler, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), + val_layout=cute.make_layout( + (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) + ), + ) + tiler_scale, tv_layout_scale = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), + val_layout=cute.make_layout((1, 1), stride=(1, 1)), + ) + + # Select the tile that belongs to this CTA, then the fragment per thread. + mX_tile = cute.local_tile(mX, tiler, (bidy, bidx)) + mO_tile = cute.local_tile(mO_row, tiler, (bidy, bidx)) + mS_tile = cute.local_tile(mS_row, tiler_scale, (bidy, bidx)) + mX_thread = cute.composition(mX_tile, tv_layout)[tidx, None] + mO_thread = cute.composition(mO_tile, tv_layout)[tidx, None] + mS_thread = cute.composition(mS_tile, tv_layout_scale)[tidx, None] + + rX_thread = cute.make_rmem_tensor( + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), + dtype=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 + mul_cvt4 = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, self.cfg.FP8_DTYPE) + rX_i32 = cute.make_tensor( + cute.recast_ptr(rX_thread.iterator, dtype=Int32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), + ) + rO_thread = cute.make_rmem_tensor( + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), + dtype=self.cfg.FP8_DTYPE, + ) + rO_u32 = cute.make_tensor( + cute.recast_ptr(rO_thread.iterator, dtype=Uint32), + cute.make_layout( + (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) + ), # Unit is Uint32, divide by 4 here + ) + + sS_thread = None + if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): + + @cute.struct + class SharedStorage: + buf: cute.struct.Align[cute.struct.MemRange[Float8E8M0FNU, CTA_Y * CTA_X], 16] + + storage = cutlass.utils.SmemAllocator().allocate(SharedStorage) + sScale = storage.buf.get_tensor(cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1))) + # sScale is (CTA_Y, CTA_X):(CTA_X, 1), which is the same layout as tv_layout_scale + # so sS_thread is really just a 1-element e8m0 buffer for this thread's scale. + sS_thread = cute.composition(sScale, tv_layout_scale)[tidx, None] + # Zero first so padding columns (cols past N/32 in the padded scale + # matrix) flush as 0 and we never read uninitialized smem. + # Raw 0x00: e8m0 has no zero value to convert from. + sS_thread[0] = Uint8(0).bitcast(Float8E8M0FNU) + cute.arch.sync_threads() + + row = bidy * self._TILE_ROWS + tidx // CTA_X + col = bidx * self._TILE_COLS + (tidx % CTA_X) * MXFP8_BLOCK_SCALING_SIZE + if row < M and col < N: + cute.autovec_copy(mX_thread, rX_thread) + amax_2x = rX_i32[0] + # Skip the first iteration because we assigned rX_i32[0] to amax_2x already + for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE // 2): + amax_2x = abs_max_x2(amax_2x, rX_i32[i]) + amax = cute.arch.fmax(fabs_f32(x2_lo_to_f32(amax_2x)), fabs_f32(x2_hi_to_f32(amax_2x))) + + biased_exp = cvt_f32_to_fp8e8m0fnu(amax * self.cfg.MAX_NORM_RCP) + if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): + sS_thread[0] = biased_exp + else: + mS_thread[0] = biased_exp + + # Rescale + FP8 cast, 4 elements per fused mul_cvt (one uint32 out), + # then a vectorized store. Mirrors CUDA's _use_cvt_4x path. + inv_scale = exp2f_rcp(biased_exp) + scale_2x = pack_f32x2(inv_scale, inv_scale) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): + rO_u32[i] = mul_cvt4(rX_i32[2 * i], rX_i32[2 * i + 1], scale_2x) + cute.autovec_copy(rO_thread, mO_thread) + + # Cooperative wide flush of the staged scales where padding columns flush as 0. + if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): + cute.arch.sync_threads() + # Use cute.size instead of .shape[1] because under the swizzled layout mode 1 is + # a nested tuple ((4, num_tiles_SC)), not a plain int. + padded_cols = cute.size(mS_row, mode=[1]) + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + # Swizzled rowwise scale layout is ((32, 4, num_tiles_M), (4, num_tiles_SC)):((16, 4, num_tiles_SC * 512), (1, 512)) + # so we have at most 4 elements continuous which is our vectorized store width + self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) + elif padded_cols % 16 == 0: + # If columns is divisible by 16, use 16 bytes as the vectorized store width + self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 16) + else: + # Otherwise use 4 bytes as the vectorized store width. + # Note our fake tensor requires 4 divisibility so this is enforced as long as you can get here + self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) + + @cute.jit + def _flush_scales_to_gmem( + self, + sScale: cute.Tensor, + mS_tile: cute.Tensor, + tidx: Int32, + bidx: Int32, + bidy: Int32, + M: Int32, + padded_cols: Int32, + width: cutlass.Constexpr[int], + ): + """Flush the staged (CTA_Y, CTA_X) scale tile to gmem with vectorized stores.""" + CTA_Y = self._TILE_ROWS + CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE + # Previously each threads has 1 byte, but now we are doing vectorized store, + # which means only a subset of threads will need to issue the store while other threads are not used. + active_threads = CTA_X // width + _, tv_flush = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, active_threads), stride=(active_threads, 1)), + val_layout=cute.make_layout((1, width), stride=(width, 1)), + ) + # We only need to use a subset of threads with shape (CTA_Y, active_threads) to write + # so if the thread is outside of this subset, it will remain inactive + if tidx < CTA_Y * active_threads: + # Absolute position of the scale vector to write in the GMEM buffer + thread_y = bidy * CTA_Y + tidx // active_threads + thread_x = bidx * CTA_X + (tidx % active_threads) * width + if thread_y < M and thread_x < padded_cols: + cute.autovec_copy( + cute.composition(sScale, tv_flush)[tidx, None], + cute.composition(mS_tile, tv_flush)[tidx, None], + ) 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 new file mode 100644 index 00000000000..95882a4161f --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -0,0 +1,368 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Two-lane GMEM-to-RMEM rowwise MXFP8 cast, ported from cast_rowwise.cuh.""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import Boolean, Float32, Int32, Int64, Uint32, Uint8 +from cuda.bindings.driver import CUstream + +from transformer_engine.common.CuTeDSL.utils import ( + abs_max_x2_bf16, + create_l2_policy, + exp2_bf16x2_rcp, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + cvt_f32_to_fp8e8m0fnu, + mul_bf16x2x2_cvt_bf16x4_to_fp8x4, +) +from .quantize_mxfp8_common import ( + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + bf16_pair_magnitude, + derive_swizzled_scale_layout, + noop_flag_is_set, +) + + +@cute.jit +def quantize_rowwise_mxfp8_g2r( + gX_tiles: cute.Tensor, + gO_tiles: cute.Tensor, + mS_tiles: cute.Tensor, + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + M: Int32, + N: Int32, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_PER_THREAD: cutlass.Constexpr[int], + NUM_TILES: cutlass.Constexpr[int], + CHECK_BOUNDS: cutlass.Constexpr[bool], + input_policy: Int64, + output_policy: Int64, +): + """Quantize a CTA's (TILE_ROWS, TILE_COLS, NUM_TILES) GMEM tiles. + + CTA_THREADS_X threads own each tile row. Each thread loads 16 BF16 values, and + two adjacent lanes share the maximum and scale for a 32-value MXFP8 block. + Load all tiles before reducing to preserve independent memory requests. + """ + TILE_ROWS = CTA_THREADS_Y + TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD + tidx, _, _ = cute.arch.thread_idx() + + _, tv_payload = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)), + val_layout=cute.make_layout((1, ELEMENTS_PER_THREAD), stride=(ELEMENTS_PER_THREAD, 1)), + ) + # Mode 0.0: two adjacent threads share one scale (stride 0). + # Mode 0.1: thread pairs are arranged as (scale columns, tile rows). + # Mode 1: each thread addresses one scale element. + tv_scale = cute.make_layout( + ((2, (CTA_THREADS_X // 2, CTA_THREADS_Y)), 1), stride=((0, (TILE_ROWS, 1)), 0) + ) + + # We load fragments from all GMEM tiles at once + rX = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, NUM_TILES), cfg.DTYPE) + # We store a fragment of one tile at a time + rO = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.FP8_DTYPE) + # RMEM flags to indicate whether each tile is out of bounds or not + # With CHECK_BOUNDS=False DCE will remove it so we don't pay extra cost when we don't use it + tile_valid_flags = cute.make_rmem_tensor((NUM_TILES,), Boolean) + + # 16 * 2 bytes = 32 bytes which is already the widest + LOAD_WIDTH = ELEMENTS_PER_THREAD * cfg.DTYPE.width + load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cfg.DTYPE, + num_bits_per_copy=LOAD_WIDTH, + invariant=True, # hint that data being read stays unchanged during the kernel + ) + store_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + cfg.FP8_DTYPE, + num_bits_per_copy=ELEMENTS_PER_THREAD * cfg.FP8_DTYPE.width, + ) + + scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(cfg.FP8_DTYPE) + + # Copy all tiles from GMEM to RMEM and potentially check their bounds + for tile_idx in cutlass.range_constexpr(NUM_TILES): + # gX_tiles is (TILE_ROWS, TILE_COLS, NUM_TILES) + gX_tile = gX_tiles[None, None, tile_idx] + gX_thread = cute.composition(gX_tile, tv_payload)[tidx, None] + valid = Boolean(True) + if cutlass.const_expr(CHECK_BOUNDS): + cta_row, cta_col, _ = cute.arch.block_idx() + row = cta_row * TILE_ROWS + tidx // CTA_THREADS_X + col = (cta_col * NUM_TILES + tile_idx) * TILE_COLS + ( + tidx % CTA_THREADS_X + ) * ELEMENTS_PER_THREAD + valid = row < M and col < N + tile_valid_flags[tile_idx] = valid + rX[None, tile_idx].fill(0) + if valid: + cute.copy(load_atom, gX_thread, rX[None, tile_idx], cache_policy=input_policy) + + @cute.jit + def extract_tile_amax(values_i32: cute.Tensor) -> Float32: + """Return the 32-value MXFP8 block amax shared by two adjacent threads. + + Balanced eight-pair tree, followed by one butterfly between lane partners. + """ + level = [ + abs_max_x2_bf16(values_i32[2 * i], values_i32[2 * i + 1]) + for i in range(cute.size(values_i32) // 2) + ] + amax_pair = abs_max_x2_bf16( + abs_max_x2_bf16(level[0], level[2]), abs_max_x2_bf16(level[1], level[3]) + ) + # Compare with the neighbor thread's amax (2 threads cooperate to process a MXFP8 block). + partner = cute.arch.shuffle_sync_bfly(amax_pair, 1) + amax_pair = abs_max_x2_bf16(amax_pair, partner) + return (bf16_pair_magnitude(amax_pair) << Uint32(16)).bitcast(Float32) + + for tile_idx in cutlass.range_constexpr(NUM_TILES): + gO_tile = gO_tiles[None, None, tile_idx] + mS_tile = mS_tiles[None, None, tile_idx] + gO_thread = cute.composition(gO_tile, tv_payload)[tidx, None] + mS_thread = cute.composition(mS_tile, tv_scale)[tidx, None] + # Reinterpret the native tensors only for the packed BF16/FP8 instructions. + rX_tile_i32 = cute.recast_tensor(rX[None, tile_idx], Int32) + rO_u32 = cute.recast_tensor(rO, Uint32) + amax = extract_tile_amax(rX_tile_i32) + exponent = cvt_f32_to_fp8e8m0fnu(amax * cfg.MAX_NORM_RCP) + + # Collect 4 scale bytes and write them using one thread + byte = Uint32(exponent.bitcast(Uint8)) + c1 = cute.arch.shuffle_sync_down(byte, 2) + c2 = cute.arch.shuffle_sync_down(byte, 4) + c3 = cute.arch.shuffle_sync_down(byte, 6) + if tidx % 8 == 0 and tile_valid_flags[tile_idx]: + scale_word = cute.make_tensor( + cute.recast_ptr(mS_thread.iterator, dtype=Uint32), cute.make_layout(1) + ) + scale_word[0] = byte | (c1 << Uint32(8)) | (c2 << Uint32(16)) | (c3 << Uint32(24)) + + # Compute and write the quantized MXFP8 values to GMEM + reciprocal_packed = exp2_bf16x2_rcp(exponent) + for i in cutlass.range_constexpr(cute.size(rO_u32)): + rO_u32[i] = scale_and_pack( + rX_tile_i32[2 * i], reciprocal_packed, rX_tile_i32[2 * i + 1], reciprocal_packed + ) + if tile_valid_flags[tile_idx]: + cute.copy(store_atom, rO, gO_thread, cache_policy=output_policy) + + +class MXFP8QuantizeG2RRowwise2LaneKernel(MXFP8QuantizeKernelBase): + """Two lanes per MXFP8 block, with a CTA sweeping NUM_TILES horizontally. + + One CTA tile has TILE_ROWS = CTA_THREADS_Y and + TILE_COLS = CTA_THREADS_X * _ELEMENTS_PER_THREAD. The CTA + processes NUM_TILES = 1 or 2 in the original input matrix. Both scale + formats use the same payload tiling; a composed layout maps swizzled scales. + _L2_CACHED_CTA_PERCENT tunes the share of CTAs whose input loads use normal + L2 retention: 0 streams all input and 100 caches all input normally. + """ + + _CTA_THREADS_X = 32 + _ELEMENTS_PER_THREAD = 16 + _L2_CACHED_CTA_PERCENT = 40 + + def __init__(self, cfg: MXFP8QuantizeConfig): + self.cfg = cfg + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + output_bytes = Int64(mX.shape[0]) * mX.shape[1] + if output_bytes <= 12 << 20: + self._launch( + mX, + mO_row, + mS_row, + mNoop, + stream, + CTA_THREADS_Y=8, + CTA_THREADS_X=self._CTA_THREADS_X, + NUM_TILES=1, + ) + elif output_bytes <= 48 << 20: + self._launch( + mX, + mO_row, + mS_row, + mNoop, + stream, + CTA_THREADS_Y=8, + CTA_THREADS_X=self._CTA_THREADS_X, + NUM_TILES=2, + ) + elif output_bytes <= 96 << 20: + self._launch( + mX, + mO_row, + mS_row, + mNoop, + stream, + CTA_THREADS_Y=4, + CTA_THREADS_X=self._CTA_THREADS_X, + NUM_TILES=2, + ) + else: + self._launch( + mX, + mO_row, + mS_row, + mNoop, + stream, + CTA_THREADS_Y=8, + CTA_THREADS_X=self._CTA_THREADS_X, + NUM_TILES=2, + ) + + @cute.jit + def _launch( + self, + mX: cute.Tensor, + mO: cute.Tensor, + mS: cute.Tensor, + mNoop: cute.Pointer, + stream: CUstream, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + NUM_TILES: cutlass.Constexpr[int], + ): + assert 0 <= self._L2_CACHED_CTA_PERCENT <= 100 + assert self._ELEMENTS_PER_THREAD == 16 + assert CTA_THREADS_X % 8 == 0 + assert CTA_THREADS_Y * CTA_THREADS_X % 32 == 0 + M, N = mX.shape + # Each CTA processes NUM_TILES horizontal tiles of (TILE_ROWS, TILE_COLS) values. + TILE_ROWS = CTA_THREADS_Y + TILE_COLS = CTA_THREADS_X * self._ELEMENTS_PER_THREAD + THREADS_PER_CTA = CTA_THREADS_Y * CTA_THREADS_X + # Total columns covered by a CTA; its row count stays TILE_ROWS. + CTA_ROWS = TILE_ROWS + CTA_COLS = TILE_COLS * NUM_TILES + + grid_rows = cute.ceil_div(M, CTA_ROWS) + grid_cols = cute.ceil_div(N, CTA_COLS) + check_bounds = M % CTA_ROWS != 0 or N % CTA_COLS != 0 + grid = Int64(grid_rows) * grid_cols + first_streaming_cta = grid * self._L2_CACHED_CTA_PERCENT // 100 + + if grid > 0: + if check_bounds: + self.kernel( + mX, + mO, + mS, + mNoop, + first_streaming_cta, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + NUM_TILES=NUM_TILES, + CHECK_BOUNDS=True, + ).launch(grid=[grid_rows, grid_cols], block=[THREADS_PER_CTA], stream=stream) + else: + self.kernel( + mX, + mO, + mS, + mNoop, + first_streaming_cta, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + NUM_TILES=NUM_TILES, + CHECK_BOUNDS=False, + ).launch(grid=[grid_rows, grid_cols], block=[THREADS_PER_CTA], stream=stream) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mO: cute.Tensor, + mS: cute.Tensor, + mNoop: cute.Pointer, + first_streaming_cta: Int64, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + NUM_TILES: cutlass.Constexpr[int], + CHECK_BOUNDS: cutlass.Constexpr[bool], + ): + """Check the noop flag before selecting or accessing any GMEM tile.""" + if not noop_flag_is_set(mNoop): + # Dispatch checks 32-byte input alignment before selecting this kernel. + gX = cute.make_tensor(mX.iterator.align(32), mX.layout) + # CUDA's x/y axes hold row/column tiles in the same order as local_tile. + cta_coord = cute.arch.block_idx()[:2] + M, N = mX.shape + TILE_ROWS = CTA_THREADS_Y + TILE_COLS = CTA_THREADS_X * self._ELEMENTS_PER_THREAD + CTA_ROWS = TILE_ROWS + CTA_COLS = TILE_COLS * NUM_TILES + + gS = mS + # Apply the swizzled scale layout if requested + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_swizzled, _ = derive_swizzled_scale_layout(M, N, True, False, mS, None) + gS = cute.composition(mS_swizzled, (cute.make_layout(M), cute.make_layout(N // 32))) + + # CTA_TILER is the block that a CTA processes in total, which consists of NUM_TILES horizontal tiles + CTA_TILER = (CTA_ROWS, CTA_COLS) + CTA_SCALE_TILER = (CTA_ROWS, CTA_COLS // 32) + gX_cta = cute.local_tile(gX, CTA_TILER, cta_coord) + gO_cta = cute.local_tile(mO, CTA_TILER, cta_coord) + gS_cta = cute.local_tile(gS, CTA_SCALE_TILER, cta_coord) + + fraction = Float32(1.0) + if cutlass.const_expr(self._L2_CACHED_CTA_PERCENT > 0): + fraction = Float32(0.0) + # A linear index is needed only for the L2 policy split. + cta_id = cute.crd2idx( + (Int64(cta_coord[0]), Int64(cta_coord[1])), cute.arch.grid_dim()[:2] + ) + if cta_id >= first_streaming_cta: + fraction = Float32(1.0) + input_policy = create_l2_policy(False, fraction) + output_policy = create_l2_policy(True) + + # Divide CTA_TILER to multiple TILER shape TILES + TILER = (TILE_ROWS, TILE_COLS) + SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) + # Mode 1's first index is 0 because CTA_ROWS == TILE_ROWS (there is only 1 tile on the row dimension). + gX_tiles = cute.local_tile(gX_cta, TILER, (0, None)) + gO_tiles = cute.local_tile(gO_cta, TILER, (0, None)) + mS_tiles = cute.local_tile(gS_cta, SCALE_TILER, (0, None)) + + quantize_rowwise_mxfp8_g2r( + gX_tiles, + gO_tiles, + mS_tiles, + self.cfg, + M, + N, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + NUM_TILES=NUM_TILES, + CHECK_BOUNDS=CHECK_BOUNDS, + input_policy=input_policy, + output_policy=output_policy, + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py new file mode 100644 index 00000000000..81a94382bc9 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py @@ -0,0 +1,712 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Specialized CuTeDSL bidimensional MXFP8 cast with staged TMA I/O.""" + +# Local @cute.struct classes are shared-memory storage descriptors. +# pylint: disable=missing-class-docstring + +from typing import Optional, Type + +import cutlass +from cutlass import cute +from cutlass import pipeline +from cutlass import Float32, Float8E8M0FNU, Int32, Int64, Uint32 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + is_packed16, + abs_max_x2_bf16, + abs_max_x2_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, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + cvt_f32_to_fp8e8m0fnu, + mul_f32x2_cvt_f32x4_to_fp8x4, + mul_f32x4_cvt_f32x4_to_fp8x4, +) +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + MXFP8_BLOCK_SCALING_SIZE, + THREADS_PER_WARP, + derive_swizzled_scale_layout, + noop_flag_is_set, +) + + +@cute.jit +def quantize_bidimensional_mxfp8_swizzled( + sX_tile: cute.Tensor, # 32x64 input tile (already sliced to this stage), Sw<3,4,3> swizzled + sO_row_tile: cute.Tensor, # 32x64 rowwise-output tile + sO_col_tile: cute.Tensor, # 32x64 colwise-output tile + sS_row_tile: cute.Tensor, # (32, 2) smem rowwise-scale staging tile (flushed in the epilogue) + sS_col_tile: cute.Tensor, # (1, 64) smem colwise-scale staging tile (flushed in the epilogue) + sColReduce_warp: cute.Tensor, # (32,) SMEM fp32 columnwise scale reduction buffer for this warp + WARPS_PER_CTA: cutlass.Constexpr[int], + MAX_NORM_RCP: cutlass.Constexpr[float], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], +): + """Quantize a pre-sliced 32x64 tile -> both rowwise and colwise MXFP8. Elements are + addressed through TV layouts: a warp = 32 lanes = 32 rows (lane == row), and each + thread owns one 32-col row segment (its 32-col block) and one output column. + + No bounds check: the caller only loops over valid tiles (`num_tiles`), and + M % 32 == 0 means a tile can only be partial along N. The partial tile's + out-of-bounds warp still runs, but harmlessly -- it reads TMA zero-filled smem, + its output columns are masked by the caller's TMA store, and its scale writes + (rowwise and colwise alike) go to staging slots whose flush targets are past-N + padding columns of the respective scale tensors. + """ + mul_cvt4 = mul_f32x2_cvt_f32x4_to_fp8x4(FP8_DTYPE) + mul_cvt4_elemwise = mul_f32x4_cvt_f32x4_to_fp8x4(FP8_DTYPE) + + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout(((MXFP8_BLOCK_SCALING_SIZE, 1), WARPS_PER_CTA)), # ((32, 1) 2) + val_layout=cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE)), + ) + _, tv_layout_rowwise_scale = cute.make_layout_tv( + thr_layout=cute.make_layout(((MXFP8_BLOCK_SCALING_SIZE, 1), WARPS_PER_CTA)), # ((32, 1) 2) + val_layout=cute.make_layout((1, 1)), + ) + _, tv_layout_colwise_scale = cute.make_layout_tv( + thr_layout=cute.make_layout((1, (MXFP8_BLOCK_SCALING_SIZE, WARPS_PER_CTA))), # (1, (32, 2)) + val_layout=cute.make_layout((1, 1)), + ) + + tidx, _, _ = cute.arch.thread_idx() + lane = tidx % MXFP8_BLOCK_SCALING_SIZE + + # Each composed [tidx, None] slice is 1-D (the value mode is flattened): the data + # slices are size 32 (this thread's row segment), the scale slices size 1. + tXsX = cute.composition(sX_tile, tv_layout)[tidx, None] # (32,) input row segment + tXsO_row = cute.composition(sO_row_tile, tv_layout)[tidx, None] # (32,) rowwise out + tXsO_col = cute.composition(sO_col_tile, tv_layout)[tidx, None] # (32,) colwise out + tSsS_row_tile = cute.composition(sS_row_tile, tv_layout_rowwise_scale)[tidx, None] # (1,) + tSsS_col_tile = cute.composition(sS_col_tile, tv_layout_colwise_scale)[tidx, None] # (1,) + + rO_row = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, FP8_DTYPE) + rO_col = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, FP8_DTYPE) + rO_row_u32 = cute.make_tensor( + cute.recast_ptr(rO_row.iterator, dtype=Uint32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,)), + ) + rO_col_u32 = cute.make_tensor( + cute.recast_ptr(rO_col.iterator, dtype=Uint32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,)), + ) + sColReduce = cute.make_tensor( + sColReduce_warp.iterator, + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4, 4), stride=(4, 1)), + ) + + if cutlass.const_expr(is_packed16(DTYPE)): + # If the input is bf16 / fp16, take this fast path and process 2 elements at a time in a packed i32 + 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 + rX = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, DTYPE) + # Do a vectorized load from SMEM to RMEM and unswizzle in the meantime. + cute.autovec_copy(tXsX, rX) + rX_2x = cute.make_tensor( + cute.recast_ptr(rX.iterator, dtype=Int32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), + ) + sColReduce_2x = cute.make_tensor( + cute.recast_ptr(sColReduce.iterator, dtype=Int64), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), + ) + + row_amax2 = rX_2x[0] + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): + pair = rX_2x[i] + # Skip the first iteration because we assigned rX_2x[0] to row_amax2 already + if cutlass.const_expr(i > 0): + row_amax2 = abs_max_x2(row_amax2, pair) + a_lo = fabs_f32(x2_lo_to_f32(pair)) + a_hi = fabs_f32(x2_hi_to_f32(pair)) + col_lo = cute.arch.warp_redux_sync(a_lo, kind="fmax") + col_hi = cute.arch.warp_redux_sync(a_hi, kind="fmax") + with cute.arch.elect_one(): + sColReduce_2x[i] = pack_f32x2(col_lo, col_hi) + + # Compute the rowwise scale factor + row_amax = cute.arch.fmax( + fabs_f32(x2_lo_to_f32(row_amax2)), fabs_f32(x2_hi_to_f32(row_amax2)) + ) + row_exp = cvt_f32_to_fp8e8m0fnu(row_amax * MAX_NORM_RCP) + row_inv = exp2f_rcp(row_exp) + tSsS_row_tile[0] = row_exp + cute.arch.sync_warp() + + # Compute the colwise scale factor (only handle the one that belongs to this thread / lane) + col_exp = cvt_f32_to_fp8e8m0fnu(sColReduce_warp[lane] * MAX_NORM_RCP) + tSsS_col_tile[0] = col_exp + sColReduce_warp[lane] = exp2f_rcp(col_exp) + cute.arch.sync_warp() + + row_scale_2x = pack_f32x2(row_inv, row_inv) + # Vectorized multiply-and-convert: 4 f32 → 4 fp8 + for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): + # Vectorized load for 4 columnwise scale factors from SMEM to RMEM + col_inv4 = cute.make_rmem_tensor(4, Float32) + cute.autovec_copy(sColReduce[j, None], col_inv4) # LDS.128, warp-broadcast + p01 = rX_2x[2 * j] + p23 = rX_2x[2 * j + 1] + f0 = x2_lo_to_f32(p01) + f1 = x2_hi_to_f32(p01) + f2 = x2_lo_to_f32(p23) + f3 = x2_hi_to_f32(p23) + # For rowwise quantized values, they use the same scale for all 4 elements, + # so we can just pass two row_inv to mul_cvt4 to apply it to all 4 elements at once. + rO_row_u32[j] = mul_cvt4(f0, f1, f2, f3, row_scale_2x) + # For columnwise quantized values, each element has its own scale; the + # elementwise variant fuses the per-element multiply into the cvt sequence. + rO_col_u32[j] = mul_cvt4_elemwise( + f0, f1, f2, f3, col_inv4[0], col_inv4[1], col_inv4[2], col_inv4[3] + ) + cute.autovec_copy(rO_row, tXsO_row) + cute.autovec_copy(rO_col, tXsO_col) + else: + # If input is fp32, take this slow path and process it normally without packing + rX = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, Float32) + cute.autovec_copy(tXsX, rX) + + row_amax = Float32(0.0) + for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + a = fabs_f32(rX[c]) + row_amax = cute.arch.fmax(row_amax, a) + col_amax = cute.arch.warp_redux_sync(a, kind="fmax") + with cute.arch.elect_one(): + sColReduce_warp[c] = col_amax + + # Compute the rowwise scale factor + row_exp = cvt_f32_to_fp8e8m0fnu(row_amax * MAX_NORM_RCP) + row_inv = exp2f_rcp(row_exp) + tSsS_row_tile[0] = row_exp # rowwise scale (this row-block, staged in smem) + cute.arch.sync_warp() + + # Compute the colwise scale factor (only handle the one that belongs to this thread / lane) + col_exp = cvt_f32_to_fp8e8m0fnu(sColReduce_warp[lane] * MAX_NORM_RCP) + tSsS_col_tile[0] = col_exp # colwise scale (this thread's column) + sColReduce_warp[lane] = exp2f_rcp(col_exp) + cute.arch.sync_warp() + + row_scale_2x = pack_f32x2(row_inv, row_inv) + # Vectorized multiply-and-convert: 4 f32 → 4 fp8 + for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): + # Vectorized load for 4 columnwise scale factors from SMEM to RMEM + col_inv4 = cute.make_rmem_tensor(4, Float32) + cute.autovec_copy(sColReduce[j, None], col_inv4) + offset = 4 * j + # For rowwise quantized values, they use the same scale for all 4 elements, + # so we can just pass two row_inv to mul_cvt4 to apply it to all 4 elements at once. + rO_row_u32[j] = mul_cvt4( + rX[offset], rX[offset + 1], rX[offset + 2], rX[offset + 3], row_scale_2x + ) + # For columnwise quantized values, each element has its own scale; the + # elementwise variant fuses the per-element multiply into the cvt sequence. + rO_col_u32[j] = mul_cvt4_elemwise( + rX[offset], + rX[offset + 1], + rX[offset + 2], + rX[offset + 3], + col_inv4[0], + col_inv4[1], + col_inv4[2], + col_inv4[3], + ) + cute.autovec_copy(rO_row, tXsO_row) + cute.autovec_copy(rO_col, tXsO_col) + + +class MXFP8QuantizeSpecializedBidimensionalKernel(MXFP8QuantizeKernelBase): + """Specialized cast-only rowwise+colwise MXFP8 kernel (swizzled TMA, one 32x32 tile per warp).""" + + _WARPS_PER_CTA = 2 + _THREADS_PER_CTA = _WARPS_PER_CTA * THREADS_PER_WARP + # A warp handles a 32x32 tile + _WARP_ROWS = MXFP8_BLOCK_SCALING_SIZE + _WARP_COLS = MXFP8_BLOCK_SCALING_SIZE + # A CTA handles a tile consisted two 32x32 warp subtiles side-by-side at a time, each handled by a warp + _TILE_ROWS = _WARP_ROWS + _TILE_COLS = _WARP_COLS * _WARPS_PER_CTA + _NUM_STAGES = 2 + _NUM_TILES_X = 4 + _NUM_TILES_Y = 1 + _NUM_TILES = _NUM_TILES_X * _NUM_TILES_Y + + # Rows and columns for the rowwise / columnwise scale factor tensor + _SCALE_ROWS = _TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE + _SCALE_COLS = _TILE_COLS // MXFP8_BLOCK_SCALING_SIZE + + def __init__(self, cfg: MXFP8QuantizeConfig): + self.cfg = cfg + if cfg.WITH_GEMM_SWIZZLED_SCALES: + self._NUM_TILES_X = 2 + self._NUM_TILES_Y = 2 + self._NUM_TILES = self._NUM_TILES_X * self._NUM_TILES_Y + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + "[CuTeDSL] MXFP8QuantizeSpecializedBidimensionalKernel.__call__() with config:" + f" {self.cfg}\n" + ) + M = mX.shape[0] + N = mX.shape[1] + + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row, mS_col = derive_swizzled_scale_layout( + M, N, self.cfg.ROWWISE, self.cfg.COLWISE, mS_row, mS_col + ) + + smem_tile_layout = cute.make_ordered_layout( + (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) + ) + cta_tiler = (self._TILE_ROWS, self._TILE_COLS) + # Apply 128B input Swizzle<3,4,3> to input tiles + in_smem_layout = cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, smem_tile_layout) + # Apply 64B output swizzle<2,4,3> for output tiles + out_smem_layout = cute.make_composed_layout(cute.make_swizzle(2, 4, 3), 0, smem_tile_layout) + op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + + # Input TMA atom + tma_atom, tma_src = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_load, + mX, + in_smem_layout, + cta_tiler, + num_multicast=1, + ) + # Rowwise output TMA atoms + tma_atom_out_row, tma_dst_out_row = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_store, + mO_row, + out_smem_layout, + cta_tiler, + num_multicast=1, + ) + # Colwise output TMA atoms + tma_atom_out_col, tma_dst_out_col = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_store, + mO_col, + out_smem_layout, + cta_tiler, + num_multicast=1, + ) + + grid = [ + # Each CTA covers a (TILE_ROWS * NUM_TILES_Y, TILE_COLS * NUM_TILES_X) area. + cute.ceil_div(Int32(N), self._TILE_COLS * self._NUM_TILES_X), + cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES_Y), + ] + block = [self._THREADS_PER_CTA] + self.kernel( + mX, + mS_row, + mS_col, + mNoop, + mX.element_type, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + ).launch(grid=grid, block=block, stream=stream) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mS_row: cute.Tensor, + mS_col: cute.Tensor, + mNoop: cute.Pointer, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom: cute.CopyAtom, + tma_src: cute.Tensor, + tma_atom_out_row: cute.CopyAtom, + tma_dst_out_row: cute.Tensor, + tma_atom_out_col: cute.CopyAtom, + tma_dst_out_col: cute.Tensor, + ): + """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" + if not noop_flag_is_set(mNoop): + self._kernel_main( + mX, + mS_row, + mS_col, + dtype, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + ) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mS_row: cute.Tensor, + mS_col: cute.Tensor, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom: cute.CopyAtom, + tma_src: cute.Tensor, + tma_atom_out_row: cute.CopyAtom, + tma_dst_out_row: cute.Tensor, + tma_atom_out_col: cute.CopyAtom, + tma_dst_out_col: cute.Tensor, + ): + """Device entry for the specialized bidimensional (rowwise+colwise) cast kernel.""" + + bidx, bidy, _ = cute.arch.block_idx() + + M = mX.shape[0] + N = mX.shape[1] + + # The CTA owns a NUM_TILES_Y x NUM_TILES_X arrangement of tiles, starting here. + tile_x_base = bidx * self._NUM_TILES_X + tile_y_base = bidy * self._NUM_TILES_Y + + num_tiles_x = cutlass.min( + self._NUM_TILES_X, + # Valid 64-col tiles remaining from this CTA's start column (bidx * span). + cute.ceil_div(Int32(N) - bidx * self._TILE_COLS * self._NUM_TILES_X, self._TILE_COLS), + ) + num_tiles_y = cutlass.min( + self._NUM_TILES_Y, + # Valid 32-row tiles remaining from this CTA's start row (bidy * span). + cute.ceil_div(Int32(M) - bidy * self._TILE_ROWS * self._NUM_TILES_Y, self._TILE_ROWS), + ) + num_tiles = num_tiles_x * num_tiles_y + + FP8_DTYPE = self.cfg.FP8_DTYPE + + @cute.struct + class SharedStorage: + mbar_storage: cute.struct.MemRange[ + cute.Int64, 2 * self._NUM_STAGES + ] # (full,empty) per stage + sX: cute.struct.Align[ + cute.struct.MemRange[dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES], + 128, + ] + sO_row: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sO_col: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + # Per-warp colwise-reduce scratchpad + sColReduce: cute.struct.Align[ + cute.struct.MemRange[Float32, THREADS_PER_WARP * self._WARPS_PER_CTA], 16 + ] + # Staged rowwise scales for the whole CTA + sScaleRow: cute.struct.Align[ + cute.struct.MemRange[ + Float8E8M0FNU, + self._TILE_ROWS * self._NUM_TILES * self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE, + ], + 16, + ] + # Staged colwise scales for the whole CTA + sScaleCol: cute.struct.Align[ + cute.struct.MemRange[Float8E8M0FNU, self._NUM_TILES * self._TILE_COLS], 16 + ] + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + tile_layout = cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + # SMEM input tile should have the same Swizzle<3,4,3> as the input TMA atom + sX = storage.sX.get_tensor(tile_layout, swizzle=cute.make_swizzle(3, 4, 3)) + # SMEM output tiles should have the same Swizzle<2,4,3> as the output TMA atom + sO_row = storage.sO_row.get_tensor(tile_layout, swizzle=cute.make_swizzle(2, 4, 3)) + sO_col = storage.sO_col.get_tensor(tile_layout, swizzle=cute.make_swizzle(2, 4, 3)) + # Reshape the per-warp colwise-reduce scratchpad to a (threads, warps) layout for easier indexing. + sColReduce = storage.sColReduce.get_tensor( + cute.make_layout((THREADS_PER_WARP, self._WARPS_PER_CTA), stride=(1, THREADS_PER_WARP)) + ) + # The whole rowwise scale factor SMEM tensor covered by a CTA + sScaleRow2D = storage.sScaleRow.get_tensor( + cute.make_ordered_layout( + ( + self._TILE_ROWS * self._NUM_TILES_Y, + self._SCALE_COLS * self._NUM_TILES_X, + ), + order=(1, 0), + ) + ) + # Divide by a single tile's shape, so it becomes ((TILE_ROWS, SCALE_COLS), TILES) for easier indexing + # where TILES = NUM_TILES_X * NUM_TILES_Y + sScaleRow = cute.zipped_divide(sScaleRow2D, (self._TILE_ROWS, self._SCALE_COLS)) + # The whole colwise scale factor SMEM tensor covered by a CTA + sScaleCol2D = storage.sScaleCol.get_tensor( + cute.make_ordered_layout( + ( + self._SCALE_ROWS * self._NUM_TILES_Y, + self._TILE_COLS * self._NUM_TILES_X, + ), + order=(1, 0), + ) + ) + # Divide by a single tile's shape, so it becomes ((SCALE_ROWS, TILE_COLS), TILES) for easier indexing + # where TILES = NUM_TILES_X * NUM_TILES_Y + sScaleCol = cute.zipped_divide(sScaleCol2D, (self._SCALE_ROWS, self._TILE_COLS)) + + # Zero the scale SMEM buffers so partial tiles / padding columns flush as 0. + tidx, _, _ = cute.arch.thread_idx() + # View rowwise staging buffers as flat uint32 and stride the CTA over them. + _ROW_SCALE_WORDS = self._TILE_ROWS * self._SCALE_COLS * self._NUM_TILES // 4 + sScaleRow_u32 = cute.make_tensor( + cute.recast_ptr(sScaleRow2D.iterator, dtype=Uint32), + cute.make_layout((_ROW_SCALE_WORDS,), stride=(1,)), + ) + for i in cutlass.range_constexpr(cute.ceil_div(_ROW_SCALE_WORDS, self._THREADS_PER_CTA)): + slot = i * self._THREADS_PER_CTA + tidx + if slot < _ROW_SCALE_WORDS: + sScaleRow_u32[slot] = Uint32(0) + # View colwise staging buffers as flat uint32 and stride the CTA over them. + _COL_SCALE_WORDS = self._SCALE_ROWS * self._TILE_COLS * self._NUM_TILES // 4 + sScaleCol_u32 = cute.make_tensor( + cute.recast_ptr(sScaleCol2D.iterator, dtype=Uint32), + cute.make_layout((_COL_SCALE_WORDS,), stride=(1,)), + ) + for i in cutlass.range_constexpr(cute.ceil_div(_COL_SCALE_WORDS, self._THREADS_PER_CTA)): + slot = i * self._THREADS_PER_CTA + tidx + if slot < _COL_SCALE_WORDS: + sScaleCol_u32[slot] = Uint32(0) + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom) + + # Only warp 0 is the producer (issues TMA) + producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + # Every warp is the consumer + consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._WARPS_PER_CTA) + # One TMA loads a tile + tx_count = self._TILE_ROWS * self._TILE_COLS * dtype.width // 8 + + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar_storage.data_ptr(), + num_stages=self._NUM_STAGES, + producer_group=producer_group, + consumer_group=consumer_group, + tx_count=tx_count, + cta_layout_vmnk=None, + ) + + prod_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._NUM_STAGES + ) + cons_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._NUM_STAGES + ) + + gX_tiled = cute.zipped_divide(tma_src, (self._TILE_ROWS, self._TILE_COLS)) + tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( + tma_atom, 0, cute.make_layout(1), sX, gX_tiled + ) + gO_row_tiled = cute.zipped_divide(tma_dst_out_row, (self._TILE_ROWS, self._TILE_COLS)) + tXsO_row, tXgO_row = cute.nvgpu.cpasync.tma_partition( + tma_atom_out_row, 0, cute.make_layout(1), sO_row, gO_row_tiled + ) + gO_col_tiled = cute.zipped_divide(tma_dst_out_col, (self._TILE_ROWS, self._TILE_COLS)) + tXsO_col, tXgO_col = cute.nvgpu.cpasync.tma_partition( + tma_atom_out_col, 0, cute.make_layout(1), sO_col, gO_col_tiled + ) + + cute.arch.sync_threads() + + # Prologue: warp 0 prefetches up to NUM_STAGES tiles to fully fill the pipeline + if warp_idx == 0: + for s in cutlass.range_constexpr(self._NUM_STAGES): + if s < num_tiles: + # This tile's position in the CTA's NUM_TILES_Y x NUM_TILES_X arrangement. + tile_y = s // num_tiles_x + tile_x = s % num_tiles_x + mainloop_pipeline.producer_acquire(prod_state) + cute.copy( + tma_atom, + tXgX[(None, (tile_y_base + tile_y, tile_x_base + tile_x))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + # Consumer: all warps fetch from the pipeline, process its tile, and issue a new load to the tile buffer it just consumed + for tile_idx in cutlass.range(num_tiles, unroll=1): + mainloop_pipeline.consumer_wait(cons_state) + # Only allow at most _NUM_STAGES-1 stages to be in-flight, because this iteration will reuse the ring buffer + # that is read _NUM_STAGES iterations ago so we must wait for whoever is reading that buffer to finish + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(self._NUM_STAGES - 1, read=True) + cute.arch.sync_threads() + # The current pipeline stage index, which is the tile index modulo the number of stages. + # This is used to index into the shared memory ring buffers that are wrapped around the number of stages. + stage_idx = cons_state.index + + # This tile's position in the CTA's NUM_TILES_Y x NUM_TILES_X arrangement. + tile_y = tile_idx // num_tiles_x + tile_x = tile_idx % num_tiles_x + + # Process the 32x64 SMEM tile for this stage + quantize_bidimensional_mxfp8_swizzled( + sX[None, stage_idx], + sO_row[None, stage_idx], + sO_col[None, stage_idx], + sScaleRow[(None, (tile_y, tile_x))], + sScaleCol[(None, (tile_y, tile_x))], + sColReduce[ + None, warp_idx + ], # Pick the per-warp colwise-reduce scratchpad for this warp + self._WARPS_PER_CTA, + self.cfg.MAX_NORM_RCP, + dtype, + self.cfg.FP8_DTYPE, + ) + + # Make the smem output writes visible to the TMA async proxy, then store. + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + + # We are done with this input pipeline SMEM buffer, signal the producer that it can write to this buffer + mainloop_pipeline.consumer_release(cons_state) + + # Issue TMA and write the output to GMEM + if warp_idx == 0: + tile_coord = (tile_y_base + tile_y, tile_x_base + tile_x) + cute.copy( + tma_atom_out_row, tXsO_row[(None, stage_idx)], tXgO_row[(None, tile_coord)] + ) + cute.copy( + tma_atom_out_col, tXsO_col[(None, stage_idx)], tXgO_col[(None, tile_coord)] + ) + cute.arch.cp_async_bulk_commit_group() + + cons_state.advance() + + # Producer: refill the buffer the consumer just freed with the tile + # NUM_STAGES ahead. + next_tile_idx = tile_idx + self._NUM_STAGES + if next_tile_idx < num_tiles: + if warp_idx == 0: + next_tile_y = next_tile_idx // num_tiles_x + next_tile_x = next_tile_idx % num_tiles_x + mainloop_pipeline.producer_acquire(prod_state) + cute.copy( + tma_atom, + tXgX[(None, (tile_y_base + next_tile_y, tile_x_base + next_tile_x))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + # TODO(kainingz): the rowwise WIDTH is hardcoded to 4. Consider using other width if we can convince + # cute.autovec_copy that wider stores are safe (e.g. 16B) and the runtime row pitch allows it. This might + # require us to change `scale_rowwise_shape`'s divisibility when we compile with fake tensors + self._flush_scales_to_gmem( + sScaleRow2D, + mS_row, + tidx, + bidx, + bidy, + ROWS=self._TILE_ROWS * self._NUM_TILES_Y, + COLS=self._SCALE_COLS * self._NUM_TILES_X, + WIDTH=4, + ) + self._flush_scales_to_gmem( + sScaleCol2D, + mS_col, + tidx, + bidx, + bidy, + ROWS=self._SCALE_ROWS * self._NUM_TILES_Y, + COLS=self._TILE_COLS * self._NUM_TILES_X, + WIDTH=1 if self.cfg.WITH_GEMM_SWIZZLED_SCALES else 16, + ) + + # Drain the final stores (their gmem writes must complete before the CTA exits). + cute.arch.cp_async_bulk_wait_group(0, read=False) + + @cute.jit + def _flush_scales_to_gmem( + self, + sScale2D: cute.Tensor, + mS: cute.Tensor, + tidx: Int32, + bidx: Int32, + bidy: Int32, + ROWS: cutlass.Constexpr[int], + COLS: cutlass.Constexpr[int], + WIDTH: cutlass.Constexpr[int], + ): + """Flush a staged (ROWS, COLS) SMEM scale block (a plain row-major 2D tensor) to its (bidy, bidx) slice of the gmem scale tensor, + where each GMEM slice has the same shape as this SMEM tile. Use `WIDTH` bytes per vectorized store. + """ + # Use cute.size instead of cute.shape because under the swizzled layout each mode is a nested tuple ((32, 4, num_tiles_M) etc.) + mS_M = cute.size(mS, mode=[0]) + mS_N = cute.size(mS, mode=[1]) + # Obtain the GMEM slice for the output scale factor block of this CTA + mS_tile = cute.local_tile(mS, (ROWS, COLS), (bidy, bidx)) + + ACTIVE_THREAD_COLS = COLS // WIDTH + _, tv_flush_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((ROWS, ACTIVE_THREAD_COLS), stride=(ACTIVE_THREAD_COLS, 1)), + val_layout=cute.make_layout((1, WIDTH), stride=(WIDTH, 1)), + ) + TOTAL_ACTIVE = ROWS * ACTIVE_THREAD_COLS + + # We may have more total active threads than threads in a CTA, so we do multiple waves of stores to flush the whole buffer + for wave in cutlass.range_constexpr(cute.ceil_div(TOTAL_ACTIVE, self._THREADS_PER_CTA)): + thread_idx = wave * self._THREADS_PER_CTA + tidx + if thread_idx < TOTAL_ACTIVE: + # Find the position for this slot's vectorized store in the GMEM scale factor buffer + thread_y = bidy * ROWS + thread_idx // ACTIVE_THREAD_COLS + thread_x = bidx * COLS + (thread_idx % ACTIVE_THREAD_COLS) * WIDTH + # For rowwise we have N divisible by 4 and WIDTH=4, and for colwise we have N divisible by 128 and WIDTH=16, + # so `thread_x < mS_N` with vectorized store is safe here. + # A thread only writes to a single row so `thread_y < mS_M` is also safe here. + if thread_y < mS_M and thread_x < mS_N: + cute.autovec_copy( + cute.composition(sScale2D, tv_flush_layout)[thread_idx, None], + cute.composition(mS_tile, tv_flush_layout)[thread_idx, None], + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py similarity index 50% rename from transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py rename to transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py index 824ed18aa59..26ff4c10b2d 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py @@ -7,6 +7,11 @@ Replicates the core logic of quantize_mxfp8.cuh: given a 2D tensor of BF16/FP16 values, quantize to MXFP8 format (FP8E4M3 data + E8M0 per-block scales). +BF16 cast-only calls with aligned input use the register-resident ports of +specialized/cast_rowwise.cuh and specialized/cast_bidim.cuh. Other configurations +use the one-lane rowwise, staged bidimensional or general kernels. All paths share the +TVM-FFI entrypoint. + Matches the C++ kernel's tile dimensions and thread layout: CHUNK_DIM_Y = 64, CHUNK_DIM_X = 64, THREADS_PER_CTA = 64 BUFF_DIM_Y = 32, BUFF_DIM_X = 64, STAGES = 2 @@ -20,21 +25,15 @@ # Local @cute.struct classes are SMEM-layout descriptors that need no docstrings. # pylint: disable=missing-class-docstring -import abc -import logging -import os from typing import Optional, Type import cutlass from cutlass import cute from cutlass import pipeline -from cutlass import Boolean, Float32, Float8E8M0FNU, Int32, Int64, Uint16, Uint32, Uint8 +from cutlass import Float32, Int32, Int64, Uint16, Uint32, Uint8 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module -import tvm_ffi from transformer_engine.common.CuTeDSL.utils import ( - device_compute_capability, - str_to_cutlass_dtype, is_packed16, abs_max_x2_bf16, abs_max_x2_f16, @@ -57,97 +56,24 @@ pack_f32x2, unpack_i64_to_i32x2, ) -from transformer_engine.common.CuTeDSL.activations import ( - act_relu, - act_gelu, - act_silu, - act_qgelu, - act_srelu, - dact_drelu, - dact_dsrelu, - dact_dsilu, - dact_dqgelu, - dact_dgelu, -) 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_f32x4_cvt_f32x4_to_fp8x4, mul_f32x2_cvt_packed16x4_to_fp8x4, ) - -CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" - -logger = logging.getLogger("transformer_engine.cutedsl.mxfp8") - -# Number of elements per MXFP8 scale block. They will share the same E8M0 scale factor -MXFP8_BLOCK_SCALING_SIZE = 32 -# How many threads are in one warp -THREADS_PER_WARP = 32 - -# FP8E4M3 max representable value -FP8E4M3_MAX_NORM = 448.0 -FP8E4M3_MAX_NORM_RCP = 1.0 / FP8E4M3_MAX_NORM -FP8E5M2_MAX_NORM = 57344.0 -FP8E5M2_MAX_NORM_RCP = 1.0 / FP8E5M2_MAX_NORM - -# If N is not divisible by 16, writing quantized FP8 output with N bytes per row will fail because -# TMA requires 16-byte alignment for each row -SYM_N_DIVISIBILITY = 16 - -SUPPORTED_ACTIVATIONS = { - "relu": act_relu, - "gelu": act_gelu, - "silu": act_silu, - "qgelu": act_qgelu, - "srelu": act_srelu, -} - -SUPPORTED_DACTIVATIONS = { - "drelu": dact_drelu, - "dgelu": dact_dgelu, - "dsilu": dact_dsilu, - "dqgelu": dact_dqgelu, - "dsrelu": dact_dsrelu, -} - - -@cute.jit -def derive_swizzled_scale_layout( - M: Int32, - N: Int32, - ROWWISE: cutlass.Constexpr[bool], - COLWISE: cutlass.Constexpr[bool], - mS_row: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], -): - """Derive the swizzled layout for the rowwise and colwise scale tensors.""" - num_scale_cols = cute.ceil_div(N, MXFP8_BLOCK_SCALING_SIZE) - num_scale_rows = cute.ceil_div(M, MXFP8_BLOCK_SCALING_SIZE) - - num_tiles_M = cute.ceil_div(M, 128) - num_tiles_SC = cute.ceil_div(num_scale_cols, 4) - num_tiles_SR = cute.ceil_div(num_scale_rows, 4) - num_tiles_N = cute.ceil_div(N, 128) - - if cutlass.const_expr(ROWWISE): - mS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout( - ((32, 4, num_tiles_M), (4, num_tiles_SC)), - stride=((16, 4, num_tiles_SC * 512), (1, 512)), - ), - ) - if cutlass.const_expr(COLWISE): - mS_col = cute.make_tensor( - mS_col.iterator, - cute.make_layout( - ((4, num_tiles_SR), (32, 4, num_tiles_N)), - stride=((1, 512), (16, 4, num_tiles_SR * 512)), - ), - ) - return mS_row, mS_col +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + MXFP8_BLOCK_SCALING_SIZE, + SUPPORTED_ACTIVATIONS, + SUPPORTED_DACTIVATIONS, + SYM_N_DIVISIBILITY, + THREADS_PER_WARP, + derive_swizzled_scale_layout, + noop_flag_is_set, +) @cute.jit @@ -609,320 +535,6 @@ def dbias_reduction_colwise( return dbias_partial -@cute.jit -def quantize_bidimensional_mxfp8_swizzled( - sX_tile: cute.Tensor, # 32x64 input tile (already sliced to this stage), Sw<3,4,3> swizzled - sO_row_tile: cute.Tensor, # 32x64 rowwise-output tile - sO_col_tile: cute.Tensor, # 32x64 colwise-output tile - sS_row_tile: cute.Tensor, # (32, 2) smem rowwise-scale staging tile (flushed in the epilogue) - sS_col_tile: cute.Tensor, # (1, 64) smem colwise-scale staging tile (flushed in the epilogue) - sColReduce_warp: cute.Tensor, # (32,) SMEM fp32 columnwise scale reduction buffer for this warp - WARPS_PER_CTA: cutlass.Constexpr[int], - MAX_NORM_RCP: cutlass.Constexpr[float], - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], -): - """Quantize a pre-sliced 32x64 tile -> both rowwise and colwise MXFP8. Elements are - addressed through TV layouts: a warp = 32 lanes = 32 rows (lane == row), and each - thread owns one 32-col row segment (its 32-col block) and one output column. - - No bounds check: the caller only loops over valid tiles (`num_tiles`), and - M % 32 == 0 means a tile can only be partial along N. The partial tile's - out-of-bounds warp still runs, but harmlessly -- it reads TMA zero-filled smem, - its output columns are masked by the caller's TMA store, and its scale writes - (rowwise and colwise alike) go to staging slots whose flush targets are past-N - padding columns of the respective scale tensors. - """ - mul_cvt4 = mul_f32x2_cvt_f32x4_to_fp8x4(FP8_DTYPE) - mul_cvt4_elemwise = mul_f32x4_cvt_f32x4_to_fp8x4(FP8_DTYPE) - - _, tv_layout = cute.make_layout_tv( - thr_layout=cute.make_layout(((MXFP8_BLOCK_SCALING_SIZE, 1), WARPS_PER_CTA)), # ((32, 1) 2) - val_layout=cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE)), - ) - _, tv_layout_rowwise_scale = cute.make_layout_tv( - thr_layout=cute.make_layout(((MXFP8_BLOCK_SCALING_SIZE, 1), WARPS_PER_CTA)), # ((32, 1) 2) - val_layout=cute.make_layout((1, 1)), - ) - _, tv_layout_colwise_scale = cute.make_layout_tv( - thr_layout=cute.make_layout((1, (MXFP8_BLOCK_SCALING_SIZE, WARPS_PER_CTA))), # (1, (32, 2)) - val_layout=cute.make_layout((1, 1)), - ) - - tidx, _, _ = cute.arch.thread_idx() - lane = tidx % MXFP8_BLOCK_SCALING_SIZE - - # Each composed [tidx, None] slice is 1-D (the value mode is flattened): the data - # slices are size 32 (this thread's row segment), the scale slices size 1. - tXsX = cute.composition(sX_tile, tv_layout)[tidx, None] # (32,) input row segment - tXsO_row = cute.composition(sO_row_tile, tv_layout)[tidx, None] # (32,) rowwise out - tXsO_col = cute.composition(sO_col_tile, tv_layout)[tidx, None] # (32,) colwise out - tSsS_row_tile = cute.composition(sS_row_tile, tv_layout_rowwise_scale)[tidx, None] # (1,) - tSsS_col_tile = cute.composition(sS_col_tile, tv_layout_colwise_scale)[tidx, None] # (1,) - - rO_row = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, FP8_DTYPE) - rO_col = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, FP8_DTYPE) - rO_row_u32 = cute.make_tensor( - cute.recast_ptr(rO_row.iterator, dtype=Uint32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,)), - ) - rO_col_u32 = cute.make_tensor( - cute.recast_ptr(rO_col.iterator, dtype=Uint32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,)), - ) - sColReduce = cute.make_tensor( - sColReduce_warp.iterator, - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4, 4), stride=(4, 1)), - ) - - if cutlass.const_expr(is_packed16(DTYPE)): - # If the input is bf16 / fp16, take this fast path and process 2 elements at a time in a packed i32 - 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 - rX = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, DTYPE) - # Do a vectorized load from SMEM to RMEM and unswizzle in the meantime. - cute.autovec_copy(tXsX, rX) - rX_2x = cute.make_tensor( - cute.recast_ptr(rX.iterator, dtype=Int32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), - ) - sColReduce_2x = cute.make_tensor( - cute.recast_ptr(sColReduce.iterator, dtype=Int64), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), - ) - - row_amax2 = rX_2x[0] - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): - pair = rX_2x[i] - # Skip the first iteration because we assigned rX_2x[0] to row_amax2 already - if cutlass.const_expr(i > 0): - row_amax2 = abs_max_x2(row_amax2, pair) - a_lo = fabs_f32(x2_lo_to_f32(pair)) - a_hi = fabs_f32(x2_hi_to_f32(pair)) - col_lo = cute.arch.warp_redux_sync(a_lo, kind="fmax") - col_hi = cute.arch.warp_redux_sync(a_hi, kind="fmax") - with cute.arch.elect_one(): - sColReduce_2x[i] = pack_f32x2(col_lo, col_hi) - - # Compute the rowwise scale factor - row_amax = cute.arch.fmax( - fabs_f32(x2_lo_to_f32(row_amax2)), fabs_f32(x2_hi_to_f32(row_amax2)) - ) - row_exp = cvt_f32_to_fp8e8m0fnu(row_amax * MAX_NORM_RCP) - row_inv = exp2f_rcp(row_exp) - tSsS_row_tile[0] = row_exp - cute.arch.sync_warp() - - # Compute the colwise scale factor (only handle the one that belongs to this thread / lane) - col_exp = cvt_f32_to_fp8e8m0fnu(sColReduce_warp[lane] * MAX_NORM_RCP) - tSsS_col_tile[0] = col_exp - sColReduce_warp[lane] = exp2f_rcp(col_exp) - cute.arch.sync_warp() - - row_scale_2x = pack_f32x2(row_inv, row_inv) - # Vectorized multiply-and-convert: 4 f32 → 4 fp8 - for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): - # Vectorized load for 4 columnwise scale factors from SMEM to RMEM - col_inv4 = cute.make_rmem_tensor(4, Float32) - cute.autovec_copy(sColReduce[j, None], col_inv4) # LDS.128, warp-broadcast - p01 = rX_2x[2 * j] - p23 = rX_2x[2 * j + 1] - f0 = x2_lo_to_f32(p01) - f1 = x2_hi_to_f32(p01) - f2 = x2_lo_to_f32(p23) - f3 = x2_hi_to_f32(p23) - # For rowwise quantized values, they use the same scale for all 4 elements, - # so we can just pass two row_inv to mul_cvt4 to apply it to all 4 elements at once. - rO_row_u32[j] = mul_cvt4(f0, f1, f2, f3, row_scale_2x) - # For columnwise quantized values, each element has its own scale; the - # elementwise variant fuses the per-element multiply into the cvt sequence. - rO_col_u32[j] = mul_cvt4_elemwise( - f0, f1, f2, f3, col_inv4[0], col_inv4[1], col_inv4[2], col_inv4[3] - ) - cute.autovec_copy(rO_row, tXsO_row) - cute.autovec_copy(rO_col, tXsO_col) - else: - # If input is fp32, take this slow path and process it normally without packing - rX = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, Float32) - cute.autovec_copy(tXsX, rX) - - row_amax = Float32(0.0) - for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - a = fabs_f32(rX[c]) - row_amax = cute.arch.fmax(row_amax, a) - col_amax = cute.arch.warp_redux_sync(a, kind="fmax") - with cute.arch.elect_one(): - sColReduce_warp[c] = col_amax - - # Compute the rowwise scale factor - row_exp = cvt_f32_to_fp8e8m0fnu(row_amax * MAX_NORM_RCP) - row_inv = exp2f_rcp(row_exp) - tSsS_row_tile[0] = row_exp # rowwise scale (this row-block, staged in smem) - cute.arch.sync_warp() - - # Compute the colwise scale factor (only handle the one that belongs to this thread / lane) - col_exp = cvt_f32_to_fp8e8m0fnu(sColReduce_warp[lane] * MAX_NORM_RCP) - tSsS_col_tile[0] = col_exp # colwise scale (this thread's column) - sColReduce_warp[lane] = exp2f_rcp(col_exp) - cute.arch.sync_warp() - - row_scale_2x = pack_f32x2(row_inv, row_inv) - # Vectorized multiply-and-convert: 4 f32 → 4 fp8 - for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): - # Vectorized load for 4 columnwise scale factors from SMEM to RMEM - col_inv4 = cute.make_rmem_tensor(4, Float32) - cute.autovec_copy(sColReduce[j, None], col_inv4) - offset = 4 * j - # For rowwise quantized values, they use the same scale for all 4 elements, - # so we can just pass two row_inv to mul_cvt4 to apply it to all 4 elements at once. - rO_row_u32[j] = mul_cvt4( - rX[offset], rX[offset + 1], rX[offset + 2], rX[offset + 3], row_scale_2x - ) - # For columnwise quantized values, each element has its own scale; the - # elementwise variant fuses the per-element multiply into the cvt sequence. - rO_col_u32[j] = mul_cvt4_elemwise( - rX[offset], - rX[offset + 1], - rX[offset + 2], - rX[offset + 3], - col_inv4[0], - col_inv4[1], - col_inv4[2], - col_inv4[3], - ) - cute.autovec_copy(rO_row, tXsO_row) - cute.autovec_copy(rO_col, tXsO_col) - - -@cute.jit -def noop_flag_is_set(mNoop: cute.Pointer) -> Boolean: - """Whether the cast_noop flag says this quantization is a no-op and must be skipped. - - mNoop is a pointer rather than a tensor so that one compiled kernel serves both a present and - an absent flag, hence the address is checked before it is dereferenced, exactly like the CUDA - C++ kernel's `noop != nullptr && noop[0] == 1.0f`. The two checks cannot be joined with `and`, - which the DSL lowers to a non-short-circuiting op that would load from the null pointer. - """ - flag_is_set = Boolean(False) - if mNoop.toint() != Int64(0): - flag_is_set = cute.make_tensor(mNoop, cute.make_layout((1,)))[0] == Float32(1.0) - return flag_is_set - - -class MXFP8QuantizeConfig: - """Configs for the compiled CuTeDSL kernel. These will be fixed once the kernel is compiled and - they will behave as const expressions. - """ - - def __init__( - self, - dtype: str, - fp8_dtype: str, - rowwise: bool, - colwise: bool, - with_gemm_swizzled_scales: bool, - with_amax: bool, - with_dbias: bool = False, - with_dact: bool = False, - with_act: bool = False, - 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) - self.DTYPE_STR = dtype # readable input-dtype token, for __str__ - if fp8_dtype not in ("Float8E4M3", "Float8E5M2"): - raise ValueError( - f"unknown FP8 dtype {fp8_dtype!r}; expected 'Float8E4M3' or 'Float8E5M2'" - ) - self.FP8_DTYPE = str_to_cutlass_dtype(fp8_dtype) - self.FP8_DTYPE_STR = fp8_dtype # readable token, for __str__ - self.ROWWISE = rowwise - self.COLWISE = colwise - if not (rowwise or colwise): - raise ValueError("at least one of rowwise or colwise must be true") - self.WITH_GEMM_SWIZZLED_SCALES = with_gemm_swizzled_scales - self.WITH_AMAX = with_amax - if not with_dact and not with_act: - if activation == "none": - self.ACTIVATION = None - else: - raise ValueError( - "activation must be none when with_dact and with_act are both False" - ) - else: - if with_dact and with_act: - raise ValueError( - "with_dact and with_act cannot be true at the same time since they are used for" - " different paths (bwd vs fwd)" - ) - if with_dact: - if activation in SUPPORTED_DACTIVATIONS: - self.ACTIVATION = activation - else: - raise ValueError( - f"unknown activation {activation!r} for with_dact=True; expected one of" - f" {sorted(SUPPORTED_DACTIVATIONS)}" - ) - elif with_act: - if activation in SUPPORTED_ACTIVATIONS: - self.ACTIVATION = activation - else: - raise ValueError( - f"unknown activation {activation!r} for with_act=True; expected one of" - f" {sorted(SUPPORTED_ACTIVATIONS)}" - ) - self.WITH_DACT = with_dact - self.WITH_ACT = with_act - # dbias is the column reduction of the (post-act/dact) element. With colwise - # output each thread owns a full column (trivial reduction); rowwise-only - # 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 - self.MAX_NORM_RCP = ( - FP8E4M3_MAX_NORM_RCP if fp8_dtype == "Float8E4M3" else FP8E5M2_MAX_NORM_RCP - ) - - def __str__(self): - return ( - f"MXFP8QuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " - 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})" - ) - - __repr__ = __str__ - - -class MXFP8QuantizeKernelBase(abc.ABC): - """Base class for MXFP8 quantize kernels.""" - - @abc.abstractmethod - def __call__( - self, - mX: cute.Tensor, - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, - mDActInput: Optional[cute.Tensor], - mWorkspace: Optional[cute.Tensor], - stream: CUstream, - ): - """ - Compiled kernel entrypoint (decorate with @cute.jit). - All MXFP8 quantize kernels must implement this interface because this is our C++ call site's contract in `quantize_mxfp8_cutedsl.cuh`. - C++ call site will pass the arguments in this exact order via tvm-ffi and the kernel must accept them even if they are not used. - """ - - class MXFP8QuantizeKernel(MXFP8QuantizeKernelBase): """The MXFP8 quantization kernel that mirrors the standard (non-specialized) MXFP8 CUDA C++ quantization kernel with multiple fusions (activation, dbias, etc.). @@ -1770,1077 +1382,3 @@ def _process_colwise( WITH_DBIAS=self.DBIAS_REDUCTION_IN_COLWISE, CACHE_ACTIVATION=self.CACHE_ACTIVATION, ) - - -class MXFP8QuantizeSpecializedRowwiseKernel(MXFP8QuantizeKernelBase): - """Specialized cast-only ROWWISE-only MXFP8 kernel. - - Requires N % 128 == 0 (full vectorizable column chunks). - - Plain rowwise-only quantize. Each thread owns one 32-element MXFP8 chunk and - uses vectorized global loads/stores (no TMA used).""" - - _TILE_ROWS = 4 - _TILE_COLS = 1024 - _THREADS_PER_CTA = 128 - - def __init__(self, cfg: MXFP8QuantizeConfig): - self.cfg = cfg - # If True, then this kernel will first write each thread's scale byte to a shared memory buffer, - # then utilize vectorized store to flush the buffer to global memory. - self._STASH_SCALE_TO_SMEM = True # Hardcode to true for now - - @cute.jit - def __call__( - self, - mX: cute.Tensor, - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], # Unused, kept for API compatibility - mAmax: Optional[cute.Tensor], # Unused, kept for API compatibility - mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device - mDActInput: Optional[cute.Tensor], # Unused, kept for API compatibility - mWorkspace: Optional[cute.Tensor], # Unused, kept for API compatibility - stream: CUstream, - ): - if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): - cute.printf( - "[CuTeDSL] MXFP8QuantizeSpecializedRowwiseKernel.__call__() with config:" - f" {self.cfg}\n" - ) - - M = mX.shape[0] - N = mX.shape[1] - - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - mS_row, _ = derive_swizzled_scale_layout(M, N, True, False, mS_row, None) - - grid = [ - cute.ceil_div(Int32(N), self._TILE_COLS), - cute.ceil_div(M, self._TILE_ROWS), - ] - block = [self._THREADS_PER_CTA] - - self.kernel( - mX, - mO_row, - mS_row, - mNoop, - mX.element_type, - ).launch(grid=grid, block=block, stream=stream) - - @cute.kernel - def kernel( - self, - mX: cute.Tensor, - mO_row: cute.Tensor, - mS_row: cute.Tensor, - mNoop: cute.Pointer, - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - ): - """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" - if not noop_flag_is_set(mNoop): - self._kernel_main(mX, mO_row, mS_row, DTYPE) - - @cute.jit - def _kernel_main( - self, - mX: cute.Tensor, - mO_row: cute.Tensor, - mS_row: cute.Tensor, - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - ): - """Device entry for the specialized rowwise-only cast kernel (vectorized global loads/stores, no TMA).""" - tidx, _, _ = cute.arch.thread_idx() - bidx, bidy, _ = cute.arch.block_idx() - M = mX.shape[0] - N = mX.shape[1] - - # Each thread handles one 32-element MXFP8 chunk (= one scale block). - # The 128 threads in the CTA are grouped as (4, 32), so they cover a - # (4, 1024) input tile and the matching (4, 32) scale tile. - CTA_Y = self._TILE_ROWS - CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE - tiler, tv_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), - val_layout=cute.make_layout( - (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) - ), - ) - tiler_scale, tv_layout_scale = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), - val_layout=cute.make_layout((1, 1), stride=(1, 1)), - ) - - # Select the tile that belongs to this CTA, then the fragment per thread. - mX_tile = cute.local_tile(mX, tiler, (bidy, bidx)) - mO_tile = cute.local_tile(mO_row, tiler, (bidy, bidx)) - mS_tile = cute.local_tile(mS_row, tiler_scale, (bidy, bidx)) - mX_thread = cute.composition(mX_tile, tv_layout)[tidx, None] - mO_thread = cute.composition(mO_tile, tv_layout)[tidx, None] - mS_thread = cute.composition(mS_tile, tv_layout_scale)[tidx, None] - - rX_thread = cute.make_rmem_tensor( - cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), - dtype=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 - mul_cvt4 = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, self.cfg.FP8_DTYPE) - rX_i32 = cute.make_tensor( - cute.recast_ptr(rX_thread.iterator, dtype=Int32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), - ) - rO_thread = cute.make_rmem_tensor( - cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), - dtype=self.cfg.FP8_DTYPE, - ) - rO_u32 = cute.make_tensor( - cute.recast_ptr(rO_thread.iterator, dtype=Uint32), - cute.make_layout( - (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) - ), # Unit is Uint32, divide by 4 here - ) - - sS_thread = None - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - - @cute.struct - class SharedStorage: - buf: cute.struct.Align[cute.struct.MemRange[Float8E8M0FNU, CTA_Y * CTA_X], 16] - - storage = cutlass.utils.SmemAllocator().allocate(SharedStorage) - sScale = storage.buf.get_tensor(cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1))) - # sScale is (CTA_Y, CTA_X):(CTA_X, 1), which is the same layout as tv_layout_scale - # so sS_thread is really just a 1-element e8m0 buffer for this thread's scale. - sS_thread = cute.composition(sScale, tv_layout_scale)[tidx, None] - # Zero first so padding columns (cols past N/32 in the padded scale - # matrix) flush as 0 and we never read uninitialized smem. - # Raw 0x00: e8m0 has no zero value to convert from. - sS_thread[0] = Uint8(0).bitcast(Float8E8M0FNU) - cute.arch.sync_threads() - - row = bidy * self._TILE_ROWS + tidx // CTA_X - col = bidx * self._TILE_COLS + (tidx % CTA_X) * MXFP8_BLOCK_SCALING_SIZE - if row < M and col < N: - cute.autovec_copy(mX_thread, rX_thread) - amax_2x = rX_i32[0] - # Skip the first iteration because we assigned rX_i32[0] to amax_2x already - for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE // 2): - amax_2x = abs_max_x2(amax_2x, rX_i32[i]) - amax = cute.arch.fmax(fabs_f32(x2_lo_to_f32(amax_2x)), fabs_f32(x2_hi_to_f32(amax_2x))) - - biased_exp = cvt_f32_to_fp8e8m0fnu(amax * self.cfg.MAX_NORM_RCP) - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - sS_thread[0] = biased_exp - else: - mS_thread[0] = biased_exp - - # Rescale + FP8 cast, 4 elements per fused mul_cvt (one uint32 out), - # then a vectorized store. Mirrors CUDA's _use_cvt_4x path. - inv_scale = exp2f_rcp(biased_exp) - scale_2x = pack_f32x2(inv_scale, inv_scale) - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): - rO_u32[i] = mul_cvt4(rX_i32[2 * i], rX_i32[2 * i + 1], scale_2x) - cute.autovec_copy(rO_thread, mO_thread) - - # Cooperative wide flush of the staged scales where padding columns flush as 0. - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - cute.arch.sync_threads() - # Use cute.size instead of .shape[1] because under the swizzled layout mode 1 is - # a nested tuple ((4, num_tiles_SC)), not a plain int. - padded_cols = cute.size(mS_row, mode=[1]) - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - # Swizzled rowwise scale layout is ((32, 4, num_tiles_M), (4, num_tiles_SC)):((16, 4, num_tiles_SC * 512), (1, 512)) - # so we have at most 4 elements continuous which is our vectorized store width - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) - elif padded_cols % 16 == 0: - # If columns is divisible by 16, use 16 bytes as the vectorized store width - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 16) - else: - # Otherwise use 4 bytes as the vectorized store width. - # Note our fake tensor requires 4 divisibility so this is enforced as long as you can get here - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) - - @cute.jit - def _flush_scales_to_gmem( - self, - sScale: cute.Tensor, - mS_tile: cute.Tensor, - tidx: Int32, - bidx: Int32, - bidy: Int32, - M: Int32, - padded_cols: Int32, - width: cutlass.Constexpr[int], - ): - """Flush the staged (CTA_Y, CTA_X) scale tile to gmem with vectorized stores.""" - CTA_Y = self._TILE_ROWS - CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE - # Previously each threads has 1 byte, but now we are doing vectorized store, - # which means only a subset of threads will need to issue the store while other threads are not used. - active_threads = CTA_X // width - _, tv_flush = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, active_threads), stride=(active_threads, 1)), - val_layout=cute.make_layout((1, width), stride=(width, 1)), - ) - # We only need to use a subset of threads with shape (CTA_Y, active_threads) to write - # so if the thread is outside of this subset, it will remain inactive - if tidx < CTA_Y * active_threads: - # Absolute position of the scale vector to write in the GMEM buffer - thread_y = bidy * CTA_Y + tidx // active_threads - thread_x = bidx * CTA_X + (tidx % active_threads) * width - if thread_y < M and thread_x < padded_cols: - cute.autovec_copy( - cute.composition(sScale, tv_flush)[tidx, None], - cute.composition(mS_tile, tv_flush)[tidx, None], - ) - - -class MXFP8QuantizeSpecializedBidimensionalKernel(MXFP8QuantizeKernelBase): - """Specialized cast-only rowwise+colwise MXFP8 kernel (swizzled TMA, one 32x32 tile per warp).""" - - _WARPS_PER_CTA = 2 - _THREADS_PER_CTA = _WARPS_PER_CTA * THREADS_PER_WARP - # A warp handles a 32x32 tile - _WARP_ROWS = MXFP8_BLOCK_SCALING_SIZE - _WARP_COLS = MXFP8_BLOCK_SCALING_SIZE - # A CTA handles a tile consisted two 32x32 warp subtiles side-by-side at a time, each handled by a warp - _TILE_ROWS = _WARP_ROWS - _TILE_COLS = _WARP_COLS * _WARPS_PER_CTA - _NUM_STAGES = 2 - _NUM_TILES_X = 4 - _NUM_TILES_Y = 1 - _NUM_TILES = _NUM_TILES_X * _NUM_TILES_Y - - # Rows and columns for the rowwise / columnwise scale factor tensor - _SCALE_ROWS = _TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE - _SCALE_COLS = _TILE_COLS // MXFP8_BLOCK_SCALING_SIZE - - def __init__(self, cfg: MXFP8QuantizeConfig): - self.cfg = cfg - if cfg.WITH_GEMM_SWIZZLED_SCALES: - self._NUM_TILES_X = 2 - self._NUM_TILES_Y = 2 - self._NUM_TILES = self._NUM_TILES_X * self._NUM_TILES_Y - - @cute.jit - def __call__( - self, - mX: cute.Tensor, - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device - mDActInput: Optional[cute.Tensor], - mWorkspace: Optional[cute.Tensor], - stream: CUstream, - ): - if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): - cute.printf( - "[CuTeDSL] MXFP8QuantizeSpecializedBidimensionalKernel.__call__() with config:" - f" {self.cfg}\n" - ) - M = mX.shape[0] - N = mX.shape[1] - - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - mS_row, mS_col = derive_swizzled_scale_layout( - M, N, self.cfg.ROWWISE, self.cfg.COLWISE, mS_row, mS_col - ) - - smem_tile_layout = cute.make_ordered_layout( - (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) - ) - cta_tiler = (self._TILE_ROWS, self._TILE_COLS) - # Apply 128B input Swizzle<3,4,3> to input tiles - in_smem_layout = cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, smem_tile_layout) - # Apply 64B output swizzle<2,4,3> for output tiles - out_smem_layout = cute.make_composed_layout(cute.make_swizzle(2, 4, 3), 0, smem_tile_layout) - op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() - op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() - - # Input TMA atom - tma_atom, tma_src = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_load, - mX, - in_smem_layout, - cta_tiler, - num_multicast=1, - ) - # Rowwise output TMA atoms - tma_atom_out_row, tma_dst_out_row = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_store, - mO_row, - out_smem_layout, - cta_tiler, - num_multicast=1, - ) - # Colwise output TMA atoms - tma_atom_out_col, tma_dst_out_col = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_store, - mO_col, - out_smem_layout, - cta_tiler, - num_multicast=1, - ) - - grid = [ - # Each CTA covers a (TILE_ROWS * NUM_TILES_Y, TILE_COLS * NUM_TILES_X) area. - cute.ceil_div(Int32(N), self._TILE_COLS * self._NUM_TILES_X), - cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES_Y), - ] - block = [self._THREADS_PER_CTA] - self.kernel( - mX, - mS_row, - mS_col, - mNoop, - mX.element_type, - tma_atom, - tma_src, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - ).launch(grid=grid, block=block, stream=stream) - - @cute.kernel - def kernel( - self, - mX: cute.Tensor, - mS_row: cute.Tensor, - mS_col: cute.Tensor, - mNoop: cute.Pointer, - dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - tma_atom: cute.CopyAtom, - tma_src: cute.Tensor, - tma_atom_out_row: cute.CopyAtom, - tma_dst_out_row: cute.Tensor, - tma_atom_out_col: cute.CopyAtom, - tma_dst_out_col: cute.Tensor, - ): - """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" - if not noop_flag_is_set(mNoop): - self._kernel_main( - mX, - mS_row, - mS_col, - dtype, - tma_atom, - tma_src, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - ) - - @cute.jit - def _kernel_main( - self, - mX: cute.Tensor, - mS_row: cute.Tensor, - mS_col: cute.Tensor, - dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - tma_atom: cute.CopyAtom, - tma_src: cute.Tensor, - tma_atom_out_row: cute.CopyAtom, - tma_dst_out_row: cute.Tensor, - tma_atom_out_col: cute.CopyAtom, - tma_dst_out_col: cute.Tensor, - ): - """Device entry for the specialized bidimensional (rowwise+colwise) cast kernel.""" - - bidx, bidy, _ = cute.arch.block_idx() - - M = mX.shape[0] - N = mX.shape[1] - - # The CTA owns a NUM_TILES_Y x NUM_TILES_X arrangement of tiles, starting here. - tile_x_base = bidx * self._NUM_TILES_X - tile_y_base = bidy * self._NUM_TILES_Y - - num_tiles_x = cutlass.min( - self._NUM_TILES_X, - # Valid 64-col tiles remaining from this CTA's start column (bidx * span). - cute.ceil_div(Int32(N) - bidx * self._TILE_COLS * self._NUM_TILES_X, self._TILE_COLS), - ) - num_tiles_y = cutlass.min( - self._NUM_TILES_Y, - # Valid 32-row tiles remaining from this CTA's start row (bidy * span). - cute.ceil_div(Int32(M) - bidy * self._TILE_ROWS * self._NUM_TILES_Y, self._TILE_ROWS), - ) - num_tiles = num_tiles_x * num_tiles_y - - FP8_DTYPE = self.cfg.FP8_DTYPE - - @cute.struct - class SharedStorage: - mbar_storage: cute.struct.MemRange[ - cute.Int64, 2 * self._NUM_STAGES - ] # (full,empty) per stage - sX: cute.struct.Align[ - cute.struct.MemRange[dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES], - 128, - ] - sO_row: cute.struct.Align[ - cute.struct.MemRange[ - FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - sO_col: cute.struct.Align[ - cute.struct.MemRange[ - FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - # Per-warp colwise-reduce scratchpad - sColReduce: cute.struct.Align[ - cute.struct.MemRange[Float32, THREADS_PER_WARP * self._WARPS_PER_CTA], 16 - ] - # Staged rowwise scales for the whole CTA - sScaleRow: cute.struct.Align[ - cute.struct.MemRange[ - Float8E8M0FNU, - self._TILE_ROWS * self._NUM_TILES * self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE, - ], - 16, - ] - # Staged colwise scales for the whole CTA - sScaleCol: cute.struct.Align[ - cute.struct.MemRange[Float8E8M0FNU, self._NUM_TILES * self._TILE_COLS], 16 - ] - - smem = cutlass.utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - - tile_layout = cute.make_layout( - ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), - stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), - ) - # SMEM input tile should have the same Swizzle<3,4,3> as the input TMA atom - sX = storage.sX.get_tensor(tile_layout, swizzle=cute.make_swizzle(3, 4, 3)) - # SMEM output tiles should have the same Swizzle<2,4,3> as the output TMA atom - sO_row = storage.sO_row.get_tensor(tile_layout, swizzle=cute.make_swizzle(2, 4, 3)) - sO_col = storage.sO_col.get_tensor(tile_layout, swizzle=cute.make_swizzle(2, 4, 3)) - # Reshape the per-warp colwise-reduce scratchpad to a (threads, warps) layout for easier indexing. - sColReduce = storage.sColReduce.get_tensor( - cute.make_layout((THREADS_PER_WARP, self._WARPS_PER_CTA), stride=(1, THREADS_PER_WARP)) - ) - # The whole rowwise scale factor SMEM tensor covered by a CTA - sScaleRow2D = storage.sScaleRow.get_tensor( - cute.make_ordered_layout( - ( - self._TILE_ROWS * self._NUM_TILES_Y, - self._SCALE_COLS * self._NUM_TILES_X, - ), - order=(1, 0), - ) - ) - # Divide by a single tile's shape, so it becomes ((TILE_ROWS, SCALE_COLS), TILES) for easier indexing - # where TILES = NUM_TILES_X * NUM_TILES_Y - sScaleRow = cute.zipped_divide(sScaleRow2D, (self._TILE_ROWS, self._SCALE_COLS)) - # The whole colwise scale factor SMEM tensor covered by a CTA - sScaleCol2D = storage.sScaleCol.get_tensor( - cute.make_ordered_layout( - ( - self._SCALE_ROWS * self._NUM_TILES_Y, - self._TILE_COLS * self._NUM_TILES_X, - ), - order=(1, 0), - ) - ) - # Divide by a single tile's shape, so it becomes ((SCALE_ROWS, TILE_COLS), TILES) for easier indexing - # where TILES = NUM_TILES_X * NUM_TILES_Y - sScaleCol = cute.zipped_divide(sScaleCol2D, (self._SCALE_ROWS, self._TILE_COLS)) - - # Zero the scale SMEM buffers so partial tiles / padding columns flush as 0. - tidx, _, _ = cute.arch.thread_idx() - # View rowwise staging buffers as flat uint32 and stride the CTA over them. - _ROW_SCALE_WORDS = self._TILE_ROWS * self._SCALE_COLS * self._NUM_TILES // 4 - sScaleRow_u32 = cute.make_tensor( - cute.recast_ptr(sScaleRow2D.iterator, dtype=Uint32), - cute.make_layout((_ROW_SCALE_WORDS,), stride=(1,)), - ) - for i in cutlass.range_constexpr(cute.ceil_div(_ROW_SCALE_WORDS, self._THREADS_PER_CTA)): - slot = i * self._THREADS_PER_CTA + tidx - if slot < _ROW_SCALE_WORDS: - sScaleRow_u32[slot] = Uint32(0) - # View colwise staging buffers as flat uint32 and stride the CTA over them. - _COL_SCALE_WORDS = self._SCALE_ROWS * self._TILE_COLS * self._NUM_TILES // 4 - sScaleCol_u32 = cute.make_tensor( - cute.recast_ptr(sScaleCol2D.iterator, dtype=Uint32), - cute.make_layout((_COL_SCALE_WORDS,), stride=(1,)), - ) - for i in cutlass.range_constexpr(cute.ceil_div(_COL_SCALE_WORDS, self._THREADS_PER_CTA)): - slot = i * self._THREADS_PER_CTA + tidx - if slot < _COL_SCALE_WORDS: - sScaleCol_u32[slot] = Uint32(0) - - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - if warp_idx == 0: - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom) - - # Only warp 0 is the producer (issues TMA) - producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) - # Every warp is the consumer - consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._WARPS_PER_CTA) - # One TMA loads a tile - tx_count = self._TILE_ROWS * self._TILE_COLS * dtype.width // 8 - - mainloop_pipeline = pipeline.PipelineTmaAsync.create( - barrier_storage=storage.mbar_storage.data_ptr(), - num_stages=self._NUM_STAGES, - producer_group=producer_group, - consumer_group=consumer_group, - tx_count=tx_count, - cta_layout_vmnk=None, - ) - - prod_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self._NUM_STAGES - ) - cons_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self._NUM_STAGES - ) - - gX_tiled = cute.zipped_divide(tma_src, (self._TILE_ROWS, self._TILE_COLS)) - tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( - tma_atom, 0, cute.make_layout(1), sX, gX_tiled - ) - gO_row_tiled = cute.zipped_divide(tma_dst_out_row, (self._TILE_ROWS, self._TILE_COLS)) - tXsO_row, tXgO_row = cute.nvgpu.cpasync.tma_partition( - tma_atom_out_row, 0, cute.make_layout(1), sO_row, gO_row_tiled - ) - gO_col_tiled = cute.zipped_divide(tma_dst_out_col, (self._TILE_ROWS, self._TILE_COLS)) - tXsO_col, tXgO_col = cute.nvgpu.cpasync.tma_partition( - tma_atom_out_col, 0, cute.make_layout(1), sO_col, gO_col_tiled - ) - - cute.arch.sync_threads() - - # Prologue: warp 0 prefetches up to NUM_STAGES tiles to fully fill the pipeline - if warp_idx == 0: - for s in cutlass.range_constexpr(self._NUM_STAGES): - if s < num_tiles: - # This tile's position in the CTA's NUM_TILES_Y x NUM_TILES_X arrangement. - tile_y = s // num_tiles_x - tile_x = s % num_tiles_x - mainloop_pipeline.producer_acquire(prod_state) - cute.copy( - tma_atom, - tXgX[(None, (tile_y_base + tile_y, tile_x_base + tile_x))], - tXsX[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - mainloop_pipeline.producer_commit(prod_state) - prod_state.advance() - - # Consumer: all warps fetch from the pipeline, process its tile, and issue a new load to the tile buffer it just consumed - for tile_idx in cutlass.range(num_tiles, unroll=1): - mainloop_pipeline.consumer_wait(cons_state) - # Only allow at most _NUM_STAGES-1 stages to be in-flight, because this iteration will reuse the ring buffer - # that is read _NUM_STAGES iterations ago so we must wait for whoever is reading that buffer to finish - if warp_idx == 0: - cute.arch.cp_async_bulk_wait_group(self._NUM_STAGES - 1, read=True) - cute.arch.sync_threads() - # The current pipeline stage index, which is the tile index modulo the number of stages. - # This is used to index into the shared memory ring buffers that are wrapped around the number of stages. - stage_idx = cons_state.index - - # This tile's position in the CTA's NUM_TILES_Y x NUM_TILES_X arrangement. - tile_y = tile_idx // num_tiles_x - tile_x = tile_idx % num_tiles_x - - # Process the 32x64 SMEM tile for this stage - quantize_bidimensional_mxfp8_swizzled( - sX[None, stage_idx], - sO_row[None, stage_idx], - sO_col[None, stage_idx], - sScaleRow[(None, (tile_y, tile_x))], - sScaleCol[(None, (tile_y, tile_x))], - sColReduce[ - None, warp_idx - ], # Pick the per-warp colwise-reduce scratchpad for this warp - self._WARPS_PER_CTA, - self.cfg.MAX_NORM_RCP, - dtype, - self.cfg.FP8_DTYPE, - ) - - # Make the smem output writes visible to the TMA async proxy, then store. - cute.arch.fence_proxy("async.shared", space="cta") - cute.arch.sync_threads() - - # We are done with this input pipeline SMEM buffer, signal the producer that it can write to this buffer - mainloop_pipeline.consumer_release(cons_state) - - # Issue TMA and write the output to GMEM - if warp_idx == 0: - tile_coord = (tile_y_base + tile_y, tile_x_base + tile_x) - cute.copy( - tma_atom_out_row, tXsO_row[(None, stage_idx)], tXgO_row[(None, tile_coord)] - ) - cute.copy( - tma_atom_out_col, tXsO_col[(None, stage_idx)], tXgO_col[(None, tile_coord)] - ) - cute.arch.cp_async_bulk_commit_group() - - cons_state.advance() - - # Producer: refill the buffer the consumer just freed with the tile - # NUM_STAGES ahead. - next_tile_idx = tile_idx + self._NUM_STAGES - if next_tile_idx < num_tiles: - if warp_idx == 0: - next_tile_y = next_tile_idx // num_tiles_x - next_tile_x = next_tile_idx % num_tiles_x - mainloop_pipeline.producer_acquire(prod_state) - cute.copy( - tma_atom, - tXgX[(None, (tile_y_base + next_tile_y, tile_x_base + next_tile_x))], - tXsX[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - mainloop_pipeline.producer_commit(prod_state) - prod_state.advance() - - # TODO(kainingz): the rowwise WIDTH is hardcoded to 4. Consider using other width if we can convince - # cute.autovec_copy that wider stores are safe (e.g. 16B) and the runtime row pitch allows it. This might - # require us to change `scale_rowwise_shape`'s divisibility when we compile with fake tensors - self._flush_scales_to_gmem( - sScaleRow2D, - mS_row, - tidx, - bidx, - bidy, - ROWS=self._TILE_ROWS * self._NUM_TILES_Y, - COLS=self._SCALE_COLS * self._NUM_TILES_X, - WIDTH=4, - ) - self._flush_scales_to_gmem( - sScaleCol2D, - mS_col, - tidx, - bidx, - bidy, - ROWS=self._SCALE_ROWS * self._NUM_TILES_Y, - COLS=self._TILE_COLS * self._NUM_TILES_X, - WIDTH=1 if self.cfg.WITH_GEMM_SWIZZLED_SCALES else 16, - ) - - # Drain the final stores (their gmem writes must complete before the CTA exits). - cute.arch.cp_async_bulk_wait_group(0, read=False) - - @cute.jit - def _flush_scales_to_gmem( - self, - sScale2D: cute.Tensor, - mS: cute.Tensor, - tidx: Int32, - bidx: Int32, - bidy: Int32, - ROWS: cutlass.Constexpr[int], - COLS: cutlass.Constexpr[int], - WIDTH: cutlass.Constexpr[int], - ): - """Flush a staged (ROWS, COLS) SMEM scale block (a plain row-major 2D tensor) to its (bidy, bidx) slice of the gmem scale tensor, - where each GMEM slice has the same shape as this SMEM tile. Use `WIDTH` bytes per vectorized store. - """ - # Use cute.size instead of cute.shape because under the swizzled layout each mode is a nested tuple ((32, 4, num_tiles_M) etc.) - mS_M = cute.size(mS, mode=[0]) - mS_N = cute.size(mS, mode=[1]) - # Obtain the GMEM slice for the output scale factor block of this CTA - mS_tile = cute.local_tile(mS, (ROWS, COLS), (bidy, bidx)) - - ACTIVE_THREAD_COLS = COLS // WIDTH - _, tv_flush_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((ROWS, ACTIVE_THREAD_COLS), stride=(ACTIVE_THREAD_COLS, 1)), - val_layout=cute.make_layout((1, WIDTH), stride=(WIDTH, 1)), - ) - TOTAL_ACTIVE = ROWS * ACTIVE_THREAD_COLS - - # We may have more total active threads than threads in a CTA, so we do multiple waves of stores to flush the whole buffer - for wave in cutlass.range_constexpr(cute.ceil_div(TOTAL_ACTIVE, self._THREADS_PER_CTA)): - thread_idx = wave * self._THREADS_PER_CTA + tidx - if thread_idx < TOTAL_ACTIVE: - # Find the position for this slot's vectorized store in the GMEM scale factor buffer - thread_y = bidy * ROWS + thread_idx // ACTIVE_THREAD_COLS - thread_x = bidx * COLS + (thread_idx % ACTIVE_THREAD_COLS) * WIDTH - # For rowwise we have N divisible by 4 and WIDTH=4, and for colwise we have N divisible by 128 and WIDTH=16, - # so `thread_x < mS_N` with vectorized store is safe here. - # A thread only writes to a single row so `thread_y < mS_M` is also safe here. - if thread_y < mS_M and thread_x < mS_N: - cute.autovec_copy( - cute.composition(sScale2D, tv_flush_layout)[thread_idx, None], - cute.composition(mS_tile, tv_flush_layout)[thread_idx, None], - ) - - -class MXFP8QuantizeEntry(MXFP8QuantizeKernelBase): - """Select the appropriate MXFP8 quantization kernel based on the configuration and runtime shapes.""" - - def __init__(self, cfg: MXFP8QuantizeConfig): - 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 = MXFP8QuantizeKernel(cfg, SKIP_INPUT_MASKING=True) - self.general_skip_input_masking_and_scale_bounds_kernel = MXFP8QuantizeKernel( - cfg, SKIP_INPUT_MASKING=True, SKIP_SCALE_BOUNDS=True - ) - # We only need to mask when WITH_ACT is enabled because with activations applied zeros filled by TMA affect block statistics - self.general_kernel = MXFP8QuantizeKernel(cfg) if self.ACT_NEED_MASKING else None - self.specialized_rowwise = ( - MXFP8QuantizeSpecializedRowwiseKernel(cfg) if cfg.ROWWISE else None - ) - self.specialized_bidim = ( - MXFP8QuantizeSpecializedBidimensionalKernel(cfg) - if cfg.ROWWISE and cfg.COLWISE - else None - ) - - @cute.jit - def __call__( - self, - mX: cute.Tensor, - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, - mDActInput: Optional[cute.Tensor], - mWorkspace: Optional[cute.Tensor], - stream: CUstream, - ): - M = mX.shape[0] - N = mX.shape[1] - # Select specialized kernels if possible - plain_cast_only = ( - not self.cfg.WITH_AMAX - and not self.cfg.WITH_DBIAS - and not self.cfg.WITH_DACT - and not self.cfg.WITH_ACT - ) - dispatched_to_specialized = False - # Only dispatch to the specialized kernels for packed16 types (bf16/fp16) - if cutlass.const_expr(plain_cast_only and is_packed16(self.cfg.DTYPE)): - if cutlass.const_expr(self.cfg.ROWWISE and not self.cfg.COLWISE): - rowwise_specialized_grid_fits = ( - cute.ceil_div(M, self.specialized_rowwise._TILE_ROWS) <= 65535 - ) - # The rowwise specialized kernel requires N divisible by 128 for vectorized stores - if N % 128 == 0 and rowwise_specialized_grid_fits: - dispatched_to_specialized = True - self.specialized_rowwise( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - # The bidimensional specialized kernel supports swizzled scales but we don't dispatch to it for now - if cutlass.const_expr( - self.cfg.ROWWISE and self.cfg.COLWISE and not self.cfg.WITH_GEMM_SWIZZLED_SCALES - ): - bidim_specialized_grid_fits = ( - cute.ceil_div( - M, - self.specialized_bidim._TILE_ROWS * self.specialized_bidim._NUM_TILES_Y, - ) - <= 65535 - ) - if bidim_specialized_grid_fits: - dispatched_to_specialized = True - self.specialized_bidim( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - # If not using a specialized kernel, fall back to the general kernel - if not dispatched_to_specialized: - # 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 == 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): - """ - Return the compiled CuTeDSL function object for the given MXFP8 quantization config. - """ - - kernel_obj = MXFP8QuantizeEntry(cfg) - sym_M = cute.sym_int32() - sym_N = cute.sym_int32(divisibility=SYM_N_DIVISIBILITY) - in_shape = out_shape = (sym_M, sym_N) - # TE allocates scale tensors at a padded shape (see - # MXFP8Quantizer::get_scale_shape in transformer_engine/pytorch/csrc): - # rowwise: (roundup(M, 128), roundup(ceildiv(N, 32), 4)) - # columnwise: (roundup(ceildiv(M, 32), 4), roundup(N, 128)) - # These padded extents are NOT M/N (and SymInt has no `//`/`+`), so give the - # scales their own fresh syms carrying the divisibility the padding - # guarantees (rowwise: 128 x 4; colwise: 4 x 128). - scale_rowwise_shape = (cute.sym_int32(divisibility=128), cute.sym_int32(divisibility=4)) - scale_colwise_shape = (cute.sym_int32(divisibility=4), cute.sym_int32(divisibility=128)) - ws_shape = (cute.sym_int32(), sym_N) # (blocks_Y, N); N ties to input N - # Native FP8/E8M0 dtypes at the FFI boundary (matches the DLPack dtype the C++ - # bridge sends). - out_dtype = cfg.FP8_DTYPE - scale_dtype = cutlass.Float8E8M0FNU - - in_fake = cute.runtime.make_fake_compact_tensor( - cfg.DTYPE, in_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=16 - ) - out_row_fake = ( - cute.runtime.make_fake_compact_tensor( - out_dtype, - out_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.ROWWISE - else None - ) - scale_row_fake = ( - cute.runtime.make_fake_compact_tensor( - scale_dtype, - scale_rowwise_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=4, - ) - if cfg.ROWWISE - else None - ) - out_col_fake = ( - cute.runtime.make_fake_compact_tensor( - out_dtype, - out_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.COLWISE - else None - ) - scale_col_fake = ( - cute.runtime.make_fake_compact_tensor( - scale_dtype, - scale_colwise_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=4, - ) - if cfg.COLWISE - else None - ) - amax_fake = ( - cute.runtime.make_fake_compact_tensor( - Float32, (1,), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4 - ) - if cfg.WITH_AMAX - else None - ) - # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so that - # the same compiled kernel serves both an absent and a present flag (see noop_flag_is_set). - # The null address here is only a compile-time placeholder for the pointer the C++ dispatcher - # passes at launch. - noop_fake = cute.runtime.nullptr(Float32, mem_space=cute.AddressSpace.gmem, assumed_align=4) - act_input_fake = ( - cute.runtime.make_fake_compact_tensor( - cfg.DTYPE, - in_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.WITH_DACT - else None - ) - workspace_fake = ( - cute.runtime.make_fake_compact_tensor( - Float32, ws_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=4 - ) - if cfg.WITH_DBIAS - else None - ) - - compiled = cute.compile( - kernel_obj, - in_fake, # mX - out_row_fake, - scale_row_fake, # mO_row, mS_row - out_col_fake, - scale_col_fake, # mO_col, mS_col - amax_fake, # mAmax - noop_fake, # mNoop (pointer to the cast_noop flag, may be null at launch) - act_input_fake, # mDActInput (backward slot, unused) - workspace_fake, # mWorkspace(backward slot, unused) - cute.runtime.make_fake_stream(), # stream (compiled as an explicit tvm-ffi - # "handle" arg; C++ passes the CUDA stream - # as void*) - options="--enable-tvm-ffi", - ) - return compiled - - -def get_mxfp8_quantization_function( - fn_name: str, - dtype: str, - fp8_dtype: str, - rowwise: bool, - colwise: bool, - with_gemm_swizzled_scales: bool, - with_amax: bool, - with_dbias: bool, - with_dact: bool, - with_act: bool, - use_2d_quantization: bool, - activation: str, -) -> bool: - """Compile the MXFP8 quantize kernel for this config and register it in the TVM-FFI global registry - under EXACTLY `fn_name` (the key the C++ dispatcher built; Python treats it as an opaque name). - Returns True if a kernel is successfully registered under `fn_name` (the C++ side then fetches it with GetGlobal(fn_name)); - False if the config is unsupported, so the caller caches the negative result and falls back to the CUDA C++ kernel. - """ - try: - # Already registered (e.g. by a prior call) -> supported. - if tvm_ffi.get_global_func(fn_name, allow_missing=True) is not None: - return True - - major, minor = device_compute_capability() - if major < 10: - logger.warning( - "CuTeDSL MXFP8 backend requires compute capability >= 10.0 (Blackwell), " - "but detected %d.%d; falling back to the CUDA C++ kernel.", - major, - minor, - ) - return False - - try: - cfg = MXFP8QuantizeConfig( - dtype=dtype, - fp8_dtype=fp8_dtype, - rowwise=rowwise, - colwise=colwise, - with_gemm_swizzled_scales=with_gemm_swizzled_scales, - with_amax=with_amax, - with_dbias=with_dbias, - with_dact=with_dact, - with_act=with_act, - use_2d_quantization=use_2d_quantization, - activation=activation, - ) - except ValueError as e: - # The exception message states exactly why the config is unsupported - # (unknown dtype/activation, dbias not implemented, ...). Surfacing it as a - # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. - logger.warning( - "CuTeDSL MXFP8 backend does not support this config, " - "falling back to the CUDA C++ kernel: %s", - e, - ) - return False - - logger.debug("Compiling CuTeDSL MXFP8 quantization kernel for %s", cfg) - compiled = compile_cutedsl_function_from_cfg(cfg) - # The returned compiled object is not neccessarily the compiled function itself. It could be a Python - # wrapper that parses arguments and calls the underlying function. This is needed in the general case - # as TVM-FFI is positional-only, however, for us this is unnecessary launch overhead as we pass the - # arguments in the correct order on the C++ side. - # Relevant issue: https://github.com/NVIDIA/cutlass/issues/3527 - # https://github.com/NVIDIA/cutlass/pull/3589 recommends using an explicit positional only marker (/) - # in the function signature, however it is not merged and it is something easy to miss and silently - # introduce unnecessarry overhead. - native = getattr(compiled, "__tvm_ffi_object__", lambda: None)() - tvm_ffi.register_global_func( - fn_name, native if native is not None else compiled, override=True - ) - return True - except Exception as e: # pylint: disable=broad-exception-caught - logger.error( - "CuTeDSL MXFP8 kernel compilation & registration failed, falling back to the CUDA" - " C++ kernel: %s", - e, - ) - # Unconditionally fallback to CUDA path because we can't tell if this exception is transient or permanent. - return False diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index ed05aa1c3c6..80a90a85ab8 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -83,6 +83,15 @@ def exp2f_rcp(scale_e8m0) -> Float32: return result +@cute.jit +def exp2_bf16x2_rcp(scale_e8m0) -> Uint32: + """Pack two BF16 copies of the E8M0 reciprocal computed by exp2f_rcp.""" + # These reciprocal powers of two are exact in BF16. Taking the upper bits + # also preserves exp2f_rcp's chosen subnormal and NaN encodings. + bf16_bits = exp2f_rcp(scale_e8m0).bitcast(Uint32) >> Uint32(16) + return bf16_bits | (bf16_bits << Uint32(16)) + + @cute.jit def pack_f32x2(lo: Float32, hi: Float32) -> Int64: """Pack two f32 scalars into a single 64-bit register (`floatx2` layout). @@ -258,3 +267,14 @@ def truncate_f32_f16(val: Float32) -> Float32: def is_packed16(dtype) -> bool: """True if `dtype` is one of the 16-bit packed input formats.""" return dtype is cutlass.BFloat16 or dtype is cutlass.Float16 + + +@cute.jit +def create_l2_policy(evict_last: cutlass.Constexpr[bool], fraction: Float32 = 1.0) -> Int64: + """Create the L2 eviction policy used by register-resident casts.""" + priority = "evict_last" if evict_last else "evict_first" + return cute.arch.inline_ptx( + f"createpolicy.fractional.L2::{priority}.b64 {{$w0}}, {{$r0}};", + write_only_types=[Int64], + read_only_args=[Float32(fraction)], + ) diff --git a/transformer_engine/common/CuTeDSL/utils_fp8.py b/transformer_engine/common/CuTeDSL/utils_fp8.py index b3b73e7998d..9a972e39469 100644 --- a/transformer_engine/common/CuTeDSL/utils_fp8.py +++ b/transformer_engine/common/CuTeDSL/utils_fp8.py @@ -180,6 +180,49 @@ def mul_f32x4_cvt_f32x4_to_fp8x4(fp8_dtype, relu: bool = False) -> Callable[..., return _build_mul_f32x4_cvt_f32x4_to_fp8x4(fp8_dtype, relu) +def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: + """Return a multiply/cast for two BF16 pairs, producing four packed FP8 bytes. + + CUDA 13.1+ supports direct BF16-to-FP8 conversion. Older compilers use + FP32 fused multiply-add, matching the corresponding ptx::mul_cvt_4x overload. + """ + out_op = "e5m2x2" if fp8_dtype is Float8E5M2 else "e4m3x2" + if cutlass.target_version(min_version="13.1"): + body = ( + ".reg.b32 y0, y1;\n\t" + "mul.rn.bf16x2 y0, {$r0}, {$r1};\n\t" + "mul.rn.bf16x2 y1, {$r2}, {$r3};\n\t" + f"cvt.rn.satfinite.{out_op}.bf16x2 z0, y0;\n\t" + f"cvt.rn.satfinite.{out_op}.bf16x2 z1, y1;\n\t" + ) + else: + body = ( + ".reg.b16 x0,x1,x2,x3,s0,s1,s2,s3;\n\t" + ".reg.f32 y0,y1,y2,y3;\n\t" + "mov.b32 {x0,x1}, {$r0};\n\t" + "mov.b32 {s0,s1}, {$r1};\n\t" + "mov.b32 {x2,x3}, {$r2};\n\t" + "mov.b32 {s2,s3}, {$r3};\n\t" + "fma.rn.f32.bf16 y0, x0, s0, 0f00000000;\n\t" + "fma.rn.f32.bf16 y1, x1, s1, 0f00000000;\n\t" + "fma.rn.f32.bf16 y2, x2, s2, 0f00000000;\n\t" + "fma.rn.f32.bf16 y3, x3, s3, 0f00000000;\n\t" + f"cvt.rn.satfinite.{out_op}.f32 z0, y1, y0;\n\t" + f"cvt.rn.satfinite.{out_op}.f32 z1, y3, y2;\n\t" + ) + asm = "{\n.reg.b16 z0,z1;\n\t" + body + "mov.b32 {$w0}, {z0,z1};\n}" + + @cute.jit + def fn(in0: Int32, scale0: Uint32, in1: Int32, scale1: Uint32) -> Uint32: + return cute.arch.inline_ptx( + asm, + write_only_types=[Uint32], + read_only_args=[in0, scale0, in1, scale1], + ) + + return fn + + def _build_mul_f32x2_cvt_packed16x4_to_fp8x4( in_dtype, fp8_dtype, relu: bool = False ) -> Callable[..., Uint32]: From 1be6447b126772e48d575508a83f1fde41d837c1 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 8 Oct 2026 21:42:05 +0000 Subject: [PATCH 02/29] wip Signed-off-by: Kaining Zhong --- docs/envvars.rst | 2 +- .../cast/mxfp8/quantize_mxfp8_common.py | 3 - .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 541 +++++++++--------- .../mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py | 2 +- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 136 +++-- .../mxfp8/quantize_mxfp8_g2s_bidimensional.py | 2 +- .../cast/mxfp8/quantize_mxfp8_general.py | 4 +- transformer_engine/common/CuTeDSL/utils.py | 47 ++ 8 files changed, 407 insertions(+), 330 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 46b70bbe46a..14f562f1069 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -150,7 +150,7 @@ General :Type: ``int`` (0 or 1) :Default: ``0`` :Description: Use CuTeDSL kernel implementations when available. This requires that Transformer Engine was built with ``NVTE_WITH_CUTEDSL=1``. - It also requires ``apache-tvm-ffi``, ``nvidia-cutlass-dsl`` to be installed and you are using CPython 3.10 or later. + It also requires ``apache-tvm-ffi`` and ``nvidia-cutlass-dsl>=4.8.0`` to be installed and you are using CPython 3.10 or later. If these requirements are not met, Transformer Engine will fall back to using CUDA C++ kernels if possible. .. envvar:: NVTE_WARN_IF_CUTEDSL_BACKEND_NOT_CHOSEN 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 f281937ffca..249e178abd6 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -43,9 +43,6 @@ # How many threads are in one warp THREADS_PER_WARP = 32 -# Prefer the current allocator API while retaining compatibility with older DSL releases. -_SmemAllocator = getattr(cutlass, "memory", cutlass.utils).SmemAllocator - # FP8E4M3 max representable value FP8E4M3_MAX_NORM = 448.0 FP8E4M3_MAX_NORM_RCP = 1.0 / FP8E4M3_MAX_NORM diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 0713d68f0fd..4a1692df68f 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -13,64 +13,22 @@ from transformer_engine.common.CuTeDSL.utils import ( abs_max_x2_bf16, + bf16_mx_scale_reciprocal, + bf16_mx_scale_reciprocal_pair, + bf16_mx_scale_bytes_pair, create_l2_policy, ) from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 from .quantize_mxfp8_common import ( + MXFP8_BLOCK_SCALING_SIZE, + THREADS_PER_WARP, MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, - _SmemAllocator, bf16_pair_magnitude, noop_flag_is_set, ) -@cute.jit -def bf16_mx_scale_reciprocal(magnitude: Uint32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: - """Exact BF16 reciprocal scale bits, including E8M0 zero, Inf and NaN cases.""" - rounded = (magnitude + Uint32(31)) & Uint32(0x7F80) - if rounded < Uint32(MAX_EXPONENT << 7): # pylint: disable=consider-using-max-builtin - rounded = Uint32(MAX_EXPONENT << 7) - reciprocal = Uint32((254 + MAX_EXPONENT) << 7) - rounded - if magnitude >= Uint32(0x7F80): - reciprocal = Uint32(0x0040) - if magnitude != Uint32(0x7F80): - reciprocal = Uint32(0x7FFF) - return reciprocal - - -@cute.jit -def bf16_mx_scale_reciprocal_pair(pair: Int32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: - """Compute two reciprocal scales without borrowing across packed BF16 halves.""" - magnitudes = Uint32(pair) & Uint32(0x7FFF7FFF) - rounded = (magnitudes + Uint32(0x001F001F)) & Uint32(0xFF80FF80) - floor = (MAX_EXPONENT << 23) | (MAX_EXPONENT << 7) - # Packed BF16 max accepts registers, so materialize the constant inside PTX. - clamped = cute.arch.inline_ptx( - "{ .reg.b32 floor_pair; " - f"mov.b32 floor_pair, {floor}; " - "max.xorsign.abs.bf16x2 {$w0}, {$r0}, floor_pair; }", - write_only_types=[Int32], - read_only_args=[rounded.bitcast(Int32)], - ) - bias = ((254 + MAX_EXPONENT) << 23) | ((254 + MAX_EXPONENT) << 7) - reciprocal = Uint32(bias) - Uint32(clamped) - if ((magnitudes + Uint32(0x00800080)) & Uint32(0x80008000)) != Uint32(0): - lo = bf16_mx_scale_reciprocal(magnitudes & Uint32(0x7FFF), MAX_EXPONENT) - hi = bf16_mx_scale_reciprocal(magnitudes >> Uint32(16), MAX_EXPONENT) - reciprocal = lo | (hi << Uint32(16)) - return reciprocal - - -@cute.jit -def bf16_mx_scale_bytes_pair(reciprocal: Uint32) -> Uint32: - """Gather the E8M0 bytes of two packed BF16 reciprocals into the low half.""" - bytes_pair = (Uint32(0x01FE01FE) - ((reciprocal >> Uint32(7)) & Uint32(0x00FF00FF))) & Uint32( - 0x00FF00FF - ) - return (bytes_pair & Uint32(0xFF)) | ((bytes_pair >> Uint32(8)) & Uint32(0xFF00)) - - @cute.jit def quantize_bidimensional_mxfp8_g2r( gX_tile: cute.Tensor, @@ -81,97 +39,117 @@ def quantize_bidimensional_mxfp8_g2r( cfg: cutlass.Constexpr[MXFP8QuantizeConfig], CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], - ROWS_PER_THREAD: cutlass.Constexpr[int], - ELEMENTS_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], cache_policy: Int64, ): """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile in both directions. - TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD = 32, so every columnwise - MXFP8 block is contained in the tile. Each warp owns ROWS_PER_THREAD - consecutive rows, and each lane owns ELEMENTS_PER_THREAD adjacent columns. + TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD = 32, so every columnwise + MXFP8 block is contained in the tile. Each warp owns ELEMENTS_Y_PER_THREAD + consecutive rows, and each lane owns ELEMENTS_X_PER_THREAD adjacent columns. Load the tile once, then reuse its registers for both quantization passes. """ - TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD - TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD - LANES_PER_MX_BLOCK = 32 // ELEMENTS_PER_THREAD - COLUMN_PAIRS = TILE_COLS // 2 - COLUMN_PAIRS_PER_THREAD = ELEMENTS_PER_THREAD // 2 - PAIRS_PER_HALF = COLUMN_PAIRS_PER_THREAD // 2 + # Shape of the tile + TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD + TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD + + # How many threads cooperate to process a single MXFP8 block (32 elements) for both quantized directions + THREADS_X_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_X_PER_THREAD + THREADS_Y_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_Y_PER_THREAD + MAX_EXPONENT = 8 if cfg.FP8_DTYPE is cutlass.Float8E4M3FN else 15 SCALE_DTYPE = cutlass.Float8E8M0FNU + tidx, _, _ = cute.arch.thread_idx() - warp, lane = tidx // CTA_THREADS_X, tidx % CTA_THREADS_X - _, tv_payload = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)), + thread_layout = cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)) + _, tv_layout = cute.make_layout_tv( + thr_layout=thread_layout, val_layout=cute.make_layout( - (ROWS_PER_THREAD, ELEMENTS_PER_THREAD), stride=(ELEMENTS_PER_THREAD, 1) + (ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD), stride=(ELEMENTS_X_PER_THREAD, 1) ), ) - # Reshape each thread's values into (columns, rows), with contiguous columns. - fragment_layout = cute.make_layout( - (ELEMENTS_PER_THREAD, ROWS_PER_THREAD), stride=(1, ELEMENTS_PER_THREAD) - ) - gX_thread = cute.composition(cute.composition(gX_tile, tv_payload)[tidx, None], fragment_layout) - gO_row_thread = cute.composition( - cute.composition(gO_row_tile, tv_payload)[tidx, None], fragment_layout - ) - gO_col_thread = cute.composition( - cute.composition(gO_col_tile, tv_payload)[tidx, None], fragment_layout - ) - # Each group of lanes shares one rowwise scale for each of its rows. + + gX_thread = cute.composition(gX_tile, tv_layout)[tidx, None] + gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] + gO_col_thread = cute.composition(gO_col_tile, tv_layout)[tidx, None] + + # Group CTA threads by THREADS_X_PER_MX_BLOCK since they share the same rowwise scale for a MXFP8 block. + # And assign stride 0 to that submode tv_row_scale = cute.make_layout( ( - (LANES_PER_MX_BLOCK, (CTA_THREADS_X // LANES_PER_MX_BLOCK, CTA_THREADS_Y)), - ROWS_PER_THREAD, + (THREADS_X_PER_MX_BLOCK, (CTA_THREADS_X // THREADS_X_PER_MX_BLOCK, CTA_THREADS_Y)), + ELEMENTS_Y_PER_THREAD, + ), + stride=( + (0, (TILE_ROWS, ELEMENTS_Y_PER_THREAD)), + 1, ), - stride=((0, (TILE_ROWS, ROWS_PER_THREAD)), 1), ) gS_row_thread = cute.composition(gS_row_tile, tv_row_scale)[tidx, None] - # Warp 0 writes one scale per column, in each lane's original column order. - gS_col_thread = cute.local_tile(gS_col_tile, (1, ELEMENTS_PER_THREAD), (0, lane))[0, None] - # Keep the complete input fragment; reuse one row's output registers. - rX = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, ROWS_PER_THREAD), cfg.DTYPE) - rO = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.FP8_DTYPE) + # Group CTA threads by THREADS_Y_PER_MX_BLOCK since they share the same columnwise scale for a MXFP8 block. + # And assign stride 0 to that submode + tv_col_scale = cute.make_layout( + ( + (CTA_THREADS_X, (THREADS_Y_PER_MX_BLOCK, CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK)), + ELEMENTS_X_PER_THREAD, + ), + stride=( + (ELEMENTS_X_PER_THREAD, (0, 1)), + 1, + ), + ) + gS_col_thread = cute.composition(gS_col_tile, tv_col_scale)[tidx, None] + + # Input fragment contains all rows it handles; output fragment contains only one row and we write it back row by row + rX = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), cfg.DTYPE) + rO = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.FP8_DTYPE) + + # For rowwise scales, we write one row at a time rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) - rS_col = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), SCALE_DTYPE) + load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), cfg.DTYPE, - num_bits_per_copy=ELEMENTS_PER_THREAD * cfg.DTYPE.width, + num_bits_per_copy=ELEMENTS_X_PER_THREAD * cfg.DTYPE.width, invariant=True, ) store_atom = cute.make_copy_atom( cute.nvgpu.CopyR2GOp(), cfg.FP8_DTYPE, - num_bits_per_copy=ELEMENTS_PER_THREAD * cfg.FP8_DTYPE.width, + num_bits_per_copy=ELEMENTS_X_PER_THREAD * cfg.FP8_DTYPE.width, ) + row_scale_atom = cute.make_copy_atom(cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8) col_scale_atom = cute.make_copy_atom( - cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=ELEMENTS_PER_THREAD * 8 + cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=ELEMENTS_X_PER_THREAD * 8 ) + scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(cfg.FP8_DTYPE) - rO_u32 = cute.recast_tensor(rO, Uint32) - for row_idx in cutlass.range_constexpr(ROWS_PER_THREAD): + # Copy all rows a thread sweeps into the RMEM fragment at once + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): cute.copy(load_atom, gX_thread[None, row_idx], rX[None, row_idx], cache_policy=cache_policy) @cute.jit - def extract_row_amax(values_i32: cute.Tensor) -> Uint32: + def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: """Reduce a 32-column MXFP8 block across its two or four lane owners.""" amax = values_i32[0] - for pair_idx in cutlass.range_constexpr(1, cute.size(values_i32)): - amax = abs_max_x2_bf16(amax, values_i32[pair_idx]) - for stage in cutlass.range_constexpr(LANES_PER_MX_BLOCK.bit_length() - 1): + for bf16x2_idx in cutlass.range_constexpr(1, cute.size(values_i32)): + amax = abs_max_x2_bf16(amax, values_i32[bf16x2_idx]) + for stage in cutlass.range_constexpr(THREADS_X_PER_MX_BLOCK.bit_length() - 1): amax = abs_max_x2_bf16(amax, cute.arch.shuffle_sync_bfly(amax, 1 << stage)) return bf16_pair_magnitude(amax) - # Rowwise pass: lane shuffles close each 32-column reduction. - for row_idx in cutlass.range_constexpr(ROWS_PER_THREAD): + # Rowwise pass: produce each row's rowwise scales and quantized outputs, and write them back to GMEM immediately. + rO_u32 = cute.recast_tensor(rO, Uint32) + # Only one thread from threads that share the MXFP8 block writes the rowwise scale + is_writer_thread = tidx % THREADS_X_PER_MX_BLOCK == 0 + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - reciprocal = bf16_mx_scale_reciprocal(extract_row_amax(rX_row_i32), MAX_EXPONENT) + reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) reciprocal_packed = reciprocal | (reciprocal << Uint32(16)) for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): rO_u32[pack_idx] = scale_and_pack( @@ -180,9 +158,10 @@ def extract_row_amax(values_i32: cute.Tensor) -> Uint32: rX_row_i32[2 * pack_idx + 1], reciprocal_packed, ) + # rO_u32 and rO share the same memory cute.copy(store_atom, rO, gO_row_thread[None, row_idx], cache_policy=cache_policy) - if lane % LANES_PER_MX_BLOCK == 0: - rS_row[0] = Uint8(Uint32(254) - (reciprocal >> Uint32(7))).bitcast(SCALE_DTYPE) + if is_writer_thread: + rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_packed)).bitcast(SCALE_DTYPE) cute.copy( row_scale_atom, rS_row, @@ -190,86 +169,113 @@ def extract_row_amax(values_i32: cute.Tensor) -> Uint32: cache_policy=cache_policy, ) - # Columnwise pass: first reduce each thread's rows into one maximum per column. - rColumnAmax = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.DTYPE) - rColumnScale = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.DTYPE) - rColumnAmax_i32 = cute.recast_tensor(rColumnAmax, Int32) - rColumnScale_u32 = cute.recast_tensor(rColumnScale, Uint32) + # Columnwise pass: only produce the colwise scale. + # We first reduce each thread's rows into one maximum per column before reduce across different threads in the CTA + rAmax_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) + + # Gather the amax for each column that this thread sweeps rX_i32 = cute.recast_tensor(rX, Int32) - for pair_idx in cutlass.range_constexpr(COLUMN_PAIRS_PER_THREAD): - amax = rX_i32[pair_idx, 0] - for row_idx in cutlass.range_constexpr(1, ROWS_PER_THREAD): - amax = abs_max_x2_bf16(amax, rX_i32[pair_idx, row_idx]) - rColumnAmax_i32[pair_idx] = amax - - allocator = _SmemAllocator() - sColumnAmax = allocator.allocate_tensor( - cfg.DTYPE, - cute.make_layout((CTA_THREADS_Y, TILE_COLS), stride=(TILE_COLS, 1)), - byte_alignment=16, + rAmax_col_i32 = cute.recast_tensor(rAmax_col, Int32) + for bf16x2_idx in cutlass.range_constexpr(cute.size(rX_i32, mode=[0])): + amax = rX_i32[bf16x2_idx, 0] + for row_idx in cutlass.range_constexpr(1, ELEMENTS_Y_PER_THREAD): + amax = abs_max_x2_bf16(amax, rX_i32[bf16x2_idx, row_idx]) + rAmax_col_i32[bf16x2_idx] = amax + + # Threads that process a columnwise MXFP8 block don't belong to the same warp so we have to use SMEM to reduce the scale + allocator = cutlass.memory.SmemAllocator() + sAmaxs_col_scratch_layout = cute.zipped_product( + rAmax_col.layout, + thread_layout, ) - sColumnScale = allocator.allocate_tensor( - cfg.DTYPE, cute.make_layout((TILE_COLS,)), byte_alignment=16 + sAmaxs_col_scratch = allocator.allocate_tensor( + cfg.DTYPE, sAmaxs_col_scratch_layout, byte_alignment=16 ) - sColumnAmax_i32 = cute.recast_tensor(sColumnAmax, Int32) - sColumnScale_u32 = cute.recast_tensor(sColumnScale, Uint32) - # Group each lane's column pairs into two halves. Placing the first halves - # of all lanes before the second halves preserves CUDA's conflict-free - # vectorized shared-memory accesses, without calculating pointer offsets. - column_pair_layout = cute.make_layout( - ((PAIRS_PER_HALF, 2), CTA_THREADS_X), - stride=((1, COLUMN_PAIRS // 2), PAIRS_PER_HALF), + sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) + sS_col_reduction_scratch = allocator.allocate_tensor( + cfg.DTYPE, (CTA_THREADS_Y, CTA_THREADS_X), byte_alignment=16 ) - sColumnAmax_thread = cute.composition(sColumnAmax_i32, (None, column_pair_layout))[ - warp, (None, lane) - ] - sColumnScale_thread = cute.composition(sColumnScale_u32, column_pair_layout)[None, lane] - for half_idx in cutlass.range_constexpr(2): - cute.autovec_copy( - cute.local_tile(rColumnAmax_i32, (PAIRS_PER_HALF,), (half_idx,)), - cute.local_tile(sColumnAmax_thread, (PAIRS_PER_HALF,), (half_idx,)), - ) - cute.arch.sync_threads() + sS_col_reduction_scratch_u32 = cute.recast_tensor(sS_col_reduction_scratch, Uint32) - # One thread reduces each column pair across all warps (all 32 tile rows). - if tidx < COLUMN_PAIRS: - amax = sColumnAmax_i32[0, tidx] - for warp_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): - amax = abs_max_x2_bf16(amax, sColumnAmax_i32[warp_idx, tidx]) - sColumnScale_u32[tidx] = bf16_mx_scale_reciprocal_pair(amax, MAX_EXPONENT) + # Every thread copies the RMEM fragment to SMEM scratch space for reduction across warps later + cute.autovec_copy( + rAmax_col, + sAmaxs_col_scratch[None, tidx], + ) + + # Reduce thread read amaxes from SMEM then calculate & write the columnwise scale to SMEM scratch space + is_reduce_thread = tidx < CTA_THREADS_X + if is_reduce_thread: + amaxX2 = sAmaxs_col_scratch_i32[0, tidx] + for scratch_row in cutlass.range_constexpr(CTA_THREADS_Y): + amaxX2 = abs_max_x2_bf16(amaxX2, sAmaxs_col_scratch_i32[scratch_row, tidx]) + sS_col_reduction_scratch_u32[tidx] = bf16_mx_scale_reciprocal(amaxX2, MAX_EXPONENT) cute.arch.sync_threads() - # The same layout returns the reciprocals to each lane's original columns. - for half_idx in cutlass.range_constexpr(2): - cute.autovec_copy( - cute.local_tile(sColumnScale_thread, (PAIRS_PER_HALF,), (half_idx,)), - cute.local_tile(rColumnScale_u32, (PAIRS_PER_HALF,), (half_idx,)), - ) - if warp == 0: + # Let the reduction threads write the RMEM columnwise scales back to GMEM + if is_reduce_thread: rS_col_u32 = cute.recast_tensor(rS_col, Uint32) - for pack_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): - rS_col_u32[pack_idx] = bf16_mx_scale_bytes_pair(rColumnScale_u32[2 * pack_idx]) | ( - bf16_mx_scale_bytes_pair(rColumnScale_u32[2 * pack_idx + 1]) << Uint32(16) + for packed_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): + rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_u32[2 * packed_idx]) | ( + bf16_mx_scale_bytes_pair(rS_col_u32[2 * packed_idx + 1]) << Uint32(16) ) cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) - for row_idx in cutlass.range_constexpr(ROWS_PER_THREAD): - rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): - rO_u32[pack_idx] = scale_and_pack( - rX_row_i32[2 * pack_idx], - rColumnScale_u32[2 * pack_idx], - rX_row_i32[2 * pack_idx + 1], - rColumnScale_u32[2 * pack_idx + 1], - ) - cute.copy(store_atom, rO, gO_col_thread[None, row_idx], cache_policy=cache_policy) + # Every thread reads the reduced columnwise scales from SMEM to RMEM + rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) + cute.autovec_copy( + sS_col_reduction_scratch[tidx], + rS_col + ) + + # And then every thread uses the columnwise scales to calculate quantized outputs and writes it to GMEM + + + + + # # One thread reduces each column pair across all warps (all 32 tile rows). + # if tidx < cute.size(sColumnScales) // 2: + # sColumnAmaxs_i32 = cute.recast_tensor(sColumnAmaxs, Int32) + # sColumnScales_u32 = cute.recast_tensor(sColumnScales, Uint32) + # amax = sColumnAmaxs_i32[0, tidx] + # for warp_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): + # amax = abs_max_x2_bf16(amax, sColumnAmaxs_i32[warp_idx, tidx]) + # sColumnScales_u32[tidx] = bf16_mx_scale_reciprocal_pair(amax, MAX_EXPONENT) + # cute.arch.sync_threads() + + # # The same layout returns the reciprocals to each lane's original columns. + # for half_idx in cutlass.range_constexpr(2): + # cute.autovec_copy( + # cute.local_tile(sColumnScale_thread, (cute.size(rColumnScales) // 2,), (half_idx,)), + # cute.local_tile(rColumnScales, (cute.size(rColumnScales) // 2,), (half_idx,)), + # ) + + # rColumnScales_u32 = cute.recast_tensor(rColumnScales, Uint32) + # if warp == 0: + # rS_col_u32 = cute.recast_tensor(rS_col, Uint32) + # for pack_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): + # rS_col_u32[pack_idx] = bf16_mx_scale_bytes_pair(rColumnScales_u32[2 * pack_idx]) | ( + # bf16_mx_scale_bytes_pair(rColumnScales_u32[2 * pack_idx + 1]) << Uint32(16) + # ) + # cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) + + # for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + # rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + # for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): + # rO_u32[pack_idx] = scale_and_pack( + # rX_row_i32[2 * pack_idx], + # rColumnScales_u32[2 * pack_idx], + # rX_row_i32[2 * pack_idx + 1], + # rColumnScales_u32[2 * pack_idx + 1], + # ) + # cute.copy(store_atom, rO, gO_col_thread[None, row_idx], cache_policy=cache_policy) class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): """Load one CTA tile into registers and quantize it along both axes. - TILE_ROWS = _CTA_THREADS_Y * _ROWS_PER_THREAD = 32. - TILE_COLS = _CTA_THREADS_X * ELEMENTS_PER_THREAD = 256 or 512. + TILE_ROWS = _CTA_THREADS_Y * _ELEMENTS_Y_PER_THREAD = 32. + TILE_COLS = _CTA_THREADS_X * ELEMENTS_X_PER_THREAD = 256 or 512. The 32-row tile contains complete columnwise MXFP8 blocks; shared memory combines the partial column maxima from the eight warps. Requires aligned BF16 input, M divisible by 32, N divisible by 256 and plain scale layouts. @@ -277,11 +283,10 @@ class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): _CTA_THREADS_Y = 8 _CTA_THREADS_X = 32 - _ROWS_PER_THREAD = 4 - _ELEMENTS_PER_THREAD = 16 def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg + self._ELEMENTS_Y_PER_THREAD = MXFP8_BLOCK_SCALING_SIZE // self._CTA_THREADS_Y # 4 @cute.jit def __call__( @@ -298,8 +303,8 @@ def __call__( stream: CUstream, ): M, N = mX.shape - TILE_ROWS = self._CTA_THREADS_Y * self._ROWS_PER_THREAD - WIDE_TILE_COLS = self._CTA_THREADS_X * self._ELEMENTS_PER_THREAD + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + WIDE_TILE_COLS = self._CTA_THREADS_X * 16 if N % WIDE_TILE_COLS == 0: if Int64(M // TILE_ROWS) * (N // WIDE_TILE_COLS) >= 4096: self._launch( @@ -310,10 +315,7 @@ def __call__( mS_col, mNoop, stream, - CTA_THREADS_Y=self._CTA_THREADS_Y, - CTA_THREADS_X=self._CTA_THREADS_X, - ROWS_PER_THREAD=self._ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=4, ) else: @@ -325,10 +327,7 @@ def __call__( mS_col, mNoop, stream, - CTA_THREADS_Y=self._CTA_THREADS_Y, - CTA_THREADS_X=self._CTA_THREADS_X, - ROWS_PER_THREAD=self._ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=1, ) else: @@ -340,10 +339,7 @@ def __call__( mS_col, mNoop, stream, - CTA_THREADS_Y=self._CTA_THREADS_Y, - CTA_THREADS_X=self._CTA_THREADS_X, - ROWS_PER_THREAD=self._ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD // 2, + ELEMENTS_X_PER_THREAD=8, CLUSTER_WIDTH=8, ) @@ -357,56 +353,74 @@ def _launch( mS_col: cute.Tensor, mNoop: cute.Pointer, stream: CUstream, - CTA_THREADS_Y: cutlass.Constexpr[int], - CTA_THREADS_X: cutlass.Constexpr[int], - ROWS_PER_THREAD: cutlass.Constexpr[int], - ELEMENTS_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], CLUSTER_WIDTH: cutlass.Constexpr[int], ): - assert CTA_THREADS_X == 32 - assert CTA_THREADS_Y * ROWS_PER_THREAD == 32 - assert ELEMENTS_PER_THREAD in (8, 16) + assert ELEMENTS_X_PER_THREAD in (8, 16) + + # Dispatch checks these alignments before selecting this kernel. + mX = cute.make_tensor(mX.iterator.align(32), mX.layout) + mS_col = cute.make_tensor(mS_col.iterator.align(16), mS_col.layout) M, N = mX.shape - TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD - TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD - THREADS_PER_CTA = CTA_THREADS_Y * CTA_THREADS_X + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + TILE_COLS = self._CTA_THREADS_X * ELEMENTS_X_PER_THREAD + THREADS_PER_CTA = self._CTA_THREADS_Y * self._CTA_THREADS_X + grid_cols = N // TILE_COLS grid_rows = M // TILE_ROWS cluster_cols = Int32(CLUSTER_WIDTH) + # Adjust cluster size so that the CTAs in the column direction can be evenly divided into clusters. while cluster_cols > 1 and grid_cols % cluster_cols != 0: cluster_cols -= 1 # Keep column CTAs on grid.x so clusters read adjacent column tiles. grid = [grid_cols, grid_rows] block = [THREADS_PER_CTA] cluster = [cluster_cols, 1, 1] - # Preserve CUDA's compile-time strides for the two common column counts. + + # Hardcode the layout using static N for some of the commonly used shapes if N == 2048: + payload_layout = cute.make_layout((M, 2048), stride=(2048, 1)) + gX = cute.make_tensor(mX.iterator, payload_layout) + gO_row = cute.make_tensor(mO_row.iterator, payload_layout) + gO_col = cute.make_tensor(mO_col.iterator, payload_layout) + gS_row = cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 64), stride=(64, 1)), + ) + gS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout((cute.size(mS_col, mode=[0]), 2048), stride=(2048, 1)), + ) self.kernel( - mX, - mO_row, - mS_row, - mO_col, - mS_col, + gX, + gO_row, + gS_row, + gO_col, + gS_col, mNoop, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ROWS_PER_THREAD=ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, - N_CONST=2048, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) elif N == 4096: + payload_layout = cute.make_layout((M, 4096), stride=(4096, 1)) + gX = cute.make_tensor(mX.iterator, payload_layout) + gO_row = cute.make_tensor(mO_row.iterator, payload_layout) + gO_col = cute.make_tensor(mO_col.iterator, payload_layout) + gS_row = cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 128), stride=(128, 1)), + ) + gS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout((cute.size(mS_col, mode=[0]), 4096), stride=(4096, 1)), + ) self.kernel( - mX, - mO_row, - mS_row, - mO_col, - mS_col, + gX, + gO_row, + gS_row, + gO_col, + gS_col, mNoop, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ROWS_PER_THREAD=ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, - N_CONST=4096, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) else: self.kernel( @@ -416,11 +430,7 @@ def _launch( mO_col, mS_col, mNoop, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ROWS_PER_THREAD=ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, - N_CONST=0, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) @cute.kernel @@ -432,59 +442,58 @@ def kernel( mO_col: cute.Tensor, mS_col: cute.Tensor, mNoop: cute.Pointer, - CTA_THREADS_Y: cutlass.Constexpr[int], - CTA_THREADS_X: cutlass.Constexpr[int], - ROWS_PER_THREAD: cutlass.Constexpr[int], - ELEMENTS_PER_THREAD: cutlass.Constexpr[int], - N_CONST: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], ): - """Check noop before constructing or accessing any CTA tile.""" + """Skip the CTA when the noop flag is set; otherwise run the kernel body.""" if not noop_flag_is_set(mNoop): - M, N = mX.shape - if cutlass.const_expr(N_CONST != 0): - N = N_CONST - payload_layout = cute.make_layout((M, N), stride=(N, 1)) - gX = cute.make_tensor(mX.iterator.align(32), payload_layout) - gO_row = cute.make_tensor(mO_row.iterator, payload_layout) - gO_col = cute.make_tensor(mO_col.iterator, payload_layout) - gS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout((cute.size(mS_row, mode=[0]), N // 32), stride=(N // 32, 1)), - ) - gS_col = cute.make_tensor( - mS_col.iterator.align(16), - cute.make_layout((cute.size(mS_col, mode=[0]), N), stride=(N, 1)), - ) - else: - # The dispatcher checks these alignments before choosing this kernel. - gX = cute.make_tensor(mX.iterator.align(32), mX.layout) - gO_row = mO_row - gO_col = mO_col - gS_row = mS_row - gS_col = cute.make_tensor(mS_col.iterator.align(16), mS_col.layout) - cta_col, cta_row, _ = cute.arch.block_idx() - cta_coord = (cta_row, cta_col) - TILE_ROWS = CTA_THREADS_Y * ROWS_PER_THREAD - TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD - TILER = (TILE_ROWS, TILE_COLS) - ROW_SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) - COL_SCALE_TILER = (TILE_ROWS // 32, TILE_COLS) - gX_tile = cute.local_tile(gX, TILER, cta_coord) - gO_row_tile = cute.local_tile(gO_row, TILER, cta_coord) - gO_col_tile = cute.local_tile(gO_col, TILER, cta_coord) - gS_row_tile = cute.local_tile(gS_row, ROW_SCALE_TILER, cta_coord) - gS_col_tile = cute.local_tile(gS_col, COL_SCALE_TILER, cta_coord) - cache_policy = create_l2_policy(True) - quantize_bidimensional_mxfp8_g2r( - gX_tile, - gO_row_tile, - gS_row_tile, - gO_col_tile, - gS_col_tile, - self.cfg, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ROWS_PER_THREAD=ROWS_PER_THREAD, - ELEMENTS_PER_THREAD=ELEMENTS_PER_THREAD, - cache_policy=cache_policy, + self._kernel_main( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: cute.Tensor, + mS_col: cute.Tensor, + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + ): + """Construct CTA tile views and run bidimensional quantization.""" + cta_col, cta_row, _ = cute.arch.block_idx() + cta_coord = (cta_row, cta_col) + + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + TILE_COLS = self._CTA_THREADS_X * ELEMENTS_X_PER_THREAD + + TILER = (TILE_ROWS, TILE_COLS) + ROW_SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) + COL_SCALE_TILER = (TILE_ROWS // 32, TILE_COLS) + + gX_tile = cute.local_tile(mX, TILER, cta_coord) + gO_row_tile = cute.local_tile(mO_row, TILER, cta_coord) + gO_col_tile = cute.local_tile(mO_col, TILER, cta_coord) + gS_row_tile = cute.local_tile(mS_row, ROW_SCALE_TILER, cta_coord) + gS_col_tile = cute.local_tile(mS_col, COL_SCALE_TILER, cta_coord) + + cache_policy = create_l2_policy(True) + + quantize_bidimensional_mxfp8_g2r( + gX_tile, + gO_row_tile, + gS_row_tile, + gO_col_tile, + gS_col_tile, + self.cfg, + CTA_THREADS_Y=self._CTA_THREADS_Y, + CTA_THREADS_X=self._CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=self._ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + cache_policy=cache_policy, + ) 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 bbdf9a4b510..552b6429465 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 @@ -177,7 +177,7 @@ def _kernel_main( class SharedStorage: buf: cute.struct.Align[cute.struct.MemRange[Float8E8M0FNU, CTA_Y * CTA_X], 16] - storage = cutlass.utils.SmemAllocator().allocate(SharedStorage) + storage = cutlass.memory.SmemAllocator().allocate(SharedStorage) sScale = storage.buf.get_tensor(cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1))) # sScale is (CTA_Y, CTA_X):(CTA_X, 1), which is the same layout as tv_layout_scale # so sS_thread is really just a 1-element e8m0 buffer for this thread's scale. 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 95882a4161f..edc6cf850f5 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 @@ -109,7 +109,7 @@ def quantize_rowwise_mxfp8_g2r( cute.copy(load_atom, gX_thread, rX[None, tile_idx], cache_policy=input_policy) @cute.jit - def extract_tile_amax(values_i32: cute.Tensor) -> Float32: + def extract_mx_block_amax(values_i32: cute.Tensor) -> Float32: """Return the 32-value MXFP8 block amax shared by two adjacent threads. Balanced eight-pair tree, followed by one butterfly between lane partners. @@ -134,7 +134,7 @@ def extract_tile_amax(values_i32: cute.Tensor) -> Float32: # Reinterpret the native tensors only for the packed BF16/FP8 instructions. rX_tile_i32 = cute.recast_tensor(rX[None, tile_idx], Int32) rO_u32 = cute.recast_tensor(rO, Uint32) - amax = extract_tile_amax(rX_tile_i32) + amax = extract_mx_block_amax(rX_tile_i32) exponent = cvt_f32_to_fp8e8m0fnu(amax * cfg.MAX_NORM_RCP) # Collect 4 scale bytes and write them using one thread @@ -306,63 +306,87 @@ def kernel( NUM_TILES: cutlass.Constexpr[int], CHECK_BOUNDS: cutlass.Constexpr[bool], ): - """Check the noop flag before selecting or accessing any GMEM tile.""" + """Skip the CTA when the noop flag is set; otherwise run the kernel body.""" if not noop_flag_is_set(mNoop): - # Dispatch checks 32-byte input alignment before selecting this kernel. - gX = cute.make_tensor(mX.iterator.align(32), mX.layout) - # CUDA's x/y axes hold row/column tiles in the same order as local_tile. - cta_coord = cute.arch.block_idx()[:2] - M, N = mX.shape - TILE_ROWS = CTA_THREADS_Y - TILE_COLS = CTA_THREADS_X * self._ELEMENTS_PER_THREAD - CTA_ROWS = TILE_ROWS - CTA_COLS = TILE_COLS * NUM_TILES - - gS = mS - # Apply the swizzled scale layout if requested - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - mS_swizzled, _ = derive_swizzled_scale_layout(M, N, True, False, mS, None) - gS = cute.composition(mS_swizzled, (cute.make_layout(M), cute.make_layout(N // 32))) - - # CTA_TILER is the block that a CTA processes in total, which consists of NUM_TILES horizontal tiles - CTA_TILER = (CTA_ROWS, CTA_COLS) - CTA_SCALE_TILER = (CTA_ROWS, CTA_COLS // 32) - gX_cta = cute.local_tile(gX, CTA_TILER, cta_coord) - gO_cta = cute.local_tile(mO, CTA_TILER, cta_coord) - gS_cta = cute.local_tile(gS, CTA_SCALE_TILER, cta_coord) - - fraction = Float32(1.0) - if cutlass.const_expr(self._L2_CACHED_CTA_PERCENT > 0): - fraction = Float32(0.0) - # A linear index is needed only for the L2 policy split. - cta_id = cute.crd2idx( - (Int64(cta_coord[0]), Int64(cta_coord[1])), cute.arch.grid_dim()[:2] - ) - if cta_id >= first_streaming_cta: - fraction = Float32(1.0) - input_policy = create_l2_policy(False, fraction) - output_policy = create_l2_policy(True) - - # Divide CTA_TILER to multiple TILER shape TILES - TILER = (TILE_ROWS, TILE_COLS) - SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) - # Mode 1's first index is 0 because CTA_ROWS == TILE_ROWS (there is only 1 tile on the row dimension). - gX_tiles = cute.local_tile(gX_cta, TILER, (0, None)) - gO_tiles = cute.local_tile(gO_cta, TILER, (0, None)) - mS_tiles = cute.local_tile(gS_cta, SCALE_TILER, (0, None)) - - quantize_rowwise_mxfp8_g2r( - gX_tiles, - gO_tiles, - mS_tiles, - self.cfg, - M, - N, + self._kernel_main( + mX, + mO, + mS, + first_streaming_cta, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, NUM_TILES=NUM_TILES, CHECK_BOUNDS=CHECK_BOUNDS, - input_policy=input_policy, - output_policy=output_policy, ) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mO: cute.Tensor, + mS: cute.Tensor, + first_streaming_cta: Int64, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + NUM_TILES: cutlass.Constexpr[int], + CHECK_BOUNDS: cutlass.Constexpr[bool], + ): + """Construct CTA tile views and run the specialized quantization.""" + # Dispatch checks 32-byte input alignment before selecting this kernel. + gX = cute.make_tensor(mX.iterator.align(32), mX.layout) + # CUDA's x/y axes hold row/column tiles in the same order as local_tile. + cta_coord = cute.arch.block_idx()[:2] + M, N = mX.shape + TILE_ROWS = CTA_THREADS_Y + TILE_COLS = CTA_THREADS_X * self._ELEMENTS_PER_THREAD + CTA_ROWS = TILE_ROWS + CTA_COLS = TILE_COLS * NUM_TILES + + gS = mS + # Apply the swizzled scale layout if requested + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_swizzled, _ = derive_swizzled_scale_layout(M, N, True, False, mS, None) + gS = cute.composition(mS_swizzled, (cute.make_layout(M), cute.make_layout(N // 32))) + + # CTA_TILER is the block that a CTA processes in total, which consists of NUM_TILES horizontal tiles + CTA_TILER = (CTA_ROWS, CTA_COLS) + CTA_SCALE_TILER = (CTA_ROWS, CTA_COLS // 32) + gX_cta = cute.local_tile(gX, CTA_TILER, cta_coord) + gO_cta = cute.local_tile(mO, CTA_TILER, cta_coord) + gS_cta = cute.local_tile(gS, CTA_SCALE_TILER, cta_coord) + + fraction = Float32(1.0) + if cutlass.const_expr(self._L2_CACHED_CTA_PERCENT > 0): + fraction = Float32(0.0) + # A linear index is needed only for the L2 policy split. + cta_id = cute.crd2idx( + (Int64(cta_coord[0]), Int64(cta_coord[1])), cute.arch.grid_dim()[:2] + ) + if cta_id >= first_streaming_cta: + fraction = Float32(1.0) + input_policy = create_l2_policy(False, fraction) + output_policy = create_l2_policy(True) + + # Divide CTA_TILER to multiple TILER shape TILES + TILER = (TILE_ROWS, TILE_COLS) + SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) + # Mode 1's first index is 0 because CTA_ROWS == TILE_ROWS (there is only 1 tile on the row dimension). + gX_tiles = cute.local_tile(gX_cta, TILER, (0, None)) + gO_tiles = cute.local_tile(gO_cta, TILER, (0, None)) + mS_tiles = cute.local_tile(gS_cta, SCALE_TILER, (0, None)) + + quantize_rowwise_mxfp8_g2r( + gX_tiles, + gO_tiles, + mS_tiles, + self.cfg, + M, + N, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + NUM_TILES=NUM_TILES, + CHECK_BOUNDS=CHECK_BOUNDS, + input_policy=input_policy, + output_policy=output_policy, + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py index 81a94382bc9..89cfff90093 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py @@ -446,7 +446,7 @@ class SharedStorage: cute.struct.MemRange[Float8E8M0FNU, self._NUM_TILES * self._TILE_COLS], 16 ] - smem = cutlass.utils.SmemAllocator() + smem = cutlass.memory.SmemAllocator() storage = smem.allocate(SharedStorage) tile_layout = cute.make_layout( diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py index 26ff4c10b2d..a2b4f5e55ff 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py @@ -867,7 +867,7 @@ class SharedStorage: ] sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] - smem = cutlass.utils.SmemAllocator() + smem = cutlass.memory.SmemAllocator() storage = smem.allocate(SharedStorage) # Apply the layout to the allocated shared memory buffers so the first rank is the tile (nested layout) # and the second rank is the pipeline stage @@ -1220,7 +1220,7 @@ def _dbias_only_colwise( @cute.jit def _dbias_reduction_rowwise_epilouge( - self, smem: cutlass.utils.SmemAllocator, tidx: Int32, rowwise_dbias_acc: cute.Tensor + self, smem: cutlass.memory.SmemAllocator, tidx: Int32, rowwise_dbias_acc: cute.Tensor ): # Pad the buffer to avoid bank conflicts. The logical shape is still the same. Only the stride is different. DBIAS_BUFF_WIDTH = ( diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index 80a90a85ab8..48f3a956544 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -92,6 +92,52 @@ def exp2_bf16x2_rcp(scale_e8m0) -> Uint32: return bf16_bits | (bf16_bits << Uint32(16)) +@cute.jit +def bf16_mx_scale_reciprocal(magnitude: Uint32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: + """Exact BF16 reciprocal scale bits, including E8M0 zero, Inf and NaN cases.""" + rounded = (magnitude + Uint32(31)) & Uint32(0x7F80) + if rounded < Uint32(MAX_EXPONENT << 7): # pylint: disable=consider-using-max-builtin + rounded = Uint32(MAX_EXPONENT << 7) + reciprocal = Uint32((254 + MAX_EXPONENT) << 7) - rounded + if magnitude >= Uint32(0x7F80): + reciprocal = Uint32(0x0040) + if magnitude != Uint32(0x7F80): + reciprocal = Uint32(0x7FFF) + return reciprocal + + +@cute.jit +def bf16_mx_scale_reciprocal_pair(pair: Int32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: + """Compute two reciprocal scales without borrowing across packed BF16 halves.""" + magnitudes = Uint32(pair) & Uint32(0x7FFF7FFF) + rounded = (magnitudes + Uint32(0x001F001F)) & Uint32(0xFF80FF80) + floor = (MAX_EXPONENT << 23) | (MAX_EXPONENT << 7) + # Packed BF16 max accepts registers, so materialize the constant inside PTX. + clamped = cute.arch.inline_ptx( + "{ .reg.b32 floor_pair; " + f"mov.b32 floor_pair, {floor}; " + "max.xorsign.abs.bf16x2 {$w0}, {$r0}, floor_pair; }", + write_only_types=[Int32], + read_only_args=[rounded.bitcast(Int32)], + ) + bias = ((254 + MAX_EXPONENT) << 23) | ((254 + MAX_EXPONENT) << 7) + reciprocal = Uint32(bias) - Uint32(clamped) + if ((magnitudes + Uint32(0x00800080)) & Uint32(0x80008000)) != Uint32(0): + lo = bf16_mx_scale_reciprocal(magnitudes & Uint32(0x7FFF), MAX_EXPONENT) + hi = bf16_mx_scale_reciprocal(magnitudes >> Uint32(16), MAX_EXPONENT) + reciprocal = lo | (hi << Uint32(16)) + return reciprocal + + +@cute.jit +def bf16_mx_scale_bytes_pair(reciprocal: Uint32) -> Uint32: + """Gather the E8M0 bytes of two packed BF16 reciprocals into the low half.""" + bytes_pair = (Uint32(0x01FE01FE) - ((reciprocal >> Uint32(7)) & Uint32(0x00FF00FF))) & Uint32( + 0x00FF00FF + ) + return (bytes_pair & Uint32(0xFF)) | ((bytes_pair >> Uint32(8)) & Uint32(0xFF00)) + + @cute.jit def pack_f32x2(lo: Float32, hi: Float32) -> Int64: """Pack two f32 scalars into a single 64-bit register (`floatx2` layout). @@ -269,6 +315,7 @@ def is_packed16(dtype) -> bool: return dtype is cutlass.BFloat16 or dtype is cutlass.Float16 +# TODO(kainingz): remove this once https://github.com/NVIDIA/cutlass/pull/3717 is merged and released @cute.jit def create_l2_policy(evict_last: cutlass.Constexpr[bool], fraction: Float32 = 1.0) -> Int64: """Create the L2 eviction policy used by register-resident casts.""" From 1995fea3cf5f32115a20a6d2880b89d466d93a0d Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 8 Oct 2026 23:39:53 +0000 Subject: [PATCH 03/29] wip Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 170 +++++++++--------- 1 file changed, 90 insertions(+), 80 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 4a1692df68f..c88dab0d4c4 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -21,7 +21,6 @@ from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 from .quantize_mxfp8_common import ( MXFP8_BLOCK_SCALING_SIZE, - THREADS_PER_WARP, MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, bf16_pair_magnitude, @@ -101,14 +100,13 @@ def quantize_bidimensional_mxfp8_g2r( 1, ), ) + # If CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK == 1, this will help compiler derive alignment + # Otherwise coalesce will not do anything which is still fine + tv_col_scale = cute.coalesce(tv_col_scale, target_profile=(1, 1)) gS_col_thread = cute.composition(gS_col_tile, tv_col_scale)[tidx, None] # Input fragment contains all rows it handles; output fragment contains only one row and we write it back row by row rX = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), cfg.DTYPE) - rO = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.FP8_DTYPE) - - # For rowwise scales, we write one row at a time - rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), @@ -131,7 +129,9 @@ def quantize_bidimensional_mxfp8_g2r( # Copy all rows a thread sweeps into the RMEM fragment at once for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - cute.copy(load_atom, gX_thread[None, row_idx], rX[None, row_idx], cache_policy=cache_policy) + cute.copy( + load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy + ) @cute.jit def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: @@ -143,23 +143,25 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: amax = abs_max_x2_bf16(amax, cute.arch.shuffle_sync_bfly(amax, 1 << stage)) return bf16_pair_magnitude(amax) - # Rowwise pass: produce each row's rowwise scales and quantized outputs, and write them back to GMEM immediately. - rO_u32 = cute.recast_tensor(rO, Uint32) + rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) + rO_row = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.FP8_DTYPE) + rO_row_u32 = cute.recast_tensor(rO_row, Uint32) + # Only one thread from threads that share the MXFP8 block writes the rowwise scale is_writer_thread = tidx % THREADS_X_PER_MX_BLOCK == 0 for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) reciprocal_packed = reciprocal | (reciprocal << Uint32(16)) - for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): - rO_u32[pack_idx] = scale_and_pack( + for pack_idx in cutlass.range_constexpr(cute.size(rO_row_u32)): + rO_row_u32[pack_idx] = scale_and_pack( rX_row_i32[2 * pack_idx], reciprocal_packed, rX_row_i32[2 * pack_idx + 1], reciprocal_packed, ) - # rO_u32 and rO share the same memory - cute.copy(store_atom, rO, gO_row_thread[None, row_idx], cache_policy=cache_policy) + # rO_row_u32 and rO_row share the same memory + cute.copy(store_atom, rO_row, gO_row_thread[(None, row_idx),], cache_policy=cache_policy) if is_writer_thread: rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_packed)).bitcast(SCALE_DTYPE) cute.copy( @@ -169,7 +171,6 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: cache_policy=cache_policy, ) - # Columnwise pass: only produce the colwise scale. # We first reduce each thread's rows into one maximum per column before reduce across different threads in the CTA rAmax_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) @@ -184,91 +185,100 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: # Threads that process a columnwise MXFP8 block don't belong to the same warp so we have to use SMEM to reduce the scale allocator = cutlass.memory.SmemAllocator() + + if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): + scratch_swizzle = cute.make_swizzle(3, 1, 5) + else: + scratch_swizzle = cute.make_swizzle(2, 1, 5) + sAmaxs_col_scratch_layout = cute.zipped_product( rAmax_col.layout, - thread_layout, + # Store fragments in flat thread-ID order while keeping the same shape. + cute.make_layout(thread_layout.shape), + ) + sAmaxs_col_scratch_layout = cute.make_composed_layout( + scratch_swizzle, 0, sAmaxs_col_scratch_layout ) sAmaxs_col_scratch = allocator.allocate_tensor( cfg.DTYPE, sAmaxs_col_scratch_layout, byte_alignment=16 ) - sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) - sS_col_reduction_scratch = allocator.allocate_tensor( - cfg.DTYPE, (CTA_THREADS_Y, CTA_THREADS_X), byte_alignment=16 - ) - sS_col_reduction_scratch_u32 = cute.recast_tensor(sS_col_reduction_scratch, Uint32) - # Every thread copies the RMEM fragment to SMEM scratch space for reduction across warps later - cute.autovec_copy( - rAmax_col, - sAmaxs_col_scratch[None, tidx], + # RMEM buffer for the reciprocal of the columnwise scale to calculate the quantized output + rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) + # SMEM buffer for colwise scale reduction + sS_col_reduction_scratch_layout = cute.zipped_product( + rS_col_rcp.layout, + cute.make_layout(CTA_THREADS_X), + ) + sS_col_reduction_scratch_layout = cute.make_composed_layout( + scratch_swizzle, 0, sS_col_reduction_scratch_layout + ) + sRcp_col_reduction_scratch = allocator.allocate_tensor( + cfg.DTYPE, sS_col_reduction_scratch_layout, byte_alignment=16 ) + sRcp_col_thread = cute.composition(sRcp_col_reduction_scratch, tv_col_scale)[tidx, None] - # Reduce thread read amaxes from SMEM then calculate & write the columnwise scale to SMEM scratch space - is_reduce_thread = tidx < CTA_THREADS_X - if is_reduce_thread: - amaxX2 = sAmaxs_col_scratch_i32[0, tidx] - for scratch_row in cutlass.range_constexpr(CTA_THREADS_Y): - amaxX2 = abs_max_x2_bf16(amaxX2, sAmaxs_col_scratch_i32[scratch_row, tidx]) - sS_col_reduction_scratch_u32[tidx] = bf16_mx_scale_reciprocal(amaxX2, MAX_EXPONENT) + # Every thread writes its partial maxima using its flat thread ID. + cute.autovec_copy(rAmax_col, sAmaxs_col_scratch[None, tidx]) cute.arch.sync_threads() - # Let the reduction threads write the RMEM columnwise scales back to GMEM - if is_reduce_thread: + sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) + + # The layout of sAmaxs_col_scratch_i32 is a zipped product of thread layout (CTA_THREADS_Y, CTA_THREADS_X) + # and each thread's value layout from rAmax_col (ELEMENTS_X_PER_THREAD,) + # And we let each thread here reduce two full column (2 bf16 as a packed i32) of CTA_THREADS_Y elements, + # so we have CTA_THREADS_X * ELEMENT_X_PER_THREAD reduction threads + reduction_threads_num = CTA_THREADS_X * ELEMENTS_X_PER_THREAD // 2 + is_reduction_thread = tidx < reduction_threads_num + + # And apply the colwise amax reduction layout to the SMEM buffer + tv_col_reduce_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((reduction_threads_num,), stride=(1,)), + val_layout=cute.make_layout((CTA_THREADS_Y,), stride=(1,)), + ) + + rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) + if is_reduction_thread: + rAmaxs_col_i32 = cute.composition(sAmaxs_col_scratch_i32, tv_col_reduce_layout)[tidx, None] + # Find the amax for all columns by comparing each row's partial amax + for bf16x2_idx in cutlass.range_constexpr(cute.size(rS_col_rcp_u32)): + amaxX2 = rAmaxs_col_i32[bf16x2_idx] + for scratch_row in cutlass.range_constexpr(1, CTA_THREADS_Y): + peer_thread = thread_layout((scratch_row, tidx)) + amaxX2 = abs_max_x2_bf16(amaxX2, rAmaxs_col_i32[bf16x2_idx, peer_thread]) + # Get the reciprocal directly from the amax + rS_col_rcp_u32[bf16x2_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) + # Write the reciprocal to the SMEM scratch buffer so that all threads can use it to quantize their columns + cute.autovec_copy(rS_col_rcp, sRcp_col_thread) + + rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) rS_col_u32 = cute.recast_tensor(rS_col, Uint32) + # Get the colwise scale from the reciprocal for packed_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): - rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_u32[2 * packed_idx]) | ( - bf16_mx_scale_bytes_pair(rS_col_u32[2 * packed_idx + 1]) << Uint32(16) + rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx]) | ( + bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx + 1]) << Uint32(16) ) + # Use RMEM buffer to stash the columnwise scale in packed bytes and flush it to GMEM using vectorized store at once cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) + cute.arch.sync_threads() - # Every thread reads the reduced columnwise scales from SMEM to RMEM - rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) - cute.autovec_copy( - sS_col_reduction_scratch[tidx], - rS_col - ) - - # And then every thread uses the columnwise scales to calculate quantized outputs and writes it to GMEM - - + # All warps read the same reciprocals for their columns through the overlapping TV layout. + cute.autovec_copy(sRcp_col_thread, rS_col_rcp) + rO_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.FP8_DTYPE) + rO_col_u32 = cute.recast_tensor(rO_col, Uint32) - # # One thread reduces each column pair across all warps (all 32 tile rows). - # if tidx < cute.size(sColumnScales) // 2: - # sColumnAmaxs_i32 = cute.recast_tensor(sColumnAmaxs, Int32) - # sColumnScales_u32 = cute.recast_tensor(sColumnScales, Uint32) - # amax = sColumnAmaxs_i32[0, tidx] - # for warp_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): - # amax = abs_max_x2_bf16(amax, sColumnAmaxs_i32[warp_idx, tidx]) - # sColumnScales_u32[tidx] = bf16_mx_scale_reciprocal_pair(amax, MAX_EXPONENT) - # cute.arch.sync_threads() - - # # The same layout returns the reciprocals to each lane's original columns. - # for half_idx in cutlass.range_constexpr(2): - # cute.autovec_copy( - # cute.local_tile(sColumnScale_thread, (cute.size(rColumnScales) // 2,), (half_idx,)), - # cute.local_tile(rColumnScales, (cute.size(rColumnScales) // 2,), (half_idx,)), - # ) - - # rColumnScales_u32 = cute.recast_tensor(rColumnScales, Uint32) - # if warp == 0: - # rS_col_u32 = cute.recast_tensor(rS_col, Uint32) - # for pack_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): - # rS_col_u32[pack_idx] = bf16_mx_scale_bytes_pair(rColumnScales_u32[2 * pack_idx]) | ( - # bf16_mx_scale_bytes_pair(rColumnScales_u32[2 * pack_idx + 1]) << Uint32(16) - # ) - # cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) - - # for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - # rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - # for pack_idx in cutlass.range_constexpr(cute.size(rO_u32)): - # rO_u32[pack_idx] = scale_and_pack( - # rX_row_i32[2 * pack_idx], - # rColumnScales_u32[2 * pack_idx], - # rX_row_i32[2 * pack_idx + 1], - # rColumnScales_u32[2 * pack_idx + 1], - # ) - # cute.copy(store_atom, rO, gO_col_thread[None, row_idx], cache_policy=cache_policy) + # Each thread quantizes its own rows with the columnwise reciprocals and write its output to GMEM. + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + for packed_idx in cutlass.range_constexpr(cute.size(rO_col_u32)): + rO_col_u32[packed_idx] = scale_and_pack( + rX_row_i32[2 * packed_idx], + rS_col_rcp_u32[2 * packed_idx], + rX_row_i32[2 * packed_idx + 1], + rS_col_rcp_u32[2 * packed_idx + 1], + ) + cute.copy(store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy) class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): From 3e7a653962a1f358eac82a44d55a2b8b5399c93c Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 9 Oct 2026 00:56:06 +0000 Subject: [PATCH 04/29] refactor Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 424 ++++++++++++------ 1 file changed, 275 insertions(+), 149 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index c88dab0d4c4..c056a104979 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -4,7 +4,7 @@ """Register-resident CuTeDSL bidimensional MXFP8 cast, ported from cast_bidim.cuh.""" -from typing import Optional +from typing import Callable, Optional import cutlass from cutlass import cute @@ -28,6 +28,280 @@ ) +@cute.jit +def _quantize_rowwise( + rX: cute.Tensor, + gO_row_thread: cute.Tensor, + gS_row_thread: cute.Tensor, + # Atoms + output_store_atom: cute.CopyAtom, + row_scale_store_atom: cute.CopyAtom, + # Miscellaneous + cache_policy: Int64, + scale_and_pack: Callable[..., Uint32], + # Constexprs + THREADS_X_PER_MX_BLOCK: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + SCALE_DTYPE: cutlass.Constexpr[cutlass.DataType], + FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], +): + tidx, _, _ = cute.arch.thread_idx() + MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 + + @cute.jit + def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: + """Reduce a 32-column MXFP8 block across its two or four lane owners.""" + amax = values_i32[0] + for bf16x2_idx in cutlass.range_constexpr(1, cute.size(values_i32)): + amax = abs_max_x2_bf16(amax, values_i32[bf16x2_idx]) + for stage in cutlass.range_constexpr(THREADS_X_PER_MX_BLOCK.bit_length() - 1): + amax = abs_max_x2_bf16(amax, cute.arch.shuffle_sync_bfly(amax, 1 << stage)) + return bf16_pair_magnitude(amax) + + rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) + rO_row = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), FP8_DTYPE) + rO_row_u32 = cute.recast_tensor(rO_row, Uint32) + + # Only one thread from threads that share the MXFP8 block writes the rowwise scale + is_writer_thread = tidx % THREADS_X_PER_MX_BLOCK == 0 + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) + reciprocal_packed = reciprocal | (reciprocal << Uint32(16)) + for pack_idx in cutlass.range_constexpr(cute.size(rO_row_u32)): + rO_row_u32[pack_idx] = scale_and_pack( + rX_row_i32[2 * pack_idx], + reciprocal_packed, + rX_row_i32[2 * pack_idx + 1], + reciprocal_packed, + ) + # rO_row_u32 and rO_row share the same memory + cute.copy(output_store_atom, rO_row, gO_row_thread[(None, row_idx),], cache_policy=cache_policy) + if is_writer_thread: + rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_packed)).bitcast(SCALE_DTYPE) + cute.copy( + row_scale_store_atom, + rS_row, + cute.local_tile(gS_row_thread, (1,), (row_idx,)), + cache_policy=cache_policy, + ) + +@cute.jit +def _gather_partial_colwise_amaxes( + rX: cute.Tensor, + # Constexprs + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + DTYPE: cutlass.Constexpr[cutlass.DataType], +) -> cute.Tensor: + """ + rX is a RMEM fragment of layout (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) that contains this thread's data + + It returns a SMEM buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)), where each thread + reduces their partial columnwise amax of shape (ELEMENTS_X_PER_THREAD,) and contributes it to the SMEM buffer + """ + tidx, _, _ = cute.arch.thread_idx() + + # We first reduce each thread's rows into one maximum per column before reduce across different threads in the CTA + rAmax_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), DTYPE) + + # Gather the amax for each column that this thread sweeps + rX_i32 = cute.recast_tensor(rX, Int32) + rAmax_col_i32 = cute.recast_tensor(rAmax_col, Int32) + for bf16x2_idx in cutlass.range_constexpr(cute.size(rX_i32, mode=[0])): + amax = rX_i32[bf16x2_idx, 0] + for row_idx in cutlass.range_constexpr(1, ELEMENTS_Y_PER_THREAD): + amax = abs_max_x2_bf16(amax, rX_i32[bf16x2_idx, row_idx]) + rAmax_col_i32[bf16x2_idx] = amax + + if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): + scratch_swizzle = cute.make_swizzle(3, 1, 5) + else: + scratch_swizzle = cute.make_swizzle(2, 1, 5) + + # Prepare amax reduction scratch space in SMEM. + sAmaxs_col_scratch_layout = cute.make_composed_layout( + scratch_swizzle, 0, cute.zipped_product( + rAmax_col.layout, # (ELEMENTS_X_PER_THREAD,) + # Store fragments in flat thread-ID order while keeping the same shape. + cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)) + ) + ) + allocator = cutlass.memory.SmemAllocator() + sAmaxs_col_scratch = allocator.allocate_tensor( + DTYPE, sAmaxs_col_scratch_layout, byte_alignment=16 + ) + + # Every thread writes its partial maxima using its flat thread ID. + cute.autovec_copy(rAmax_col, sAmaxs_col_scratch[None, tidx]) + cute.arch.sync_threads() + + return sAmaxs_col_scratch + +@cute.jit +def _reduce_partial_colwise_amaxes( + rS_col_rcp: cute.Tensor, + sAmaxs_col_scratch: cute.Tensor, + # Constexprs + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + DTYPE: cutlass.Constexpr[cutlass.DataType], + FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], +) -> cute.Tensor: + """ + sAmaxs_col_scratch: the SMEM buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) + that contains each threads' own columnwise amax of its ELEMENTS_X_PER_THREAD columns + rS_col_rcp: the RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) + + This function reduces the partial amax buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) + along the Y axis, to a single columnwise scale SMEM buffer of layout (ELEMENTS_X_PER_THREAD * CTA_THREADS_X), + and write the columnwise scale reciprocals belonging to this thread to its RMEM fragment + """ + tidx, _, _ = cute.arch.thread_idx() + + TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD + + # The layout of sAmaxs_col_scratch_i32 is a zipped product of thread layout (CTA_THREADS_Y, CTA_THREADS_X) + # and each thread's value layout from rAmax_col (ELEMENTS_X_PER_THREAD,) + # And we let each thread here reduce two full column (2 bf16 as a packed i32) of CTA_THREADS_Y elements, + # so we have CTA_THREADS_X * ELEMENT_X_PER_THREAD reduction threads + REDUCTION_THREADS_NUM = CTA_THREADS_X * ELEMENTS_X_PER_THREAD // 2 + is_reduction_thread = tidx < REDUCTION_THREADS_NUM + # How many bf16x2 pairs we need to reduce + REDUCTION_i32_PACKS = TILE_COLS // REDUCTION_THREADS_NUM // 2 + + _, tv_col_reduce_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), + val_layout=cute.make_layout((CTA_THREADS_Y, REDUCTION_i32_PACKS), stride=(REDUCTION_i32_PACKS, 1)), + ) + + # SMEM buffer for colwise scale reduction where mode 0 is how many reciprocals a thread has, + # and mode 1 is how many threads we have in X axis, and the layout size is just TILE_COLS + if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): + scratch_swizzle = cute.make_swizzle(3, 1, 5) + else: + scratch_swizzle = cute.make_swizzle(2, 1, 5) + sS_col_layout = cute.make_composed_layout( + scratch_swizzle, 0, cute.zipped_product( + cute.make_layout((ELEMENTS_X_PER_THREAD,)), + cute.make_layout((CTA_THREADS_X,)), + ) + ) + allocator = cutlass.memory.SmemAllocator() + sS_col = allocator.allocate_tensor( + DTYPE, sS_col_layout, byte_alignment=16 + ) + + # View sS_col using the TV layout from reduction threads. + # REDUCTION_THREADS_NUM * ELEMENTS_X_PER_THREAD = TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD + sS_col_i32 = cute.recast_tensor(sS_col, Int32) + sS_col_i32_reduction_view = cute.composition(sS_col_i32, cute.make_tv_layout( + thr_layout=cute.make_layout((1, REDUCTION_THREADS_NUM), stride=(1, 1)), + val_layout=cute.make_layout((1, ELEMENTS_X_PER_THREAD,), stride=(1, 1)), + )) + + # Each reduction process `REDUCTION_i32_PACKS` bf16x2 packs by visiting all rows (CTA_THREADS_Y) + # to extract the amax of that column + sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) + MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 + if is_reduction_thread: + sRcp_col_i32_reduction_frag = cute.composition(sS_col_i32, sS_col_i32_reduction_view)[tidx, None] + sAmaxs_col_frag_i32 = cute.composition(sAmaxs_col_scratch_i32, tv_col_reduce_layout)[tidx, None] + for pack_idx in cutlass.range_constexpr(REDUCTION_i32_PACKS): + sAmaxs_col_frag_i32_rows = sAmaxs_col_frag_i32[None, pack_idx] + amaxX2 = sAmaxs_col_frag_i32_rows[0] + for scratch_row_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): + amaxX2 = abs_max_x2_bf16(amaxX2, sAmaxs_col_frag_i32_rows[scratch_row_idx]) + sRcp_col_i32_reduction_frag[pack_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) + cute.arch.sync_threads() + + cute.autovec_copy(sS_col[None, tidx], rS_col_rcp) + + +@cute.jit +def _write_colwise_scale( + rS_col_rcp: cute.Tensor, + gS_col_thread: cute.Tensor, + # Atoms + col_scale_atom: cute.CopyAtom, + # Miscellaneous + cache_policy: Int64, + # Constexprs + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + DTYPE: cutlass.Constexpr[cutlass.DataType], + FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], + SCALE_DTYPE: cutlass.Constexpr[cutlass.DataType], +): + """ + rS_col_rcp: RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) + """ + tidx, _, _ = cute.arch.thread_idx() + + rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), DTYPE) + rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) + + # Only threads from CTA's the first row writes the scale fragments, retaining vectorized GMEM stores. + if tidx < CTA_THREADS_X: + rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) + rS_col_u32 = cute.recast_tensor(rS_col, Uint32) + # Get the colwise scale from the reciprocal + for packed_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): + rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx]) | ( + bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx + 1]) << Uint32(16) + ) + # Use RMEM buffer to stash the columnwise scale in packed bytes and flush it to GMEM using vectorized store at once + cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) + + +@cute.jit +def _quantize_colwise( + rX: cute.Tensor, + rS_col_rcp: cute.Tensor, + gO_col_thread: cute.Tensor, + # Atoms + store_atom: cute.CopyAtom, + # Miscellaneous + cache_policy: Int64, + scale_and_pack: Callable[..., Uint32], + # Constexprs + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], +): + """ + rX: RMEM fragment of this thread's rows with shape (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) + rS_col_rcp: RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) + gO_col_thread: GMEM fragment of this thread's columnwise quantized output with shape + (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) + + This function quantizes its fragments using the input and columnwise scales, and write to its output fragment + """ + rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) + + rO_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), FP8_DTYPE) + rO_col_u32 = cute.recast_tensor(rO_col, Uint32) + + # Each thread quantizes its own rows with the columnwise reciprocals and write its output to GMEM. + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + for packed_idx in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): # packed in 2 + rO_col_u32[packed_idx] = scale_and_pack( + rX_row_i32[2 * packed_idx], + rS_col_rcp_u32[2 * packed_idx], + rX_row_i32[2 * packed_idx + 1], + rS_col_rcp_u32[2 * packed_idx + 1], + ) + cute.copy(store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy) + + @cute.jit def quantize_bidimensional_mxfp8_g2r( gX_tile: cute.Tensor, @@ -51,13 +325,11 @@ def quantize_bidimensional_mxfp8_g2r( """ # Shape of the tile TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD - TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD # How many threads cooperate to process a single MXFP8 block (32 elements) for both quantized directions THREADS_X_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_X_PER_THREAD THREADS_Y_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_Y_PER_THREAD - MAX_EXPONENT = 8 if cfg.FP8_DTYPE is cutlass.Float8E4M3FN else 15 SCALE_DTYPE = cutlass.Float8E8M0FNU tidx, _, _ = cute.arch.thread_idx() @@ -133,152 +405,6 @@ def quantize_bidimensional_mxfp8_g2r( load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy ) - @cute.jit - def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: - """Reduce a 32-column MXFP8 block across its two or four lane owners.""" - amax = values_i32[0] - for bf16x2_idx in cutlass.range_constexpr(1, cute.size(values_i32)): - amax = abs_max_x2_bf16(amax, values_i32[bf16x2_idx]) - for stage in cutlass.range_constexpr(THREADS_X_PER_MX_BLOCK.bit_length() - 1): - amax = abs_max_x2_bf16(amax, cute.arch.shuffle_sync_bfly(amax, 1 << stage)) - return bf16_pair_magnitude(amax) - - rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) - rO_row = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.FP8_DTYPE) - rO_row_u32 = cute.recast_tensor(rO_row, Uint32) - - # Only one thread from threads that share the MXFP8 block writes the rowwise scale - is_writer_thread = tidx % THREADS_X_PER_MX_BLOCK == 0 - for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) - reciprocal_packed = reciprocal | (reciprocal << Uint32(16)) - for pack_idx in cutlass.range_constexpr(cute.size(rO_row_u32)): - rO_row_u32[pack_idx] = scale_and_pack( - rX_row_i32[2 * pack_idx], - reciprocal_packed, - rX_row_i32[2 * pack_idx + 1], - reciprocal_packed, - ) - # rO_row_u32 and rO_row share the same memory - cute.copy(store_atom, rO_row, gO_row_thread[(None, row_idx),], cache_policy=cache_policy) - if is_writer_thread: - rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_packed)).bitcast(SCALE_DTYPE) - cute.copy( - row_scale_atom, - rS_row, - cute.local_tile(gS_row_thread, (1,), (row_idx,)), - cache_policy=cache_policy, - ) - - # We first reduce each thread's rows into one maximum per column before reduce across different threads in the CTA - rAmax_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) - - # Gather the amax for each column that this thread sweeps - rX_i32 = cute.recast_tensor(rX, Int32) - rAmax_col_i32 = cute.recast_tensor(rAmax_col, Int32) - for bf16x2_idx in cutlass.range_constexpr(cute.size(rX_i32, mode=[0])): - amax = rX_i32[bf16x2_idx, 0] - for row_idx in cutlass.range_constexpr(1, ELEMENTS_Y_PER_THREAD): - amax = abs_max_x2_bf16(amax, rX_i32[bf16x2_idx, row_idx]) - rAmax_col_i32[bf16x2_idx] = amax - - # Threads that process a columnwise MXFP8 block don't belong to the same warp so we have to use SMEM to reduce the scale - allocator = cutlass.memory.SmemAllocator() - - if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): - scratch_swizzle = cute.make_swizzle(3, 1, 5) - else: - scratch_swizzle = cute.make_swizzle(2, 1, 5) - - sAmaxs_col_scratch_layout = cute.zipped_product( - rAmax_col.layout, - # Store fragments in flat thread-ID order while keeping the same shape. - cute.make_layout(thread_layout.shape), - ) - sAmaxs_col_scratch_layout = cute.make_composed_layout( - scratch_swizzle, 0, sAmaxs_col_scratch_layout - ) - sAmaxs_col_scratch = allocator.allocate_tensor( - cfg.DTYPE, sAmaxs_col_scratch_layout, byte_alignment=16 - ) - - # RMEM buffer for the reciprocal of the columnwise scale to calculate the quantized output - rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) - # SMEM buffer for colwise scale reduction - sS_col_reduction_scratch_layout = cute.zipped_product( - rS_col_rcp.layout, - cute.make_layout(CTA_THREADS_X), - ) - sS_col_reduction_scratch_layout = cute.make_composed_layout( - scratch_swizzle, 0, sS_col_reduction_scratch_layout - ) - sRcp_col_reduction_scratch = allocator.allocate_tensor( - cfg.DTYPE, sS_col_reduction_scratch_layout, byte_alignment=16 - ) - sRcp_col_thread = cute.composition(sRcp_col_reduction_scratch, tv_col_scale)[tidx, None] - - # Every thread writes its partial maxima using its flat thread ID. - cute.autovec_copy(rAmax_col, sAmaxs_col_scratch[None, tidx]) - cute.arch.sync_threads() - - sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) - - # The layout of sAmaxs_col_scratch_i32 is a zipped product of thread layout (CTA_THREADS_Y, CTA_THREADS_X) - # and each thread's value layout from rAmax_col (ELEMENTS_X_PER_THREAD,) - # And we let each thread here reduce two full column (2 bf16 as a packed i32) of CTA_THREADS_Y elements, - # so we have CTA_THREADS_X * ELEMENT_X_PER_THREAD reduction threads - reduction_threads_num = CTA_THREADS_X * ELEMENTS_X_PER_THREAD // 2 - is_reduction_thread = tidx < reduction_threads_num - - # And apply the colwise amax reduction layout to the SMEM buffer - tv_col_reduce_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((reduction_threads_num,), stride=(1,)), - val_layout=cute.make_layout((CTA_THREADS_Y,), stride=(1,)), - ) - - rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) - if is_reduction_thread: - rAmaxs_col_i32 = cute.composition(sAmaxs_col_scratch_i32, tv_col_reduce_layout)[tidx, None] - # Find the amax for all columns by comparing each row's partial amax - for bf16x2_idx in cutlass.range_constexpr(cute.size(rS_col_rcp_u32)): - amaxX2 = rAmaxs_col_i32[bf16x2_idx] - for scratch_row in cutlass.range_constexpr(1, CTA_THREADS_Y): - peer_thread = thread_layout((scratch_row, tidx)) - amaxX2 = abs_max_x2_bf16(amaxX2, rAmaxs_col_i32[bf16x2_idx, peer_thread]) - # Get the reciprocal directly from the amax - rS_col_rcp_u32[bf16x2_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) - # Write the reciprocal to the SMEM scratch buffer so that all threads can use it to quantize their columns - cute.autovec_copy(rS_col_rcp, sRcp_col_thread) - - rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) - rS_col_u32 = cute.recast_tensor(rS_col, Uint32) - # Get the colwise scale from the reciprocal - for packed_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): - rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx]) | ( - bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx + 1]) << Uint32(16) - ) - # Use RMEM buffer to stash the columnwise scale in packed bytes and flush it to GMEM using vectorized store at once - cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) - cute.arch.sync_threads() - - # All warps read the same reciprocals for their columns through the overlapping TV layout. - cute.autovec_copy(sRcp_col_thread, rS_col_rcp) - - rO_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.FP8_DTYPE) - rO_col_u32 = cute.recast_tensor(rO_col, Uint32) - - # Each thread quantizes its own rows with the columnwise reciprocals and write its output to GMEM. - for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - for packed_idx in cutlass.range_constexpr(cute.size(rO_col_u32)): - rO_col_u32[packed_idx] = scale_and_pack( - rX_row_i32[2 * packed_idx], - rS_col_rcp_u32[2 * packed_idx], - rX_row_i32[2 * packed_idx + 1], - rS_col_rcp_u32[2 * packed_idx + 1], - ) - cute.copy(store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy) class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): From 79c5d5325104ccc209a997a0207a70edb4e45c38 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 9 Oct 2026 01:28:20 +0000 Subject: [PATCH 05/29] refactor Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 197 +++++++++++++----- 1 file changed, 142 insertions(+), 55 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index c056a104979..9db22913ddd 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -4,7 +4,7 @@ """Register-resident CuTeDSL bidimensional MXFP8 cast, ported from cast_bidim.cuh.""" -from typing import Callable, Optional +from typing import Callable, Optional, Type import cutlass from cutlass import cute @@ -38,13 +38,13 @@ def _quantize_rowwise( row_scale_store_atom: cute.CopyAtom, # Miscellaneous cache_policy: Int64, - scale_and_pack: Callable[..., Uint32], + scale_and_pack: cutlass.Constexpr[Callable[..., Uint32]], # Constexprs THREADS_X_PER_MX_BLOCK: cutlass.Constexpr[int], ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - SCALE_DTYPE: cutlass.Constexpr[cutlass.DataType], - FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], + SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): tidx, _, _ = cute.arch.thread_idx() MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 @@ -77,7 +77,9 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: reciprocal_packed, ) # rO_row_u32 and rO_row share the same memory - cute.copy(output_store_atom, rO_row, gO_row_thread[(None, row_idx),], cache_policy=cache_policy) + cute.copy( + output_store_atom, rO_row, gO_row_thread[(None, row_idx),], cache_policy=cache_policy + ) if is_writer_thread: rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_packed)).bitcast(SCALE_DTYPE) cute.copy( @@ -87,6 +89,7 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: cache_policy=cache_policy, ) + @cute.jit def _gather_partial_colwise_amaxes( rX: cute.Tensor, @@ -95,7 +98,7 @@ def _gather_partial_colwise_amaxes( CTA_THREADS_X: cutlass.Constexpr[int], ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - DTYPE: cutlass.Constexpr[cutlass.DataType], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ) -> cute.Tensor: """ rX is a RMEM fragment of layout (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) that contains this thread's data @@ -124,11 +127,12 @@ def _gather_partial_colwise_amaxes( # Prepare amax reduction scratch space in SMEM. sAmaxs_col_scratch_layout = cute.make_composed_layout( - scratch_swizzle, 0, cute.zipped_product( - rAmax_col.layout, # (ELEMENTS_X_PER_THREAD,) - # Store fragments in flat thread-ID order while keeping the same shape. - cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)) - ) + scratch_swizzle, + 0, + cute.zipped_product( + rAmax_col.layout, # Each thread contributes (ELEMENTS_X_PER_THREAD,) + cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X)), # Total threads in the CTA + ), ) allocator = cutlass.memory.SmemAllocator() sAmaxs_col_scratch = allocator.allocate_tensor( @@ -141,86 +145,103 @@ def _gather_partial_colwise_amaxes( return sAmaxs_col_scratch + @cute.jit def _reduce_partial_colwise_amaxes( - rS_col_rcp: cute.Tensor, sAmaxs_col_scratch: cute.Tensor, # Constexprs CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - DTYPE: cutlass.Constexpr[cutlass.DataType], - FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], -) -> cute.Tensor: + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], +) -> None: """ - sAmaxs_col_scratch: the SMEM buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) + sAmaxs_col_scratch: the SMEM buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) that contains each threads' own columnwise amax of its ELEMENTS_X_PER_THREAD columns rS_col_rcp: the RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) - This function reduces the partial amax buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) + This function reduces the partial amax buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) along the Y axis, to a single columnwise scale SMEM buffer of layout (ELEMENTS_X_PER_THREAD * CTA_THREADS_X), - and write the columnwise scale reciprocals belonging to this thread to its RMEM fragment + which is the return value """ tidx, _, _ = cute.arch.thread_idx() TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD - # The layout of sAmaxs_col_scratch_i32 is a zipped product of thread layout (CTA_THREADS_Y, CTA_THREADS_X) - # and each thread's value layout from rAmax_col (ELEMENTS_X_PER_THREAD,) - # And we let each thread here reduce two full column (2 bf16 as a packed i32) of CTA_THREADS_Y elements, - # so we have CTA_THREADS_X * ELEMENT_X_PER_THREAD reduction threads + # Each reducer owns one column pair and visits its partial maximum from + # every warp. Column pairs precede warp rows in the scratch's logical order. REDUCTION_THREADS_NUM = CTA_THREADS_X * ELEMENTS_X_PER_THREAD // 2 is_reduction_thread = tidx < REDUCTION_THREADS_NUM # How many bf16x2 pairs we need to reduce REDUCTION_i32_PACKS = TILE_COLS // REDUCTION_THREADS_NUM // 2 - _, tv_col_reduce_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), - val_layout=cute.make_layout((CTA_THREADS_Y, REDUCTION_i32_PACKS), stride=(REDUCTION_i32_PACKS, 1)), - ) - - # SMEM buffer for colwise scale reduction where mode 0 is how many reciprocals a thread has, - # and mode 1 is how many threads we have in X axis, and the layout size is just TILE_COLS + # Allocate colwise scale SMEM buffer with swizzling if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): scratch_swizzle = cute.make_swizzle(3, 1, 5) else: scratch_swizzle = cute.make_swizzle(2, 1, 5) + # SMEM buffer for colwise scale reduction where mode 0 is how many reciprocals a thread has, + # and mode 1 is how many threads we have in X axis, and the layout size is just TILE_COLS sS_col_layout = cute.make_composed_layout( - scratch_swizzle, 0, cute.zipped_product( + scratch_swizzle, + 0, + cute.zipped_product( cute.make_layout((ELEMENTS_X_PER_THREAD,)), cute.make_layout((CTA_THREADS_X,)), - ) + ), ) allocator = cutlass.memory.SmemAllocator() - sS_col = allocator.allocate_tensor( - DTYPE, sS_col_layout, byte_alignment=16 - ) - - # View sS_col using the TV layout from reduction threads. - # REDUCTION_THREADS_NUM * ELEMENTS_X_PER_THREAD = TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD + sS_col = allocator.allocate_tensor(DTYPE, sS_col_layout, byte_alignment=16) sS_col_i32 = cute.recast_tensor(sS_col, Int32) - sS_col_i32_reduction_view = cute.composition(sS_col_i32, cute.make_tv_layout( - thr_layout=cute.make_layout((1, REDUCTION_THREADS_NUM), stride=(1, 1)), - val_layout=cute.make_layout((1, ELEMENTS_X_PER_THREAD,), stride=(1, 1)), - )) + + # Colwise amaxes layout: each reduction thread owns CTA_THREADS_Y rows, with REDUCTION_i32_PACKS column packs + _, tv_col_amaxes_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), + val_layout=cute.make_layout( + (REDUCTION_i32_PACKS, CTA_THREADS_Y), stride=(1, REDUCTION_i32_PACKS) + ), + ) + # Colwise rcp layout: each reduction thread owns only 1 row, with REDUCTION_i32_PACKS column packs + _, tv_col_rcp_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), + val_layout=cute.make_layout((REDUCTION_i32_PACKS, 1), stride=(1, REDUCTION_i32_PACKS)), + ) + + sS_col_i32 = cute.composition(sS_col_i32, tv_col_rcp_layout) # Each reduction process `REDUCTION_i32_PACKS` bf16x2 packs by visiting all rows (CTA_THREADS_Y) # to extract the amax of that column sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 if is_reduction_thread: - sRcp_col_i32_reduction_frag = cute.composition(sS_col_i32, sS_col_i32_reduction_view)[tidx, None] - sAmaxs_col_frag_i32 = cute.composition(sAmaxs_col_scratch_i32, tv_col_reduce_layout)[tidx, None] + sAmaxs_col_frag_i32 = cute.composition(sAmaxs_col_scratch_i32, tv_col_amaxes_layout)[ + tidx, None + ] + # Expose the thread's value mode as (partial rows, column pairs). + sAmaxs_col_frag_i32 = cute.composition( + sAmaxs_col_frag_i32, + cute.make_layout((CTA_THREADS_Y, REDUCTION_i32_PACKS), stride=(REDUCTION_i32_PACKS, 1)), + ) + + # Iterative over all packed columns for pack_idx in cutlass.range_constexpr(REDUCTION_i32_PACKS): + # Select the packed column sAmaxs_col_frag_i32_rows = sAmaxs_col_frag_i32[None, pack_idx] + + # Reduce across all rows' partial amaxes to obtain the columnwise amax for this column pack amaxX2 = sAmaxs_col_frag_i32_rows[0] for scratch_row_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): amaxX2 = abs_max_x2_bf16(amaxX2, sAmaxs_col_frag_i32_rows[scratch_row_idx]) - sRcp_col_i32_reduction_frag[pack_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) + + # Obtain the columnwise scale reciprocal for this column pack + sS_col_i32[tidx, pack_idx] = bf16_mx_scale_reciprocal_pair( + amaxX2, MAX_EXPONENT + ) cute.arch.sync_threads() - cute.autovec_copy(sS_col[None, tidx], rS_col_rcp) + return sS_col @cute.jit @@ -236,19 +257,18 @@ def _write_colwise_scale( CTA_THREADS_X: cutlass.Constexpr[int], ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - DTYPE: cutlass.Constexpr[cutlass.DataType], - FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], - SCALE_DTYPE: cutlass.Constexpr[cutlass.DataType], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): """ rS_col_rcp: RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) """ tidx, _, _ = cute.arch.thread_idx() - rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), DTYPE) rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) - # Only threads from CTA's the first row writes the scale fragments, retaining vectorized GMEM stores. + # Only threads from CTA's the first row writes the scale fragments if tidx < CTA_THREADS_X: rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) rS_col_u32 = cute.recast_tensor(rS_col, Uint32) @@ -257,7 +277,8 @@ def _write_colwise_scale( rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx]) | ( bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx + 1]) << Uint32(16) ) - # Use RMEM buffer to stash the columnwise scale in packed bytes and flush it to GMEM using vectorized store at once + # Use RMEM buffer to stash the columnwise scale in packed bytes and flush it to GMEM + # using vectorized store at once cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) @@ -270,18 +291,18 @@ def _quantize_colwise( store_atom: cute.CopyAtom, # Miscellaneous cache_policy: Int64, - scale_and_pack: Callable[..., Uint32], + scale_and_pack: cutlass.Constexpr[Callable[..., Uint32]], # Constexprs ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - FP8_DTYPE: cutlass.Constexpr[cutlass.DataType], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): """ rX: RMEM fragment of this thread's rows with shape (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) rS_col_rcp: RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) gO_col_thread: GMEM fragment of this thread's columnwise quantized output with shape (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) - + This function quantizes its fragments using the input and columnwise scales, and write to its output fragment """ rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) @@ -292,7 +313,7 @@ def _quantize_colwise( # Each thread quantizes its own rows with the columnwise reciprocals and write its output to GMEM. for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - for packed_idx in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): # packed in 2 + for packed_idx in cutlass.range_constexpr(cute.size(rO_col_u32)): # 4 values in one pack rO_col_u32[packed_idx] = scale_and_pack( rX_row_i32[2 * packed_idx], rS_col_rcp_u32[2 * packed_idx], @@ -405,6 +426,71 @@ def quantize_bidimensional_mxfp8_g2r( load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy ) + _quantize_rowwise( + rX, + gO_row_thread, + gS_row_thread, + output_store_atom=store_atom, + row_scale_store_atom=row_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + ) + + sAmaxs_col_scratch = _gather_partial_colwise_amaxes( + rX, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + DTYPE=cfg.DTYPE, + ) + + sS_col = _reduce_partial_colwise_amaxes( + rS_col_rcp, + sAmaxs_col_scratch, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + DTYPE=cfg.DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + ) + + # Copy the reduced columnwise scale reciprocals to RMEM. + # Threads that quantize different rows share the same columnwise scale + rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) + cute.autovec_copy(sS_col[None, tidx % CTA_THREADS_X], rS_col_rcp) + + _write_colwise_scale( + rS_col_rcp, + gS_col_thread, + col_scale_atom=col_scale_atom, + cache_policy=cache_policy, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + DTYPE=cfg.DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + SCALE_DTYPE=SCALE_DTYPE, + ) + + _quantize_colwise( + rX, + rS_col_rcp, + gO_col_thread, + store_atom=store_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + FP8_DTYPE=cfg.FP8_DTYPE, + ) class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): @@ -422,6 +508,7 @@ class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg + # One CTA only quantize one columnwise MXFP8 block at a time self._ELEMENTS_Y_PER_THREAD = MXFP8_BLOCK_SCALING_SIZE // self._CTA_THREADS_Y # 4 @cute.jit From c51a3b8e8b408db1440336905a0d2468a5a63d65 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 9 Oct 2026 04:50:18 +0000 Subject: [PATCH 06/29] refactor Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 84 +++++++++---------- 1 file changed, 38 insertions(+), 46 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 9db22913ddd..5e2d1a6a8aa 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -68,20 +68,20 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) - reciprocal_packed = reciprocal | (reciprocal << Uint32(16)) - for pack_idx in cutlass.range_constexpr(cute.size(rO_row_u32)): + reciprocal_pair = reciprocal | (reciprocal << Uint32(16)) + for pack_idx in cutlass.range_constexpr(cute.size(rO_row_u32)): # 4 output bytes in one pack rO_row_u32[pack_idx] = scale_and_pack( rX_row_i32[2 * pack_idx], - reciprocal_packed, + reciprocal_pair, rX_row_i32[2 * pack_idx + 1], - reciprocal_packed, + reciprocal_pair, ) # rO_row_u32 and rO_row share the same memory cute.copy( output_store_atom, rO_row, gO_row_thread[(None, row_idx),], cache_policy=cache_policy ) if is_writer_thread: - rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_packed)).bitcast(SCALE_DTYPE) + rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_pair)).bitcast(SCALE_DTYPE) cute.copy( row_scale_store_atom, rS_row, @@ -120,10 +120,9 @@ def _gather_partial_colwise_amaxes( amax = abs_max_x2_bf16(amax, rX_i32[bf16x2_idx, row_idx]) rAmax_col_i32[bf16x2_idx] = amax - if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): - scratch_swizzle = cute.make_swizzle(3, 1, 5) - else: - scratch_swizzle = cute.make_swizzle(2, 1, 5) + # Keep eight BF16 values (16 bytes) contiguous. XOR BF16 offset bit 6 into + # bit 3 so each 128-bit warp transaction uses all 32 banks without overlap. + scratch_swizzle = cute.make_swizzle(1, 3, 3) # Prepare amax reduction scratch space in SMEM. sAmaxs_col_scratch_layout = cute.make_composed_layout( @@ -131,7 +130,7 @@ def _gather_partial_colwise_amaxes( 0, cute.zipped_product( rAmax_col.layout, # Each thread contributes (ELEMENTS_X_PER_THREAD,) - cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X)), # Total threads in the CTA + cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X)), # Total threads in the CTA ), ) allocator = cutlass.memory.SmemAllocator() @@ -156,11 +155,10 @@ def _reduce_partial_colwise_amaxes( ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], -) -> None: +) -> cute.Tensor: """ sAmaxs_col_scratch: the SMEM buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) that contains each threads' own columnwise amax of its ELEMENTS_X_PER_THREAD columns - rS_col_rcp: the RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) This function reduces the partial amax buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) along the Y axis, to a single columnwise scale SMEM buffer of layout (ELEMENTS_X_PER_THREAD * CTA_THREADS_X), @@ -175,13 +173,11 @@ def _reduce_partial_colwise_amaxes( REDUCTION_THREADS_NUM = CTA_THREADS_X * ELEMENTS_X_PER_THREAD // 2 is_reduction_thread = tidx < REDUCTION_THREADS_NUM # How many bf16x2 pairs we need to reduce - REDUCTION_i32_PACKS = TILE_COLS // REDUCTION_THREADS_NUM // 2 + REDUCTION_i32_PAIRS = TILE_COLS // REDUCTION_THREADS_NUM // 2 - # Allocate colwise scale SMEM buffer with swizzling - if cutlass.const_expr(ELEMENTS_X_PER_THREAD == 16): - scratch_swizzle = cute.make_swizzle(3, 1, 5) - else: - scratch_swizzle = cute.make_swizzle(2, 1, 5) + # Use the same 16-byte swizzle as the partial-max scratch. Consecutive + # column-pair reducers still access distinct banks through this layout. + scratch_swizzle = cute.make_swizzle(1, 3, 3) # SMEM buffer for colwise scale reduction where mode 0 is how many reciprocals a thread has, # and mode 1 is how many threads we have in X axis, and the layout size is just TILE_COLS sS_col_layout = cute.make_composed_layout( @@ -196,22 +192,22 @@ def _reduce_partial_colwise_amaxes( sS_col = allocator.allocate_tensor(DTYPE, sS_col_layout, byte_alignment=16) sS_col_i32 = cute.recast_tensor(sS_col, Int32) - # Colwise amaxes layout: each reduction thread owns CTA_THREADS_Y rows, with REDUCTION_i32_PACKS column packs + # Colwise amaxes layout: each reduction thread owns CTA_THREADS_Y rows, with REDUCTION_i32_PAIRS column pairs _, tv_col_amaxes_layout = cute.make_layout_tv( thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), val_layout=cute.make_layout( - (REDUCTION_i32_PACKS, CTA_THREADS_Y), stride=(1, REDUCTION_i32_PACKS) + (REDUCTION_i32_PAIRS, CTA_THREADS_Y), stride=(1, REDUCTION_i32_PAIRS) ), ) - # Colwise rcp layout: each reduction thread owns only 1 row, with REDUCTION_i32_PACKS column packs + # Colwise rcp layout: each reduction thread owns only 1 row, with REDUCTION_i32_PAIRS column pairs _, tv_col_rcp_layout = cute.make_layout_tv( thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), - val_layout=cute.make_layout((REDUCTION_i32_PACKS, 1), stride=(1, REDUCTION_i32_PACKS)), + val_layout=cute.make_layout((REDUCTION_i32_PAIRS, 1), stride=(1, REDUCTION_i32_PAIRS)), ) - + sS_col_i32 = cute.composition(sS_col_i32, tv_col_rcp_layout) - # Each reduction process `REDUCTION_i32_PACKS` bf16x2 packs by visiting all rows (CTA_THREADS_Y) + # Each reduction process `REDUCTION_i32_PAIRS` bf16x2 pairs by visiting all rows (CTA_THREADS_Y) # to extract the amax of that column sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 @@ -222,23 +218,21 @@ def _reduce_partial_colwise_amaxes( # Expose the thread's value mode as (partial rows, column pairs). sAmaxs_col_frag_i32 = cute.composition( sAmaxs_col_frag_i32, - cute.make_layout((CTA_THREADS_Y, REDUCTION_i32_PACKS), stride=(REDUCTION_i32_PACKS, 1)), + cute.make_layout((CTA_THREADS_Y, REDUCTION_i32_PAIRS), stride=(REDUCTION_i32_PAIRS, 1)), ) - # Iterative over all packed columns - for pack_idx in cutlass.range_constexpr(REDUCTION_i32_PACKS): - # Select the packed column - sAmaxs_col_frag_i32_rows = sAmaxs_col_frag_i32[None, pack_idx] + # Iterative over all column pairs + for pair_idx in cutlass.range_constexpr(REDUCTION_i32_PAIRS): + # Select the pair column + sAmaxs_col_frag_i32_rows = sAmaxs_col_frag_i32[None, pair_idx] - # Reduce across all rows' partial amaxes to obtain the columnwise amax for this column pack + # Reduce across all rows' partial amaxes to obtain the columnwise amax for this column pair amaxX2 = sAmaxs_col_frag_i32_rows[0] for scratch_row_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): amaxX2 = abs_max_x2_bf16(amaxX2, sAmaxs_col_frag_i32_rows[scratch_row_idx]) - # Obtain the columnwise scale reciprocal for this column pack - sS_col_i32[tidx, pack_idx] = bf16_mx_scale_reciprocal_pair( - amaxX2, MAX_EXPONENT - ) + # Obtain the columnwise scale reciprocal for this column pair + sS_col_i32[tidx, pair_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) cute.arch.sync_threads() return sS_col @@ -273,12 +267,11 @@ def _write_colwise_scale( rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) rS_col_u32 = cute.recast_tensor(rS_col, Uint32) # Get the colwise scale from the reciprocal - for packed_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): - rS_col_u32[packed_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx]) | ( - bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * packed_idx + 1]) << Uint32(16) + for pack_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): # 4 output bytes in one pack + rS_col_u32[pack_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * pack_idx]) | ( + bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * pack_idx + 1]) << Uint32(16) ) - # Use RMEM buffer to stash the columnwise scale in packed bytes and flush it to GMEM - # using vectorized store at once + # Use RMEM buffer to stash the columnwise scale and flush it to GMEM using vectorized store at once cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) @@ -313,12 +306,12 @@ def _quantize_colwise( # Each thread quantizes its own rows with the columnwise reciprocals and write its output to GMEM. for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - for packed_idx in cutlass.range_constexpr(cute.size(rO_col_u32)): # 4 values in one pack - rO_col_u32[packed_idx] = scale_and_pack( - rX_row_i32[2 * packed_idx], - rS_col_rcp_u32[2 * packed_idx], - rX_row_i32[2 * packed_idx + 1], - rS_col_rcp_u32[2 * packed_idx + 1], + for pack_idx in cutlass.range_constexpr(cute.size(rO_col_u32)): # 4 output bytes in one pack + rO_col_u32[pack_idx] = scale_and_pack( + rX_row_i32[2 * pack_idx], + rS_col_rcp_u32[2 * pack_idx], + rX_row_i32[2 * pack_idx + 1], + rS_col_rcp_u32[2 * pack_idx + 1], ) cute.copy(store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy) @@ -451,7 +444,6 @@ def quantize_bidimensional_mxfp8_g2r( ) sS_col = _reduce_partial_colwise_amaxes( - rS_col_rcp, sAmaxs_col_scratch, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, From f14e433781ba19ca84105e023fdc796ac12dc4fd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 04:52:08 +0000 Subject: [PATCH 07/29] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 5e2d1a6a8aa..886b9eb27bd 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -69,7 +69,9 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) reciprocal_pair = reciprocal | (reciprocal << Uint32(16)) - for pack_idx in cutlass.range_constexpr(cute.size(rO_row_u32)): # 4 output bytes in one pack + for pack_idx in cutlass.range_constexpr( + cute.size(rO_row_u32) + ): # 4 output bytes in one pack rO_row_u32[pack_idx] = scale_and_pack( rX_row_i32[2 * pack_idx], reciprocal_pair, @@ -267,7 +269,9 @@ def _write_colwise_scale( rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) rS_col_u32 = cute.recast_tensor(rS_col, Uint32) # Get the colwise scale from the reciprocal - for pack_idx in cutlass.range_constexpr(cute.size(rS_col_u32)): # 4 output bytes in one pack + for pack_idx in cutlass.range_constexpr( + cute.size(rS_col_u32) + ): # 4 output bytes in one pack rS_col_u32[pack_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * pack_idx]) | ( bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * pack_idx + 1]) << Uint32(16) ) @@ -306,7 +310,9 @@ def _quantize_colwise( # Each thread quantizes its own rows with the columnwise reciprocals and write its output to GMEM. for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - for pack_idx in cutlass.range_constexpr(cute.size(rO_col_u32)): # 4 output bytes in one pack + for pack_idx in cutlass.range_constexpr( + cute.size(rO_col_u32) + ): # 4 output bytes in one pack rO_col_u32[pack_idx] = scale_and_pack( rX_row_i32[2 * pack_idx], rS_col_rcp_u32[2 * pack_idx], From 5f9ece21f3bc3241bd1817493a489a96ccd07133 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 9 Oct 2026 13:57:11 -0700 Subject: [PATCH 08/29] fix Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_common.py | 12 +- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 112 +++++++++++------- 2 files changed, 72 insertions(+), 52 deletions(-) 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 249e178abd6..14d3b88ca66 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -329,15 +329,9 @@ def __call__( scale_stride_matches = Boolean(True) if cutlass.const_expr(not self.cfg.WITH_GEMM_SWIZZLED_SCALES): scale_stride_matches = mS_row.shape[1] == N // 32 - # Column tiles use CUDA's grid.y axis. Inputs this wide always use two tiles. - grid_cols_fit = ( - N - <= 65535 - * self.g2r_rowwise_2lane._CTA_THREADS_X - * self.g2r_rowwise_2lane._ELEMENTS_PER_THREAD - * 2 - ) - if N % 128 == 0 and input_is_aligned and scale_stride_matches and grid_cols_fit: + # Column tiles use grid.x and row tiles spill from grid.y into grid.z, so + # no shape exceeds the grid limits. + if N % 128 == 0 and input_is_aligned and scale_stride_matches: dispatched_to_specialized = True self.g2r_rowwise_2lane( mX, 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 edc6cf850f5..a264b578575 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 @@ -8,7 +8,7 @@ import cutlass from cutlass import cute -from cutlass import Boolean, Float32, Int32, Int64, Uint32, Uint8 +from cutlass import Boolean, Float32, Int32, Int64, Uint32 from cuda.bindings.driver import CUstream from transformer_engine.common.CuTeDSL.utils import ( @@ -44,6 +44,8 @@ def quantize_rowwise_mxfp8_g2r( CHECK_BOUNDS: cutlass.Constexpr[bool], input_policy: Int64, output_policy: Int64, + cta_row: Int32, + cta_col: Int32, ): """Quantize a CTA's (TILE_ROWS, TILE_COLS, NUM_TILES) GMEM tiles. @@ -97,7 +99,6 @@ def quantize_rowwise_mxfp8_g2r( gX_thread = cute.composition(gX_tile, tv_payload)[tidx, None] valid = Boolean(True) if cutlass.const_expr(CHECK_BOUNDS): - cta_row, cta_col, _ = cute.arch.block_idx() row = cta_row * TILE_ROWS + tidx // CTA_THREADS_X col = (cta_col * NUM_TILES + tile_idx) * TILE_COLS + ( tidx % CTA_THREADS_X @@ -137,16 +138,14 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Float32: amax = extract_mx_block_amax(rX_tile_i32) exponent = cvt_f32_to_fp8e8m0fnu(amax * cfg.MAX_NORM_RCP) - # Collect 4 scale bytes and write them using one thread - byte = Uint32(exponent.bitcast(Uint8)) - c1 = cute.arch.shuffle_sync_down(byte, 2) - c2 = cute.arch.shuffle_sync_down(byte, 4) - c3 = cute.arch.shuffle_sync_down(byte, 6) - if tidx % 8 == 0 and tile_valid_flags[tile_idx]: - scale_word = cute.make_tensor( - cute.recast_ptr(mS_thread.iterator, dtype=Uint32), cute.make_layout(1) - ) - scale_word[0] = byte | (c1 << Uint32(8)) | (c2 << Uint32(16)) | (c3 << Uint32(24)) + # The even lane of each pair stores its block's scale byte, as cast_rowwise.cuh does. + # A warp's stores coalesce in either layout: one 16-byte run when packed, four + # 4-byte runs when GEMM-swizzled. Gathering the bytes into 32-bit words first + # touches the same sectors but adds shuffles and a branch (+12% instructions). + # TODO(kainingz): not using vectorized word store seems to be faster in CuTeDSL. + # Check if we should port this to CUDA + if tidx % 2 == 0 and tile_valid_flags[tile_idx]: + mS_thread[0] = exponent # Compute and write the quantized MXFP8 values to GMEM reciprocal_packed = exp2_bf16x2_rcp(exponent) @@ -164,14 +163,17 @@ class MXFP8QuantizeG2RRowwise2LaneKernel(MXFP8QuantizeKernelBase): One CTA tile has TILE_ROWS = CTA_THREADS_Y and TILE_COLS = CTA_THREADS_X * _ELEMENTS_PER_THREAD. The CTA processes NUM_TILES = 1 or 2 in the original input matrix. Both scale - formats use the same payload tiling; a composed layout maps swizzled scales. + formats use the same payload tiling; swizzled scale offsets are computed directly. _L2_CACHED_CTA_PERCENT tunes the share of CTAs whose input loads use normal L2 retention: 0 streams all input and 100 caches all input normally. """ _CTA_THREADS_X = 32 _ELEMENTS_PER_THREAD = 16 + # cast_rowwise.cuh's tier table: the two smaller tiers stream all of their input. _L2_CACHED_CTA_PERCENT = 40 + _SMALL_TIER_CACHED_PERCENT = 0 + _MAX_GRID_Y = 65535 def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg @@ -201,6 +203,7 @@ def __call__( CTA_THREADS_Y=8, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=1, + L2_CACHED_CTA_PERCENT=self._SMALL_TIER_CACHED_PERCENT, ) elif output_bytes <= 48 << 20: self._launch( @@ -212,6 +215,7 @@ def __call__( CTA_THREADS_Y=8, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=2, + L2_CACHED_CTA_PERCENT=self._SMALL_TIER_CACHED_PERCENT, ) elif output_bytes <= 96 << 20: self._launch( @@ -223,6 +227,7 @@ def __call__( CTA_THREADS_Y=4, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=2, + L2_CACHED_CTA_PERCENT=self._L2_CACHED_CTA_PERCENT, ) else: self._launch( @@ -234,6 +239,7 @@ def __call__( CTA_THREADS_Y=8, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=2, + L2_CACHED_CTA_PERCENT=self._L2_CACHED_CTA_PERCENT, ) @cute.jit @@ -247,8 +253,9 @@ def _launch( CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], NUM_TILES: cutlass.Constexpr[int], + L2_CACHED_CTA_PERCENT: cutlass.Constexpr[int], ): - assert 0 <= self._L2_CACHED_CTA_PERCENT <= 100 + assert 0 <= L2_CACHED_CTA_PERCENT <= 100 assert self._ELEMENTS_PER_THREAD == 16 assert CTA_THREADS_X % 8 == 0 assert CTA_THREADS_Y * CTA_THREADS_X % 32 == 0 @@ -265,8 +272,15 @@ def _launch( grid_cols = cute.ceil_div(N, CTA_COLS) check_bounds = M % CTA_ROWS != 0 or N % CTA_COLS != 0 grid = Int64(grid_rows) * grid_cols - first_streaming_cta = grid * self._L2_CACHED_CTA_PERCENT // 100 - + first_streaming_cta = grid * L2_CACHED_CTA_PERCENT // 100 + + # Column tiles vary fastest, so consecutive CTAs read consecutive memory along a + # row, as cast_rowwise.cuh's flat grid does. Putting row tiles on x instead + # sends concurrent CTAs down a column strip, one 2 KiB piece per row, and costs + # up to 9% at large shapes. Row tiles past grid.y's limit spill into grid.z. + grid_y = cutlass.min(grid_rows, self._MAX_GRID_Y) + grid_z = cute.ceil_div(grid_rows, self._MAX_GRID_Y) + launch_grid = [grid_cols, grid_y, grid_z] if grid > 0: if check_bounds: self.kernel( @@ -278,8 +292,9 @@ def _launch( CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, NUM_TILES=NUM_TILES, + L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, CHECK_BOUNDS=True, - ).launch(grid=[grid_rows, grid_cols], block=[THREADS_PER_CTA], stream=stream) + ).launch(grid=launch_grid, block=[THREADS_PER_CTA], stream=stream) else: self.kernel( mX, @@ -290,8 +305,9 @@ def _launch( CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, NUM_TILES=NUM_TILES, + L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, CHECK_BOUNDS=False, - ).launch(grid=[grid_rows, grid_cols], block=[THREADS_PER_CTA], stream=stream) + ).launch(grid=launch_grid, block=[THREADS_PER_CTA], stream=stream) @cute.kernel def kernel( @@ -304,6 +320,7 @@ def kernel( CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], NUM_TILES: cutlass.Constexpr[int], + L2_CACHED_CTA_PERCENT: cutlass.Constexpr[int], CHECK_BOUNDS: cutlass.Constexpr[bool], ): """Skip the CTA when the noop flag is set; otherwise run the kernel body.""" @@ -316,6 +333,7 @@ def kernel( CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, NUM_TILES=NUM_TILES, + L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, CHECK_BOUNDS=CHECK_BOUNDS, ) @@ -329,13 +347,16 @@ def _kernel_main( CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], NUM_TILES: cutlass.Constexpr[int], + L2_CACHED_CTA_PERCENT: cutlass.Constexpr[int], CHECK_BOUNDS: cutlass.Constexpr[bool], ): """Construct CTA tile views and run the specialized quantization.""" # Dispatch checks 32-byte input alignment before selecting this kernel. gX = cute.make_tensor(mX.iterator.align(32), mX.layout) - # CUDA's x/y axes hold row/column tiles in the same order as local_tile. - cta_coord = cute.arch.block_idx()[:2] + # grid.x walks column tiles; grid.y and grid.z together walk row tiles. + bidx, bidy, bidz = cute.arch.block_idx() + cta_row = bidz * self._MAX_GRID_Y + bidy + cta_coord = (cta_row, bidx) M, N = mX.shape TILE_ROWS = CTA_THREADS_Y TILE_COLS = CTA_THREADS_X * self._ELEMENTS_PER_THREAD @@ -343,10 +364,11 @@ def _kernel_main( CTA_COLS = TILE_COLS * NUM_TILES gS = mS - # Apply the swizzled scale layout if requested + # Use the swizzled layout as is. Composing it with (M, N // 32) layouts to flatten + # its shape made CuTe emit runtime integer div/mod on every scale address, + # about 7x this kernel's instruction count. if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - mS_swizzled, _ = derive_swizzled_scale_layout(M, N, True, False, mS, None) - gS = cute.composition(mS_swizzled, (cute.make_layout(M), cute.make_layout(N // 32))) + gS, _ = derive_swizzled_scale_layout(M, N, True, False, mS, None) # CTA_TILER is the block that a CTA processes in total, which consists of NUM_TILES horizontal tiles CTA_TILER = (CTA_ROWS, CTA_COLS) @@ -356,12 +378,12 @@ def _kernel_main( gS_cta = cute.local_tile(gS, CTA_SCALE_TILER, cta_coord) fraction = Float32(1.0) - if cutlass.const_expr(self._L2_CACHED_CTA_PERCENT > 0): + if cutlass.const_expr(L2_CACHED_CTA_PERCENT > 0): fraction = Float32(0.0) - # A linear index is needed only for the L2 policy split. - cta_id = cute.crd2idx( - (Int64(cta_coord[0]), Int64(cta_coord[1])), cute.arch.grid_dim()[:2] - ) + # A linear index is needed only for the L2 policy split. It follows launch + # order, so the cached CTAs are the first ones scheduled. + gdx, _, _ = cute.arch.grid_dim() + cta_id = Int64(cta_row) * gdx + bidx if cta_id >= first_streaming_cta: fraction = Float32(1.0) input_policy = create_l2_policy(False, fraction) @@ -375,18 +397,22 @@ def _kernel_main( gO_tiles = cute.local_tile(gO_cta, TILER, (0, None)) mS_tiles = cute.local_tile(gS_cta, SCALE_TILER, (0, None)) - quantize_rowwise_mxfp8_g2r( - gX_tiles, - gO_tiles, - mS_tiles, - self.cfg, - M, - N, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, - NUM_TILES=NUM_TILES, - CHECK_BOUNDS=CHECK_BOUNDS, - input_policy=input_policy, - output_policy=output_policy, - ) + # The last grid.z slice can overshoot the row tiles; those CTAs have no work. + if cta_row * CTA_ROWS < M: + quantize_rowwise_mxfp8_g2r( + gX_tiles, + gO_tiles, + mS_tiles, + self.cfg, + M, + N, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + NUM_TILES=NUM_TILES, + CHECK_BOUNDS=CHECK_BOUNDS, + input_policy=input_policy, + output_policy=output_policy, + cta_row=cta_coord[0], + cta_col=cta_coord[1], + ) From 870b23cbe23e3073c0d3d5898cccba015d378595 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 9 Oct 2026 15:22:49 -0700 Subject: [PATCH 09/29] disable invariant temporarily Signed-off-by: Kaining Zhong --- .../CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py | 4 +++- .../CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 886b9eb27bd..2f194e804cf 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -404,7 +404,9 @@ def quantize_bidimensional_mxfp8_g2r( cute.nvgpu.CopyG2ROp(), cfg.DTYPE, num_bits_per_copy=ELEMENTS_X_PER_THREAD * cfg.DTYPE.width, - invariant=True, + # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using + # invariant=True here will drop the cache policy + invariant=False, ) store_atom = cute.make_copy_atom( cute.nvgpu.CopyR2GOp(), 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 a264b578575..dc39ecfb02d 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 @@ -82,7 +82,9 @@ def quantize_rowwise_mxfp8_g2r( cute.nvgpu.CopyG2ROp(), cfg.DTYPE, num_bits_per_copy=LOAD_WIDTH, - invariant=True, # hint that data being read stays unchanged during the kernel + # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using + # invariant=True here will drop the cache policy + invariant=False, ) store_atom = cute.make_copy_atom( cute.nvgpu.CopyR2GOp(), From 4ef3404dfe6eb1efc3cbb51df83790afae6af9c7 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 9 Oct 2026 17:02:12 -0700 Subject: [PATCH 10/29] fix Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 54 +++++++++++++++++-- 1 file changed, 51 insertions(+), 3 deletions(-) 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 dc39ecfb02d..7f285677357 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 @@ -46,6 +46,7 @@ def quantize_rowwise_mxfp8_g2r( output_policy: Int64, cta_row: Int32, cta_col: Int32, + num_elems: Int64, ): """Quantize a CTA's (TILE_ROWS, TILE_COLS, NUM_TILES) GMEM tiles. @@ -105,7 +106,10 @@ def quantize_rowwise_mxfp8_g2r( col = (cta_col * NUM_TILES + tile_idx) * TILE_COLS + ( tidx % CTA_THREADS_X ) * ELEMENTS_PER_THREAD - valid = row < M and col < N + # When we apply the (-1, 1024) manually constructed layout to the input tensor, + # the last row might not be a full row so we need to check OOB with element number + # instead of just M and N + valid = col < N and Int64(row) * N + col < num_elems tile_valid_flags[tile_idx] = valid rX[None, tile_idx].fill(0) if valid: @@ -194,7 +198,40 @@ def __call__( mWorkspace: Optional[cute.Tensor], stream: CUstream, ): - output_bytes = Int64(mX.shape[0]) * mX.shape[1] + M, N = mX.shape + num_elems = Int64(M) * N + # View the input tensor as a layout with last dimension = 1024 when not swizzled and the number of elements can be expressed using i32 + # However, M*N might not be divisible by 1024, so the last row of this layout might be a partial row with fewer than 1024 elements, + # which is why we need to pass num_elems to the kernel so OOB checks will identify these regions as invalid offsets correctly + if cutlass.const_expr(not self.cfg.WITH_GEMM_SWIZZLED_SCALES): + if num_elems <= 2**31 - 1: + rows = cute.ceil_div(num_elems, 1024) + flat = cute.make_layout((rows, 1024), stride=(1024, 1)) + self._launch_tiers( + cute.make_tensor(mX.iterator, flat), + cute.make_tensor(mO_row.iterator, flat), + cute.make_tensor(mS_row.iterator, cute.make_layout((rows, 32), stride=(32, 1))), + mNoop, + stream, + num_elems, + ) + else: + self._launch_tiers(mX, mO_row, mS_row, mNoop, stream, num_elems) + else: + self._launch_tiers(mX, mO_row, mS_row, mNoop, stream, num_elems) + + @cute.jit + def _launch_tiers( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mNoop: cute.Pointer, + stream: CUstream, + num_elems: Int64, + ): + """Pick the launch configuration for this many elements (cast_rowwise.cuh's tiers).""" + output_bytes = num_elems if output_bytes <= 12 << 20: self._launch( mX, @@ -202,6 +239,7 @@ def __call__( mS_row, mNoop, stream, + num_elems, CTA_THREADS_Y=8, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=1, @@ -214,6 +252,7 @@ def __call__( mS_row, mNoop, stream, + num_elems, CTA_THREADS_Y=8, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=2, @@ -226,6 +265,7 @@ def __call__( mS_row, mNoop, stream, + num_elems, CTA_THREADS_Y=4, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=2, @@ -238,6 +278,7 @@ def __call__( mS_row, mNoop, stream, + num_elems, CTA_THREADS_Y=8, CTA_THREADS_X=self._CTA_THREADS_X, NUM_TILES=2, @@ -252,6 +293,7 @@ def _launch( mS: cute.Tensor, mNoop: cute.Pointer, stream: CUstream, + num_elems: Int64, CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], NUM_TILES: cutlass.Constexpr[int], @@ -272,7 +314,7 @@ def _launch( grid_rows = cute.ceil_div(M, CTA_ROWS) grid_cols = cute.ceil_div(N, CTA_COLS) - check_bounds = M % CTA_ROWS != 0 or N % CTA_COLS != 0 + check_bounds = M % CTA_ROWS != 0 or N % CTA_COLS != 0 or num_elems != Int64(M) * N grid = Int64(grid_rows) * grid_cols first_streaming_cta = grid * L2_CACHED_CTA_PERCENT // 100 @@ -291,6 +333,7 @@ def _launch( mS, mNoop, first_streaming_cta, + num_elems, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, NUM_TILES=NUM_TILES, @@ -304,6 +347,7 @@ def _launch( mS, mNoop, first_streaming_cta, + num_elems, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, NUM_TILES=NUM_TILES, @@ -319,6 +363,7 @@ def kernel( mS: cute.Tensor, mNoop: cute.Pointer, first_streaming_cta: Int64, + num_elems: Int64, CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], NUM_TILES: cutlass.Constexpr[int], @@ -332,6 +377,7 @@ def kernel( mO, mS, first_streaming_cta, + num_elems, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, NUM_TILES=NUM_TILES, @@ -346,6 +392,7 @@ def _kernel_main( mO: cute.Tensor, mS: cute.Tensor, first_streaming_cta: Int64, + num_elems: Int64, CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], NUM_TILES: cutlass.Constexpr[int], @@ -417,4 +464,5 @@ def _kernel_main( output_policy=output_policy, cta_row=cta_coord[0], cta_col=cta_coord[1], + num_elems=num_elems, ) From ef9877df4a10d5efa09b1461688fb8612b4d3704 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 18:45:39 +0000 Subject: [PATCH 11/29] refactor Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_common.py | 146 +++++------------- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 69 ++++++--- .../mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py | 37 ++++- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 85 ++++++---- .../mxfp8/quantize_mxfp8_g2s_bidimensional.py | 61 +++++--- 5 files changed, 212 insertions(+), 186 deletions(-) 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 14d3b88ca66..b8859210d0b 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -19,7 +19,6 @@ str_to_cutlass_dtype, abs_max_x2_bf16, device_compute_capability, - is_packed16, ) from transformer_engine.common.CuTeDSL.activations import ( act_relu, @@ -211,9 +210,32 @@ def __str__(self): __repr__ = __str__ + 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) + + class MXFP8QuantizeKernelBase(abc.ABC): """Base class for MXFP8 quantize kernels.""" + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: # pylint: disable=unused-argument + """Compile time check: whether this specialized kernel handles cfg at all.""" + return False + + @staticmethod + def eligible_to_dispatch( + cfg: 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, + ) -> Boolean: + """Runtime check: whether we can dispatch to the instantiated kernel with this call's shapes and noops.""" + raise NotImplementedError + @abc.abstractmethod def __call__( self, @@ -281,16 +303,18 @@ def __init__(self, cfg: MXFP8QuantizeConfig): ) # We only need to mask when WITH_ACT is enabled because with activations applied zeros filled by TMA affect block statistics self.general_kernel = MXFP8QuantizeKernel(cfg) if self.ACT_NEED_MASKING else None - self.g2r_rowwise_1lane = MXFP8QuantizeG2RRowwise1LaneKernel(cfg) if cfg.ROWWISE else None - self.g2s_bidim = ( - MXFP8QuantizeSpecializedBidimensionalKernel(cfg) - if cfg.ROWWISE and cfg.COLWISE - else None - ) - self.g2r_rowwise_2lane = MXFP8QuantizeG2RRowwise2LaneKernel(cfg) if cfg.ROWWISE else None - self.g2r_bidim = ( - MXFP8QuantizeRegisterBidimensionalKernel(cfg) if cfg.ROWWISE and cfg.COLWISE else None - ) + + # Specialized cast-only kernels in a specific dispatch priority order so we prefer to dispatch to kernels that come first in the list. + self.specialized_kernels = [ + kernel_cls(cfg) + for kernel_cls in ( + MXFP8QuantizeG2RRowwise2LaneKernel, + MXFP8QuantizeRegisterBidimensionalKernel, + MXFP8QuantizeG2RRowwise1LaneKernel, + MXFP8QuantizeSpecializedBidimensionalKernel, + ) + if kernel_cls.supports(cfg) + ] @cute.jit def __call__( @@ -306,104 +330,12 @@ def __call__( mWorkspace: Optional[cute.Tensor], stream: CUstream, ): - M = mX.shape[0] - N = mX.shape[1] - # Select specialized kernels if possible - plain_cast_only = ( - not self.cfg.WITH_AMAX - and not self.cfg.WITH_DBIAS - and not self.cfg.WITH_DACT - and not self.cfg.WITH_ACT - ) dispatched_to_specialized = False - # The two-lane rowwise and g2r bidimensional kernels use BF16 packed math and 32-byte loads. - # Keep the existing kernels for FP16, fusions, ragged bidimensional tiles, - # misaligned input views and compilers predating 256-bit global loads. - if cutlass.const_expr( - plain_cast_only - and self.cfg.DTYPE is cutlass.BFloat16 - and cutlass.target_version(min_version="12.9") - ): - input_is_aligned = mX.iterator.toint() % 32 == 0 - if cutlass.const_expr(self.cfg.ROWWISE and not self.cfg.COLWISE): - scale_stride_matches = Boolean(True) - if cutlass.const_expr(not self.cfg.WITH_GEMM_SWIZZLED_SCALES): - scale_stride_matches = mS_row.shape[1] == N // 32 - # Column tiles use grid.x and row tiles spill from grid.y into grid.z, so - # no shape exceeds the grid limits. - if N % 128 == 0 and input_is_aligned and scale_stride_matches: - dispatched_to_specialized = True - self.g2r_rowwise_2lane( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - if cutlass.const_expr( - self.cfg.ROWWISE and self.cfg.COLWISE and not self.cfg.WITH_GEMM_SWIZZLED_SCALES - ): - columnwise_scales_are_aligned = mS_col.iterator.toint() % 16 == 0 - if ( - M % 32 == 0 - and N % 256 == 0 - and M // 32 <= 65535 - and input_is_aligned - and columnwise_scales_are_aligned - ): - dispatched_to_specialized = True - self.g2r_bidim( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - # The one-lane rowwise and g2s bidimensional kernels support packed16 types (bf16/fp16). - if cutlass.const_expr(plain_cast_only and is_packed16(self.cfg.DTYPE)): - if cutlass.const_expr(self.cfg.ROWWISE and not self.cfg.COLWISE): - g2r_rowwise_1lane_grid_fits = ( - cute.ceil_div(M, self.g2r_rowwise_1lane._TILE_ROWS) <= 65535 - ) - # The one-lane rowwise kernel requires N divisible by 128 for vectorized stores. - if not dispatched_to_specialized and N % 128 == 0 and g2r_rowwise_1lane_grid_fits: - dispatched_to_specialized = True - self.g2r_rowwise_1lane( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - # The g2s bidimensional kernel supports swizzled scales but we don't dispatch to it for now - if cutlass.const_expr( - self.cfg.ROWWISE and self.cfg.COLWISE and not self.cfg.WITH_GEMM_SWIZZLED_SCALES - ): - g2s_bidim_grid_fits = ( - cute.ceil_div( - M, - self.g2s_bidim._TILE_ROWS * self.g2s_bidim._NUM_TILES_Y, - ) - <= 65535 - ) - if not dispatched_to_specialized and g2s_bidim_grid_fits: + for kernel in self.specialized_kernels: + if not dispatched_to_specialized: + if kernel.eligible_to_dispatch(self.cfg, mX, mO_row, mS_row, mO_col, mS_col, mNoop): dispatched_to_specialized = True - self.g2s_bidim( + kernel( mX, mO_row, mS_row, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 2f194e804cf..27e43988450 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -8,7 +8,7 @@ import cutlass from cutlass import cute -from cutlass import Int32, Int64, Uint32, Uint8 +from cutlass import Boolean, Int32, Int64, Uint32, Uint8 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( @@ -20,11 +20,11 @@ ) from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 from .quantize_mxfp8_common import ( + is_plain_cast, MXFP8_BLOCK_SCALING_SIZE, MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, bf16_pair_magnitude, - noop_flag_is_set, ) @@ -506,6 +506,41 @@ class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): _CTA_THREADS_Y = 8 _CTA_THREADS_X = 32 + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + # BF16 packed math and 256-bit loads; packed scales only (no colwise GEMM swizzle yet). + return ( + cfg.is_plain_cast() + and cfg.DTYPE is cutlass.BFloat16 + and cfg.ROWWISE + and cfg.COLWISE + and not cfg.WITH_GEMM_SWIZZLED_SCALES + and cutlass.target_version(min_version="12.9") + ) + + @staticmethod + @cute.jit + def eligible_to_dispatch( + 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, + ) -> Boolean: + M, N = mX.shape + # Whole 32-row by 256-column tiles, grid.y within limits, and aligned input and + # colwise-scale pointers. + return ( + M % 32 == 0 + and N % 256 == 0 + and M // 32 <= 65535 + and mX.iterator.toint() % 32 == 0 + and mS_col.iterator.toint() % 16 == 0 + and mNoop.toint() == 0 + ) + def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg # One CTA only quantize one columnwise MXFP8 block at a time @@ -520,7 +555,7 @@ def __call__( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, + mNoop: cute.Pointer, # Unused: the dispatcher routes noop calls to the general kernel mDActInput: Optional[cute.Tensor], mWorkspace: Optional[cute.Tensor], stream: CUstream, @@ -536,7 +571,6 @@ def __call__( mS_row, mO_col, mS_col, - mNoop, stream, ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=4, @@ -548,7 +582,6 @@ def __call__( mS_row, mO_col, mS_col, - mNoop, stream, ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=1, @@ -560,7 +593,6 @@ def __call__( mS_row, mO_col, mS_col, - mNoop, stream, ELEMENTS_X_PER_THREAD=8, CLUSTER_WIDTH=8, @@ -574,7 +606,6 @@ def _launch( mS_row: cute.Tensor, mO_col: cute.Tensor, mS_col: cute.Tensor, - mNoop: cute.Pointer, stream: CUstream, ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], CLUSTER_WIDTH: cutlass.Constexpr[int], @@ -620,7 +651,6 @@ def _launch( gS_row, gO_col, gS_col, - mNoop, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) elif N == 4096: @@ -642,7 +672,6 @@ def _launch( gS_row, gO_col, gS_col, - mNoop, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) else: @@ -652,7 +681,6 @@ def _launch( mS_row, mO_col, mS_col, - mNoop, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) @@ -664,19 +692,18 @@ def kernel( mS_row: cute.Tensor, mO_col: cute.Tensor, mS_col: cute.Tensor, - mNoop: cute.Pointer, ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], ): - """Skip the CTA when the noop flag is set; otherwise run the kernel body.""" - if not noop_flag_is_set(mNoop): - self._kernel_main( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ) + """Run the kernel body. Like the CUDA specialized cast kernels it does not consume the + cast_noop flag: the dispatcher routes calls with one to the general kernel.""" + self._kernel_main( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ) @cute.jit def _kernel_main( 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 552b6429465..fd76e039fa6 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 @@ -11,10 +11,11 @@ import cutlass from cutlass import cute -from cutlass import Float8E8M0FNU, Int32, Uint32, Uint8 +from cutlass import Boolean, Float8E8M0FNU, Int32, Uint32, Uint8 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( + is_packed16, abs_max_x2_bf16, abs_max_x2_f16, x2_lo_to_f32_bf16, @@ -35,7 +36,6 @@ MXFP8QuantizeKernelBase, MXFP8_BLOCK_SCALING_SIZE, derive_swizzled_scale_layout, - noop_flag_is_set, ) @@ -51,6 +51,29 @@ class MXFP8QuantizeG2RRowwise1LaneKernel(MXFP8QuantizeKernelBase): _TILE_COLS = 1024 _THREADS_PER_CTA = 128 + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + return cfg.is_plain_cast() and is_packed16(cfg.DTYPE) and cfg.ROWWISE and not cfg.COLWISE + + @staticmethod + @cute.jit + def eligible_to_dispatch( + 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, + ) -> Boolean: + M, N = mX.shape + # N divisible by 128 for vectorized stores, and grid.y within limits. + return ( + N % 128 == 0 + and cute.ceil_div(M, MXFP8QuantizeG2RRowwise1LaneKernel._TILE_ROWS) <= 65535 + and mNoop.toint() == 0 + ) + def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg # If True, then this kernel will first write each thread's scale byte to a shared memory buffer, @@ -66,7 +89,7 @@ def __call__( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], # Unused, kept for API compatibility mAmax: Optional[cute.Tensor], # Unused, kept for API compatibility - mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device + mNoop: cute.Pointer, # Unused: the dispatcher routes noop calls to the general kernel mDActInput: Optional[cute.Tensor], # Unused, kept for API compatibility mWorkspace: Optional[cute.Tensor], # Unused, kept for API compatibility stream: CUstream, @@ -92,7 +115,6 @@ def __call__( mX, mO_row, mS_row, - mNoop, mX.element_type, ).launch(grid=grid, block=block, stream=stream) @@ -102,12 +124,11 @@ def kernel( mX: cute.Tensor, mO_row: cute.Tensor, mS_row: cute.Tensor, - mNoop: cute.Pointer, DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): - """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" - if not noop_flag_is_set(mNoop): - self._kernel_main(mX, mO_row, mS_row, DTYPE) + """Device entry. Like the CUDA specialized cast kernels it does not consume the + cast_noop flag: the dispatcher routes calls with one to the general kernel.""" + self._kernel_main(mX, mO_row, mS_row, DTYPE) @cute.jit def _kernel_main( 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 7f285677357..1304a47dd00 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 @@ -9,7 +9,7 @@ import cutlass from cutlass import cute from cutlass import Boolean, Float32, Int32, Int64, Uint32 -from cuda.bindings.driver import CUstream +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( abs_max_x2_bf16, @@ -25,7 +25,6 @@ MXFP8QuantizeKernelBase, bf16_pair_magnitude, derive_swizzled_scale_layout, - noop_flag_is_set, ) @@ -35,7 +34,6 @@ def quantize_rowwise_mxfp8_g2r( gO_tiles: cute.Tensor, mS_tiles: cute.Tensor, cfg: cutlass.Constexpr[MXFP8QuantizeConfig], - M: Int32, N: Int32, CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], @@ -181,6 +179,42 @@ class MXFP8QuantizeG2RRowwise2LaneKernel(MXFP8QuantizeKernelBase): _SMALL_TIER_CACHED_PERCENT = 0 _MAX_GRID_Y = 65535 + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + # BF16 packed math and 256-bit loads; compilers predating those keep the other kernels. + return ( + cfg.is_plain_cast() + and cfg.DTYPE is cutlass.BFloat16 + and cfg.ROWWISE + and not cfg.COLWISE + and cutlass.target_version(min_version="12.9") + ) + + @staticmethod + @cute.jit + def eligible_to_dispatch( + 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, + ) -> Boolean: + N = mX.shape[1] + scale_stride_matches = Boolean(True) + if cutlass.const_expr(not cfg.WITH_GEMM_SWIZZLED_SCALES): + # The flat view needs the packed scales to be contiguous. + scale_stride_matches = mS_row.shape[1] == N // 32 + # Column tiles use grid.x and row tiles spill from grid.y into grid.z, so no shape + # exceeds the grid limits. + return ( + N % 128 == 0 + and mX.iterator.toint() % 32 == 0 + and scale_stride_matches + and mNoop.toint() == 0 + ) + def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg @@ -193,7 +227,7 @@ def __call__( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, + mNoop: cute.Pointer, # Unused: the dispatcher routes noop calls to the general kernel mDActInput: Optional[cute.Tensor], mWorkspace: Optional[cute.Tensor], stream: CUstream, @@ -211,14 +245,13 @@ def __call__( cute.make_tensor(mX.iterator, flat), cute.make_tensor(mO_row.iterator, flat), cute.make_tensor(mS_row.iterator, cute.make_layout((rows, 32), stride=(32, 1))), - mNoop, stream, num_elems, ) else: - self._launch_tiers(mX, mO_row, mS_row, mNoop, stream, num_elems) + self._launch_tiers(mX, mO_row, mS_row, stream, num_elems) else: - self._launch_tiers(mX, mO_row, mS_row, mNoop, stream, num_elems) + self._launch_tiers(mX, mO_row, mS_row, stream, num_elems) @cute.jit def _launch_tiers( @@ -226,7 +259,6 @@ def _launch_tiers( mX: cute.Tensor, mO_row: cute.Tensor, mS_row: cute.Tensor, - mNoop: cute.Pointer, stream: CUstream, num_elems: Int64, ): @@ -237,7 +269,6 @@ def _launch_tiers( mX, mO_row, mS_row, - mNoop, stream, num_elems, CTA_THREADS_Y=8, @@ -250,7 +281,6 @@ def _launch_tiers( mX, mO_row, mS_row, - mNoop, stream, num_elems, CTA_THREADS_Y=8, @@ -263,7 +293,6 @@ def _launch_tiers( mX, mO_row, mS_row, - mNoop, stream, num_elems, CTA_THREADS_Y=4, @@ -276,7 +305,6 @@ def _launch_tiers( mX, mO_row, mS_row, - mNoop, stream, num_elems, CTA_THREADS_Y=8, @@ -291,7 +319,6 @@ def _launch( mX: cute.Tensor, mO: cute.Tensor, mS: cute.Tensor, - mNoop: cute.Pointer, stream: CUstream, num_elems: Int64, CTA_THREADS_Y: cutlass.Constexpr[int], @@ -331,7 +358,6 @@ def _launch( mX, mO, mS, - mNoop, first_streaming_cta, num_elems, CTA_THREADS_Y=CTA_THREADS_Y, @@ -345,7 +371,6 @@ def _launch( mX, mO, mS, - mNoop, first_streaming_cta, num_elems, CTA_THREADS_Y=CTA_THREADS_Y, @@ -361,7 +386,6 @@ def kernel( mX: cute.Tensor, mO: cute.Tensor, mS: cute.Tensor, - mNoop: cute.Pointer, first_streaming_cta: Int64, num_elems: Int64, CTA_THREADS_Y: cutlass.Constexpr[int], @@ -370,20 +394,20 @@ def kernel( L2_CACHED_CTA_PERCENT: cutlass.Constexpr[int], CHECK_BOUNDS: cutlass.Constexpr[bool], ): - """Skip the CTA when the noop flag is set; otherwise run the kernel body.""" - if not noop_flag_is_set(mNoop): - self._kernel_main( - mX, - mO, - mS, - first_streaming_cta, - num_elems, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - NUM_TILES=NUM_TILES, - L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, - CHECK_BOUNDS=CHECK_BOUNDS, - ) + """Run the kernel body. Like cast_rowwise.cuh, there is no noop handling: the + dispatcher only selects this kernel when no cast_noop flag is passed.""" + self._kernel_main( + mX, + mO, + mS, + first_streaming_cta, + num_elems, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + NUM_TILES=NUM_TILES, + L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, + CHECK_BOUNDS=CHECK_BOUNDS, + ) @cute.jit def _kernel_main( @@ -453,7 +477,6 @@ def _kernel_main( gO_tiles, mS_tiles, self.cfg, - M, N, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py index 89cfff90093..f6c6d391d99 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py @@ -12,7 +12,7 @@ import cutlass from cutlass import cute from cutlass import pipeline -from cutlass import Float32, Float8E8M0FNU, Int32, Int64, Uint32 +from cutlass import Boolean, Float32, Float8E8M0FNU, Int32, Int64, Uint32 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( @@ -39,7 +39,6 @@ MXFP8_BLOCK_SCALING_SIZE, THREADS_PER_WARP, derive_swizzled_scale_layout, - noop_flag_is_set, ) @@ -249,6 +248,32 @@ class MXFP8QuantizeSpecializedBidimensionalKernel(MXFP8QuantizeKernelBase): _SCALE_ROWS = _TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE _SCALE_COLS = _TILE_COLS // MXFP8_BLOCK_SCALING_SIZE + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + # This kernel supports swizzled scales but we don't dispatch to it for now. + return ( + cfg.is_plain_cast() + and is_packed16(cfg.DTYPE) + and cfg.ROWWISE + and cfg.COLWISE + and not cfg.WITH_GEMM_SWIZZLED_SCALES + ) + + @staticmethod + @cute.jit + def eligible_to_dispatch( + 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, + ) -> Boolean: + cls = MXFP8QuantizeSpecializedBidimensionalKernel + rows_per_cta = cls._TILE_ROWS * cls._NUM_TILES_Y # packed scales: one tile row + return cute.ceil_div(mX.shape[0], rows_per_cta) <= 65535 and mNoop.toint() == 0 + def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg if cfg.WITH_GEMM_SWIZZLED_SCALES: @@ -265,7 +290,7 @@ def __call__( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device + mNoop: cute.Pointer, # Unused: the dispatcher routes noop calls to the general kernel mDActInput: Optional[cute.Tensor], mWorkspace: Optional[cute.Tensor], stream: CUstream, @@ -329,7 +354,6 @@ def __call__( mX, mS_row, mS_col, - mNoop, mX.element_type, tma_atom, tma_src, @@ -345,7 +369,6 @@ def kernel( mX: cute.Tensor, mS_row: cute.Tensor, mS_col: cute.Tensor, - mNoop: cute.Pointer, dtype: cutlass.Constexpr[Type[cutlass.Numeric]], tma_atom: cute.CopyAtom, tma_src: cute.Tensor, @@ -354,20 +377,20 @@ def kernel( tma_atom_out_col: cute.CopyAtom, tma_dst_out_col: cute.Tensor, ): - """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" - if not noop_flag_is_set(mNoop): - self._kernel_main( - mX, - mS_row, - mS_col, - dtype, - tma_atom, - tma_src, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - ) + """Device entry. Like the CUDA specialized cast kernels it does not consume the + cast_noop flag: the dispatcher routes calls with one to the general kernel.""" + self._kernel_main( + mX, + mS_row, + mS_col, + dtype, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + ) @cute.jit def _kernel_main( From d6bb25862eb793ce1d99455cc8363ac1b9343806 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 12 Oct 2026 18:47:00 +0000 Subject: [PATCH 12/29] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py | 1 - 1 file changed, 1 deletion(-) 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 b8859210d0b..748f19756f4 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -209,7 +209,6 @@ def __str__(self): __repr__ = __str__ - 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) From fffc87490c028869d45f13a3994c1b383b713033 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 18:48:29 +0000 Subject: [PATCH 13/29] fix Signed-off-by: Kaining Zhong --- .../CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 27e43988450..fa1c2d12807 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -20,7 +20,6 @@ ) from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 from .quantize_mxfp8_common import ( - is_plain_cast, MXFP8_BLOCK_SCALING_SIZE, MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, From 2fd0aca375bda99e0bc11757550ce5842439d06a Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 11:59:00 -0700 Subject: [PATCH 14/29] fix Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 18 ++++++++++++++--- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 20 +++++++++++++++---- 2 files changed, 31 insertions(+), 7 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index fa1c2d12807..e0df0abf804 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -402,7 +402,10 @@ def quantize_bidimensional_mxfp8_g2r( load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), cfg.DTYPE, - num_bits_per_copy=ELEMENTS_X_PER_THREAD * cfg.DTYPE.width, + num_bits_per_copy=min( + ELEMENTS_X_PER_THREAD * cfg.DTYPE.width, + MXFP8QuantizeRegisterBidimensionalKernel.max_num_bits_per_copy(), + ), # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using # invariant=True here will drop the cache policy invariant=False, @@ -505,16 +508,25 @@ class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): _CTA_THREADS_Y = 8 _CTA_THREADS_X = 32 + @staticmethod + @cute.jit + def max_num_bits_per_copy() -> int: + """Widest global load to issue: 256-bit loads (`ld.global.v8.b32`) need PTX ISA 8.8, + i.e. CUDA 12.9. Older toolchains split each one into two 128-bit loads, as + cast_rowwise.cuh / cast_bidim.cuh do via util/ptx.cuh's ld_global_nc_b32x8.""" + if cutlass.const_expr(cutlass.target_version(min_version="12.9")): + return 256 + return 128 + @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: - # BF16 packed math and 256-bit loads; packed scales only (no colwise GEMM swizzle yet). + # BF16 packed math; packed scales only (no colwise GEMM swizzle yet). return ( cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 and cfg.ROWWISE and cfg.COLWISE and not cfg.WITH_GEMM_SWIZZLED_SCALES - and cutlass.target_version(min_version="12.9") ) @staticmethod 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 1304a47dd00..8df894de373 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 @@ -75,8 +75,11 @@ def quantize_rowwise_mxfp8_g2r( # With CHECK_BOUNDS=False DCE will remove it so we don't pay extra cost when we don't use it tile_valid_flags = cute.make_rmem_tensor((NUM_TILES,), Boolean) - # 16 * 2 bytes = 32 bytes which is already the widest - LOAD_WIDTH = ELEMENTS_PER_THREAD * cfg.DTYPE.width + # 16 * 2 bytes = 32 bytes, one 256-bit load where the toolchain supports it + LOAD_WIDTH = min( + ELEMENTS_PER_THREAD * cfg.DTYPE.width, + MXFP8QuantizeG2RRowwise2LaneKernel.max_num_bits_per_copy(), + ) load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), cfg.DTYPE, @@ -179,15 +182,24 @@ class MXFP8QuantizeG2RRowwise2LaneKernel(MXFP8QuantizeKernelBase): _SMALL_TIER_CACHED_PERCENT = 0 _MAX_GRID_Y = 65535 + @staticmethod + @cute.jit + def max_num_bits_per_copy() -> int: + """Widest global load to issue: 256-bit loads (`ld.global.v8.b32`) need PTX ISA 8.8, + i.e. CUDA 12.9. Older toolchains split each one into two 128-bit loads, as + cast_rowwise.cuh / cast_bidim.cuh do via util/ptx.cuh's ld_global_nc_b32x8.""" + if cutlass.const_expr(cutlass.target_version(min_version="12.9")): + return 256 + return 128 + @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: - # BF16 packed math and 256-bit loads; compilers predating those keep the other kernels. + # BF16 packed math. return ( cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 and cfg.ROWWISE and not cfg.COLWISE - and cutlass.target_version(min_version="12.9") ) @staticmethod From 8afd230a0d7455903ac7f24d3e3851fb2b32d5d5 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 13:11:25 -0700 Subject: [PATCH 15/29] also support swizzled scale for g2r bidim Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 63 +++++++++++-------- 1 file changed, 36 insertions(+), 27 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index e0df0abf804..1b9f6ee4efb 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -24,6 +24,7 @@ MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, bf16_pair_magnitude, + derive_swizzled_scale_layout, ) @@ -417,8 +418,12 @@ def quantize_bidimensional_mxfp8_g2r( ) row_scale_atom = cute.make_copy_atom(cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8) + # GEMM-swizzled colwise scales are not contiguous (adjacent columns land 16 bytes apart), + # so they are stored a byte at a time, as the general kernel does. col_scale_atom = cute.make_copy_atom( - cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=ELEMENTS_X_PER_THREAD * 8 + cute.nvgpu.CopyR2GOp(), + SCALE_DTYPE, + num_bits_per_copy=8 if cfg.WITH_GEMM_SWIZZLED_SCALES else ELEMENTS_X_PER_THREAD * 8, ) scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(cfg.FP8_DTYPE) @@ -520,14 +525,8 @@ def max_num_bits_per_copy() -> int: @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: - # BF16 packed math; packed scales only (no colwise GEMM swizzle yet). - return ( - cfg.is_plain_cast() - and cfg.DTYPE is cutlass.BFloat16 - and cfg.ROWWISE - and cfg.COLWISE - and not cfg.WITH_GEMM_SWIZZLED_SCALES - ) + # BF16 packed math. + return cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 and cfg.ROWWISE and cfg.COLWISE @staticmethod @cute.jit @@ -648,14 +647,17 @@ def _launch( gX = cute.make_tensor(mX.iterator, payload_layout) gO_row = cute.make_tensor(mO_row.iterator, payload_layout) gO_col = cute.make_tensor(mO_col.iterator, payload_layout) - gS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout((cute.size(mS_row, mode=[0]), 64), stride=(64, 1)), - ) - gS_col = cute.make_tensor( - mS_col.iterator, - cute.make_layout((cute.size(mS_col, mode=[0]), 2048), stride=(2048, 1)), - ) + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + gS_row, gS_col = derive_swizzled_scale_layout(M, 2048, True, True, mS_row, mS_col) + else: + gS_row = cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 64), stride=(64, 1)), + ) + gS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout((cute.size(mS_col, mode=[0]), 2048), stride=(2048, 1)), + ) self.kernel( gX, gO_row, @@ -669,14 +671,17 @@ def _launch( gX = cute.make_tensor(mX.iterator, payload_layout) gO_row = cute.make_tensor(mO_row.iterator, payload_layout) gO_col = cute.make_tensor(mO_col.iterator, payload_layout) - gS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout((cute.size(mS_row, mode=[0]), 128), stride=(128, 1)), - ) - gS_col = cute.make_tensor( - mS_col.iterator, - cute.make_layout((cute.size(mS_col, mode=[0]), 4096), stride=(4096, 1)), - ) + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + gS_row, gS_col = derive_swizzled_scale_layout(M, 4096, True, True, mS_row, mS_col) + else: + gS_row = cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 128), stride=(128, 1)), + ) + gS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout((cute.size(mS_col, mode=[0]), 4096), stride=(4096, 1)), + ) self.kernel( gX, gO_row, @@ -686,12 +691,16 @@ def _launch( ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) else: + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + gS_row, gS_col = derive_swizzled_scale_layout(M, N, True, True, mS_row, mS_col) + else: + gS_row, gS_col = mS_row, mS_col self.kernel( mX, mO_row, - mS_row, + gS_row, mO_col, - mS_col, + gS_col, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) From 87ab24e659e8e6c73e7b60558e6dc95628f37c44 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 13:20:47 -0700 Subject: [PATCH 16/29] lint Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index 1b9f6ee4efb..e6794283a03 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -153,7 +153,6 @@ def _reduce_partial_colwise_amaxes( # Constexprs CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], - ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], @@ -249,12 +248,8 @@ def _write_colwise_scale( # Miscellaneous cache_policy: Int64, # Constexprs - CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], - ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): """ @@ -462,7 +457,6 @@ def quantize_bidimensional_mxfp8_g2r( sAmaxs_col_scratch, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, DTYPE=cfg.DTYPE, FP8_DTYPE=cfg.FP8_DTYPE, @@ -478,12 +472,8 @@ def quantize_bidimensional_mxfp8_g2r( gS_col_thread, col_scale_atom=col_scale_atom, cache_policy=cache_policy, - CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - DTYPE=cfg.DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, SCALE_DTYPE=SCALE_DTYPE, ) From cae94f577d812a6ae42ea890dca23d9f9956084a Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 21:05:15 +0000 Subject: [PATCH 17/29] add tests Signed-off-by: Kaining Zhong --- tests/jax/test_mxfp8_cutedsl_backend.py | 31 ++++++++++++ .../mxfp8/test_mxfp8_cutedsl_backend.py | 50 ++++++++++++++++++- 2 files changed, 80 insertions(+), 1 deletion(-) diff --git a/tests/jax/test_mxfp8_cutedsl_backend.py b/tests/jax/test_mxfp8_cutedsl_backend.py index 4671ea06869..10f242fb96c 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -283,3 +283,34 @@ def test_dtypes(method, act_type, act_desc, fp8_dtype, in_dtype): in_dtype, fp8_dtype, ) + + +# Cover the specialized kernels' shapes +# Note that JAX only emits linear scales and passes no noop tensor +SPECIALIZED_CASES = { + # G2RRowwise2Lane: its four launch tiers (by output size). 4064x4480 leaves a partial last + # row in the flat 1024-column view. + "g2r_rowwise2lane_tier0": ((1024, 4096), QuantizeLayout.ROWWISE, jnp.bfloat16), + "g2r_rowwise2lane_tier1_flat_tail": ((4064, 4480), QuantizeLayout.ROWWISE, jnp.bfloat16), + "g2r_rowwise2lane_tier2": ((8192, 8192), QuantizeLayout.ROWWISE, jnp.bfloat16), + "g2r_rowwise2lane_tier3": ((16384, 8192), QuantizeLayout.ROWWISE, jnp.bfloat16), + # G2RRowwise1Lane + "g2r_rowwise1lane_fp16": ((1024, 4096), QuantizeLayout.ROWWISE, jnp.float16), + # RegisterBidimensional: static N == 2048 / 4096, the wide tile with and without the + # deep cluster, and the narrow tile (N % 512 != 0). + "g2r_bidim_n2048": ((256, 2048), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + "g2r_bidim_n4096": ((256, 4096), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + "g2r_bidim_wide": ((256, 6144), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + "g2r_bidim_wide_cluster": ((4096, 16384), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + "g2r_bidim_narrow": ((160, 2304), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + # SpecializedBidimensional (g2s): fp16, or bf16 with N % 256 != 0. + "g2s_bidim_fp16": ((256, 1024), QuantizeLayout.ROWWISE_COLWISE, jnp.float16), + "g2s_bidim_bf16_n384": ((256, 384), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), +} + + +@pytest.mark.parametrize( + "shape,q_layout,in_dtype", SPECIALIZED_CASES.values(), ids=SPECIALIZED_CASES.keys() +) +def test_specialized_kernels(shape, q_layout, in_dtype): + run_test_case("CAST_ONLY", *NO_ACT, shape, q_layout, in_dtype, jnp.float8_e4m3fn) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index f409f76993d..84d70563501 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -9,6 +9,7 @@ import ctypes import os +import re import subprocess import sys import textwrap @@ -59,7 +60,8 @@ (96, 128), ] # (block_rows, block_cols): (1,32)=rowwise, (32,1)=colwise, (32,32)=both. -BLOCK_SIZES = [(1, 32), (32, 1), (32, 32)] +ROW, COL, BOTH = (1, 32), (32, 1), (32, 32) +BLOCK_SIZES = [ROW, COL, BOTH] # Every activation the CuTeDSL backend implements; "desc" is the name it goes by in the # config key (Activation in common/util/cutedsl_utils.h). Gated variants are absent because @@ -370,3 +372,49 @@ def test_sizes(swizzled, method, act, block_size, shape): @pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) def test_dtypes(swizzled, method, act, fp8_dtype, in_dtype): run_test_case(method, act, (256, 384), (32, 32), in_dtype, fp8_dtype, swizzled) + + +# Cover the specialized kernels' shapes + +SPECIALIZED_PATH_CASES = [ + # G2RRowwise2Lane: the four launch tiers (by output size) on the flat packed-scale view. + ("g2r_rowwise2lane_tier0", (1024, 4096), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), + # M*N % 1024 != 0: the flat view's last row is partial. + ("g2r_rowwise2lane_tier1_flat_tail", (4064, 4480), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_tier2", (8192, 8192), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_tier3", (16384, 8192), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), + # Swizzled scales use the 2-D grid instead of the flat view. + ("g2r_rowwise2lane_swizzled", (4096, 4096), ROW, torch.bfloat16, True, "G2RRowwise2Lane", 0, False), + # Partial row and column tiles, and scale padding. + ("g2r_rowwise2lane_swizzled_partial", (160, 384), ROW, torch.bfloat16, True, "G2RRowwise2Lane", 0, False), + # Row tiles past grid.y's limit spill into grid.z. + ("g2r_rowwise2lane_swizzled_grid_z", (524352, 128), ROW, torch.bfloat16, True, "G2RRowwise2Lane", 0, False), + # G2RRowwise1Lane: input only 16-byte aligned, or fp16. + ("g2r_rowwise1lane_misaligned_bf16", (1024, 4096), ROW, torch.bfloat16, False, "G2RRowwise1Lane", 8, False), + ("g2r_rowwise1lane_fp16", (1024, 4096), ROW, torch.float16, False, "G2RRowwise1Lane", 0, False), + # G2RBidim: static N == 2048 / 4096, the wide tile with and without the + # deep cluster, and the narrow tile (N % 512 != 0). + ("g2r_bidim_n2048", (256, 2048), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), + ("g2r_bidim_n2048_swizzled", (256, 2048), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), + ("g2r_bidim_n4096_swizzled", (256, 4096), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), + ("g2r_bidim_wide", (256, 6144), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), + ("g2r_bidim_wide_cluster", (4096, 16384), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), + ("g2r_bidim_narrow", (160, 2304), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), + ("g2r_bidim_narrow_swizzled", (160, 2304), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), + # M % 128 != 0: swizzled scale padding. + ("g2r_bidim_wide_swizzled_pad", (4128, 6144), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), + # G2SBidim: fp16, or bf16 with N % 256 != 0. + ("g2s_bidim_fp16", (256, 1024), BOTH, torch.float16, False, "G2SBidim", 0, False), + ("g2s_bidim_bf16_n384", (256, 384), BOTH, torch.bfloat16, False, "G2SBidim", 0, False), +] + + +@pytest.mark.parametrize( + "shape,block_size,in_dtype,swizzled,expected_kernel,input_offset,noop", + [case[1:] for case in SPECIALIZED_PATH_CASES], + ids=[case[0] for case in SPECIALIZED_PATH_CASES], +) +def test_specialized_kernels( + shape, block_size, in_dtype, swizzled +): + run_test_case("CAST_ONLY", IDENTITY, shape, block_size, in_dtype, tex.DType.kFloat8E4M3,swizzled) From f67b1e5945d194d3e90c5209673166a8b78761af Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 14:10:58 -0700 Subject: [PATCH 18/29] fix Signed-off-by: Kaining Zhong --- .../mxfp8/test_mxfp8_cutedsl_backend.py | 66 +++++++++++-------- 1 file changed, 39 insertions(+), 27 deletions(-) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index 84d70563501..8d99ccd7dfa 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -9,7 +9,6 @@ import ctypes import os -import re import subprocess import sys import textwrap @@ -269,14 +268,20 @@ def extract_quantized_output(out, rowwise, columnwise, swizzled): return parts -def run_test_case(method, act, shape, block_size, in_dtype, fp8_dtype, swizzled=False): +def run_test_case( + method, act, shape, block_size, in_dtype, fp8_dtype, swizzled=False, input_offset=0 +): """Assert the CuTeDSL and CUDA backends produce bit-identical outputs for the - same input and config. + same input and config. input_offset (in elements) shifts x's data pointer past the + start of its allocation, to test misaligned inputs. """ M, N = shape rowwise = block_size[1] != 1 columnwise = block_size[0] != 1 x, act_input = generate_inputs(M, N, in_dtype) + if input_offset: + buf = torch.empty(input_offset + M * N, dtype=in_dtype, device="cuda") + x = buf[input_offset:].view(M, N).copy_(x) set_cutedsl_backend(False) out_cuda, dbias_cuda = run_quantize( @@ -378,43 +383,50 @@ def test_dtypes(swizzled, method, act, fp8_dtype, in_dtype): SPECIALIZED_PATH_CASES = [ # G2RRowwise2Lane: the four launch tiers (by output size) on the flat packed-scale view. - ("g2r_rowwise2lane_tier0", (1024, 4096), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_tier0", (1024, 4096), ROW, torch.bfloat16, False, 0), # M*N % 1024 != 0: the flat view's last row is partial. - ("g2r_rowwise2lane_tier1_flat_tail", (4064, 4480), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), - ("g2r_rowwise2lane_tier2", (8192, 8192), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), - ("g2r_rowwise2lane_tier3", (16384, 8192), ROW, torch.bfloat16, False, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_tier1_flat_tail", (4064, 4480), ROW, torch.bfloat16, False, 0), + ("g2r_rowwise2lane_tier2", (8192, 8192), ROW, torch.bfloat16, False, 0), + ("g2r_rowwise2lane_tier3", (16384, 8192), ROW, torch.bfloat16, False, 0), # Swizzled scales use the 2-D grid instead of the flat view. - ("g2r_rowwise2lane_swizzled", (4096, 4096), ROW, torch.bfloat16, True, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_swizzled", (4096, 4096), ROW, torch.bfloat16, True, 0), # Partial row and column tiles, and scale padding. - ("g2r_rowwise2lane_swizzled_partial", (160, 384), ROW, torch.bfloat16, True, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_swizzled_partial", (160, 384), ROW, torch.bfloat16, True, 0), # Row tiles past grid.y's limit spill into grid.z. - ("g2r_rowwise2lane_swizzled_grid_z", (524352, 128), ROW, torch.bfloat16, True, "G2RRowwise2Lane", 0, False), + ("g2r_rowwise2lane_swizzled_grid_z", (524352, 128), ROW, torch.bfloat16, True, 0), # G2RRowwise1Lane: input only 16-byte aligned, or fp16. - ("g2r_rowwise1lane_misaligned_bf16", (1024, 4096), ROW, torch.bfloat16, False, "G2RRowwise1Lane", 8, False), - ("g2r_rowwise1lane_fp16", (1024, 4096), ROW, torch.float16, False, "G2RRowwise1Lane", 0, False), + ("g2r_rowwise1lane_misaligned_bf16", (1024, 4096), ROW, torch.bfloat16, False, 8), + ("g2r_rowwise1lane_fp16", (1024, 4096), ROW, torch.float16, False, 0), # G2RBidim: static N == 2048 / 4096, the wide tile with and without the # deep cluster, and the narrow tile (N % 512 != 0). - ("g2r_bidim_n2048", (256, 2048), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), - ("g2r_bidim_n2048_swizzled", (256, 2048), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), - ("g2r_bidim_n4096_swizzled", (256, 4096), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), - ("g2r_bidim_wide", (256, 6144), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), - ("g2r_bidim_wide_cluster", (4096, 16384), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), - ("g2r_bidim_narrow", (160, 2304), BOTH, torch.bfloat16, False, "G2RBidim", 0, False), - ("g2r_bidim_narrow_swizzled", (160, 2304), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), + ("g2r_bidim_n2048", (256, 2048), BOTH, torch.bfloat16, False, 0), + ("g2r_bidim_n2048_swizzled", (256, 2048), BOTH, torch.bfloat16, True, 0), + ("g2r_bidim_n4096_swizzled", (256, 4096), BOTH, torch.bfloat16, True, 0), + ("g2r_bidim_wide", (256, 6144), BOTH, torch.bfloat16, False, 0), + ("g2r_bidim_wide_cluster", (4096, 16384), BOTH, torch.bfloat16, False, 0), + ("g2r_bidim_narrow", (160, 2304), BOTH, torch.bfloat16, False, 0), + ("g2r_bidim_narrow_swizzled", (160, 2304), BOTH, torch.bfloat16, True, 0), # M % 128 != 0: swizzled scale padding. - ("g2r_bidim_wide_swizzled_pad", (4128, 6144), BOTH, torch.bfloat16, True, "G2RBidim", 0, False), + ("g2r_bidim_wide_swizzled_pad", (4128, 6144), BOTH, torch.bfloat16, True, 0), # G2SBidim: fp16, or bf16 with N % 256 != 0. - ("g2s_bidim_fp16", (256, 1024), BOTH, torch.float16, False, "G2SBidim", 0, False), - ("g2s_bidim_bf16_n384", (256, 384), BOTH, torch.bfloat16, False, "G2SBidim", 0, False), + ("g2s_bidim_fp16", (256, 1024), BOTH, torch.float16, False, 0), + ("g2s_bidim_bf16_n384", (256, 384), BOTH, torch.bfloat16, False, 0), ] @pytest.mark.parametrize( - "shape,block_size,in_dtype,swizzled,expected_kernel,input_offset,noop", + "shape,block_size,in_dtype,swizzled,input_offset", [case[1:] for case in SPECIALIZED_PATH_CASES], ids=[case[0] for case in SPECIALIZED_PATH_CASES], ) -def test_specialized_kernels( - shape, block_size, in_dtype, swizzled -): - run_test_case("CAST_ONLY", IDENTITY, shape, block_size, in_dtype, tex.DType.kFloat8E4M3,swizzled) +def test_specialized_kernels(shape, block_size, in_dtype, swizzled, input_offset): + run_test_case( + "CAST_ONLY", + IDENTITY, + shape, + block_size, + in_dtype, + tex.DType.kFloat8E4M3, + swizzled, + input_offset, + ) From 177647e1ca58620db7dfd2a50bda4f0171ed0bf9 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 12 Oct 2026 15:24:59 -0700 Subject: [PATCH 19/29] nit Signed-off-by: Kaining Zhong --- tests/cpp/operator/test_cast_mxfp8.cu | 3 +++ tests/jax/test_mxfp8_cutedsl_backend.py | 5 ++--- .../pytorch/mxfp8/test_mxfp8_cutedsl_backend.py | 3 +-- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 16 ++++------------ .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 13 ++----------- .../CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py | 7 +++---- transformer_engine/common/CuTeDSL/utils.py | 10 ++++++++++ 7 files changed, 25 insertions(+), 32 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8.cu b/tests/cpp/operator/test_cast_mxfp8.cu index 9368d4030b5..bc62d333b5f 100644 --- a/tests/cpp/operator/test_cast_mxfp8.cu +++ b/tests/cpp/operator/test_cast_mxfp8.cu @@ -699,6 +699,9 @@ std::vector> matrix_sizes = { {1, 16}, {16, 48}, {128, 128}, + // This case will cover cast_rowwise's last row being partial case, in which + // the kernel will reshape the input tensor with last dim being 1024 so the last row + // is not a full row due to 993*512 % 1024 != 0 {993, 512}, {1024}, {8, 32, 1024}, diff --git a/tests/jax/test_mxfp8_cutedsl_backend.py b/tests/jax/test_mxfp8_cutedsl_backend.py index 10f242fb96c..bd492c967a1 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -288,10 +288,9 @@ def test_dtypes(method, act_type, act_desc, fp8_dtype, in_dtype): # Cover the specialized kernels' shapes # Note that JAX only emits linear scales and passes no noop tensor SPECIALIZED_CASES = { - # G2RRowwise2Lane: its four launch tiers (by output size). 4064x4480 leaves a partial last - # row in the flat 1024-column view. + # G2RRowwise2Lane: its four launch tiers (by output size). "g2r_rowwise2lane_tier0": ((1024, 4096), QuantizeLayout.ROWWISE, jnp.bfloat16), - "g2r_rowwise2lane_tier1_flat_tail": ((4064, 4480), QuantizeLayout.ROWWISE, jnp.bfloat16), + "g2r_rowwise2lane_tier1": ((4064, 4480), QuantizeLayout.ROWWISE, jnp.bfloat16), "g2r_rowwise2lane_tier2": ((8192, 8192), QuantizeLayout.ROWWISE, jnp.bfloat16), "g2r_rowwise2lane_tier3": ((16384, 8192), QuantizeLayout.ROWWISE, jnp.bfloat16), # G2RRowwise1Lane diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index 8d99ccd7dfa..a886b63c088 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -384,8 +384,7 @@ def test_dtypes(swizzled, method, act, fp8_dtype, in_dtype): SPECIALIZED_PATH_CASES = [ # G2RRowwise2Lane: the four launch tiers (by output size) on the flat packed-scale view. ("g2r_rowwise2lane_tier0", (1024, 4096), ROW, torch.bfloat16, False, 0), - # M*N % 1024 != 0: the flat view's last row is partial. - ("g2r_rowwise2lane_tier1_flat_tail", (4064, 4480), ROW, torch.bfloat16, False, 0), + ("g2r_rowwise2lane_tier1", (4064, 4480), ROW, torch.bfloat16, False, 0), ("g2r_rowwise2lane_tier2", (8192, 8192), ROW, torch.bfloat16, False, 0), ("g2r_rowwise2lane_tier3", (16384, 8192), ROW, torch.bfloat16, False, 0), # Swizzled scales use the 2-D grid instead of the flat view. diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index e6794283a03..b636d04edfb 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -17,6 +17,7 @@ bf16_mx_scale_reciprocal_pair, bf16_mx_scale_bytes_pair, create_l2_policy, + max_num_bits_per_copy_ld_global, ) from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 from .quantize_mxfp8_common import ( @@ -400,7 +401,7 @@ def quantize_bidimensional_mxfp8_g2r( cfg.DTYPE, num_bits_per_copy=min( ELEMENTS_X_PER_THREAD * cfg.DTYPE.width, - MXFP8QuantizeRegisterBidimensionalKernel.max_num_bits_per_copy(), + max_num_bits_per_copy_ld_global(), ), # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using # invariant=True here will drop the cache policy @@ -497,22 +498,13 @@ class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): TILE_COLS = _CTA_THREADS_X * ELEMENTS_X_PER_THREAD = 256 or 512. The 32-row tile contains complete columnwise MXFP8 blocks; shared memory combines the partial column maxima from the eight warps. Requires aligned - BF16 input, M divisible by 32, N divisible by 256 and plain scale layouts. + BF16 input, M divisible by 32 and N divisible by 256; writes packed or GEMM-swizzled + scales. """ _CTA_THREADS_Y = 8 _CTA_THREADS_X = 32 - @staticmethod - @cute.jit - def max_num_bits_per_copy() -> int: - """Widest global load to issue: 256-bit loads (`ld.global.v8.b32`) need PTX ISA 8.8, - i.e. CUDA 12.9. Older toolchains split each one into two 128-bit loads, as - cast_rowwise.cuh / cast_bidim.cuh do via util/ptx.cuh's ld_global_nc_b32x8.""" - if cutlass.const_expr(cutlass.target_version(min_version="12.9")): - return 256 - return 128 - @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. 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 8df894de373..68c69040cfc 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 @@ -15,6 +15,7 @@ abs_max_x2_bf16, create_l2_policy, exp2_bf16x2_rcp, + max_num_bits_per_copy_ld_global, ) from transformer_engine.common.CuTeDSL.utils_fp8 import ( cvt_f32_to_fp8e8m0fnu, @@ -78,7 +79,7 @@ def quantize_rowwise_mxfp8_g2r( # 16 * 2 bytes = 32 bytes, one 256-bit load where the toolchain supports it LOAD_WIDTH = min( ELEMENTS_PER_THREAD * cfg.DTYPE.width, - MXFP8QuantizeG2RRowwise2LaneKernel.max_num_bits_per_copy(), + max_num_bits_per_copy_ld_global(), ) load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), @@ -182,16 +183,6 @@ class MXFP8QuantizeG2RRowwise2LaneKernel(MXFP8QuantizeKernelBase): _SMALL_TIER_CACHED_PERCENT = 0 _MAX_GRID_Y = 65535 - @staticmethod - @cute.jit - def max_num_bits_per_copy() -> int: - """Widest global load to issue: 256-bit loads (`ld.global.v8.b32`) need PTX ISA 8.8, - i.e. CUDA 12.9. Older toolchains split each one into two 128-bit loads, as - cast_rowwise.cuh / cast_bidim.cuh do via util/ptx.cuh's ld_global_nc_b32x8.""" - if cutlass.const_expr(cutlass.target_version(min_version="12.9")): - return 256 - return 128 - @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py index a2b4f5e55ff..aca02a253b1 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py @@ -7,10 +7,9 @@ Replicates the core logic of quantize_mxfp8.cuh: given a 2D tensor of BF16/FP16 values, quantize to MXFP8 format (FP8E4M3 data + E8M0 per-block scales). -BF16 cast-only calls with aligned input use the register-resident ports of -specialized/cast_rowwise.cuh and specialized/cast_bidim.cuh. Other configurations -use the one-lane rowwise, staged bidimensional or general kernels. All paths share the -TVM-FFI entrypoint. +This is the general kernel: it handles every config. MXFP8QuantizeEntry in +quantize_mxfp8_common.py dispatches plain casts to the specialized kernels when they +apply and falls back to this kernel otherwise. Matches the C++ kernel's tile dimensions and thread layout: CHUNK_DIM_Y = 64, CHUNK_DIM_X = 64, THREADS_PER_CTA = 64 diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index 48f3a956544..d5ffd4abfe9 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -27,6 +27,16 @@ logger = logging.getLogger("transformer_engine.cutedsl.utils") +@cute.jit +def max_num_bits_per_copy_ld_global() -> int: + """Widest global load to issue: 256-bit loads (`ld.global.v8.b32`) need PTX ISA 8.8, + i.e. CUDA 12.9. Older toolchains split each one into two 128-bit loads, as + cast_rowwise.cuh / cast_bidim.cuh do via util/ptx.cuh's ld_global_nc_b32x8.""" + if cutlass.const_expr(cutlass.target_version(min_version="12.9")): + return 256 + return 128 + + def device_compute_capability(device_id: Optional[int] = None) -> tuple: """(major, minor) compute capability of a CUDA device (current by default), or (0, 0) if it can't be queried.""" from cuda.core import Device # pylint: disable=no-name-in-module From 8f4c54d6b124fa863a3282d28496c709154191fd Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 13 Oct 2026 22:47:11 +0000 Subject: [PATCH 20/29] refactor Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_bidimensional.py | 281 +++++++++------ .../mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py | 329 +++++++++-------- .../cast/mxfp8/quantize_mxfp8_general.py | 335 +++++++++--------- 3 files changed, 532 insertions(+), 413 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py index b636d04edfb..75ad5799cbf 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py @@ -276,7 +276,7 @@ def _write_colwise_scale( @cute.jit -def _quantize_colwise( +def _apply_colwise_scales( rX: cute.Tensor, rS_col_rcp: cute.Tensor, gO_col_thread: cute.Tensor, @@ -318,26 +318,100 @@ def _quantize_colwise( cute.copy(store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy) +@cute.jit +def _quantize_colwise( + rX: cute.Tensor, + gO_col_thread: cute.Tensor, + gS_col_thread: cute.Tensor, + # Atoms + output_store_atom: cute.CopyAtom, + col_scale_store_atom: cute.CopyAtom, + # Miscellaneous + cache_policy: Int64, + scale_and_pack: cutlass.Constexpr[Callable[..., Uint32]], + # Constexprs + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], +): + """Reduce columnwise amaxes across the CTA, write the columnwise scales, and quantize rX with them.""" + tidx, _, _ = cute.arch.thread_idx() + + sAmaxs_col_scratch = _gather_partial_colwise_amaxes( + rX, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + DTYPE=DTYPE, + ) + + sS_col = _reduce_partial_colwise_amaxes( + sAmaxs_col_scratch, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + ) + + # Copy the reduced columnwise scale reciprocals to RMEM. + # Threads that quantize different rows share the same columnwise scale + rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), DTYPE) + cute.autovec_copy(sS_col[None, tidx % CTA_THREADS_X], rS_col_rcp) + + _write_colwise_scale( + rS_col_rcp, + gS_col_thread, + col_scale_atom=col_scale_store_atom, + cache_policy=cache_policy, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + ) + + _apply_colwise_scales( + rX, + rS_col_rcp, + gO_col_thread, + store_atom=output_store_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + FP8_DTYPE=FP8_DTYPE, + ) + + @cute.jit def quantize_bidimensional_mxfp8_g2r( gX_tile: cute.Tensor, - gO_row_tile: cute.Tensor, - gS_row_tile: cute.Tensor, - gO_col_tile: cute.Tensor, - gS_col_tile: cute.Tensor, - cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + gO_row_tile: Optional[cute.Tensor], + gS_row_tile: Optional[cute.Tensor], + gO_col_tile: Optional[cute.Tensor], + gS_col_tile: Optional[cute.Tensor], + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], cache_policy: Int64, ): - """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile in both directions. + """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile in the directions enabled by ROWWISE / COLWISE. TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD = 32, so every columnwise MXFP8 block is contained in the tile. Each warp owns ELEMENTS_Y_PER_THREAD consecutive rows, and each lane owns ELEMENTS_X_PER_THREAD adjacent columns. - Load the tile once, then reuse its registers for both quantization passes. + Load the tile once, then reuse its registers for each enabled quantization pass. + The tensors of a disabled direction are unused and may be None. """ # Shape of the tile TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD @@ -359,48 +433,15 @@ def quantize_bidimensional_mxfp8_g2r( ) gX_thread = cute.composition(gX_tile, tv_layout)[tidx, None] - gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] - gO_col_thread = cute.composition(gO_col_tile, tv_layout)[tidx, None] - - # Group CTA threads by THREADS_X_PER_MX_BLOCK since they share the same rowwise scale for a MXFP8 block. - # And assign stride 0 to that submode - tv_row_scale = cute.make_layout( - ( - (THREADS_X_PER_MX_BLOCK, (CTA_THREADS_X // THREADS_X_PER_MX_BLOCK, CTA_THREADS_Y)), - ELEMENTS_Y_PER_THREAD, - ), - stride=( - (0, (TILE_ROWS, ELEMENTS_Y_PER_THREAD)), - 1, - ), - ) - gS_row_thread = cute.composition(gS_row_tile, tv_row_scale)[tidx, None] - - # Group CTA threads by THREADS_Y_PER_MX_BLOCK since they share the same columnwise scale for a MXFP8 block. - # And assign stride 0 to that submode - tv_col_scale = cute.make_layout( - ( - (CTA_THREADS_X, (THREADS_Y_PER_MX_BLOCK, CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK)), - ELEMENTS_X_PER_THREAD, - ), - stride=( - (ELEMENTS_X_PER_THREAD, (0, 1)), - 1, - ), - ) - # If CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK == 1, this will help compiler derive alignment - # Otherwise coalesce will not do anything which is still fine - tv_col_scale = cute.coalesce(tv_col_scale, target_profile=(1, 1)) - gS_col_thread = cute.composition(gS_col_tile, tv_col_scale)[tidx, None] # Input fragment contains all rows it handles; output fragment contains only one row and we write it back row by row - rX = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), cfg.DTYPE) + rX = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), DTYPE) load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), - cfg.DTYPE, + DTYPE, num_bits_per_copy=min( - ELEMENTS_X_PER_THREAD * cfg.DTYPE.width, + ELEMENTS_X_PER_THREAD * DTYPE.width, max_num_bits_per_copy_ld_global(), ), # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using @@ -409,20 +450,11 @@ def quantize_bidimensional_mxfp8_g2r( ) store_atom = cute.make_copy_atom( cute.nvgpu.CopyR2GOp(), - cfg.FP8_DTYPE, - num_bits_per_copy=ELEMENTS_X_PER_THREAD * cfg.FP8_DTYPE.width, + FP8_DTYPE, + num_bits_per_copy=ELEMENTS_X_PER_THREAD * FP8_DTYPE.width, ) - row_scale_atom = cute.make_copy_atom(cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8) - # GEMM-swizzled colwise scales are not contiguous (adjacent columns land 16 bytes apart), - # so they are stored a byte at a time, as the general kernel does. - col_scale_atom = cute.make_copy_atom( - cute.nvgpu.CopyR2GOp(), - SCALE_DTYPE, - num_bits_per_copy=8 if cfg.WITH_GEMM_SWIZZLED_SCALES else ELEMENTS_X_PER_THREAD * 8, - ) - - scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(cfg.FP8_DTYPE) + scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(FP8_DTYPE) # Copy all rows a thread sweeps into the RMEM fragment at once for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): @@ -430,65 +462,86 @@ def quantize_bidimensional_mxfp8_g2r( load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy ) - _quantize_rowwise( - rX, - gO_row_thread, - gS_row_thread, - output_store_atom=store_atom, - row_scale_store_atom=row_scale_atom, - cache_policy=cache_policy, - scale_and_pack=scale_and_pack, - THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - SCALE_DTYPE=SCALE_DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, - ) - - sAmaxs_col_scratch = _gather_partial_colwise_amaxes( - rX, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - DTYPE=cfg.DTYPE, - ) + if cutlass.const_expr(ROWWISE): + gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] + + # Group CTA threads by THREADS_X_PER_MX_BLOCK since they share the same rowwise scale for a MXFP8 block. + # And assign stride 0 to that submode + tv_row_scale = cute.make_layout( + ( + (THREADS_X_PER_MX_BLOCK, (CTA_THREADS_X // THREADS_X_PER_MX_BLOCK, CTA_THREADS_Y)), + ELEMENTS_Y_PER_THREAD, + ), + stride=( + (0, (TILE_ROWS, ELEMENTS_Y_PER_THREAD)), + 1, + ), + ) + gS_row_thread = cute.composition(gS_row_tile, tv_row_scale)[tidx, None] - sS_col = _reduce_partial_colwise_amaxes( - sAmaxs_col_scratch, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - DTYPE=cfg.DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, - ) + row_scale_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8 + ) - # Copy the reduced columnwise scale reciprocals to RMEM. - # Threads that quantize different rows share the same columnwise scale - rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), cfg.DTYPE) - cute.autovec_copy(sS_col[None, tidx % CTA_THREADS_X], rS_col_rcp) + _quantize_rowwise( + rX, + gO_row_thread, + gS_row_thread, + output_store_atom=store_atom, + row_scale_store_atom=row_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=FP8_DTYPE, + ) - _write_colwise_scale( - rS_col_rcp, - gS_col_thread, - col_scale_atom=col_scale_atom, - cache_policy=cache_policy, - CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - SCALE_DTYPE=SCALE_DTYPE, - ) + if cutlass.const_expr(COLWISE): + gO_col_thread = cute.composition(gO_col_tile, tv_layout)[tidx, None] + + # Group CTA threads by THREADS_Y_PER_MX_BLOCK since they share the same columnwise scale for a MXFP8 block. + # And assign stride 0 to that submode + tv_col_scale = cute.make_layout( + ( + (CTA_THREADS_X, (THREADS_Y_PER_MX_BLOCK, CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK)), + ELEMENTS_X_PER_THREAD, + ), + stride=( + (ELEMENTS_X_PER_THREAD, (0, 1)), + 1, + ), + ) + # If CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK == 1, this will help compiler derive alignment + # Otherwise coalesce will not do anything which is still fine + tv_col_scale = cute.coalesce(tv_col_scale, target_profile=(1, 1)) + gS_col_thread = cute.composition(gS_col_tile, tv_col_scale)[tidx, None] + + # GEMM-swizzled colwise scales are not contiguous (adjacent columns land 16 bytes apart), + # so they are stored a byte at a time, as the general kernel does. + col_scale_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + SCALE_DTYPE, + num_bits_per_copy=8 if WITH_GEMM_SWIZZLED_SCALES else ELEMENTS_X_PER_THREAD * 8, + ) - _quantize_colwise( - rX, - rS_col_rcp, - gO_col_thread, - store_atom=store_atom, - cache_policy=cache_policy, - scale_and_pack=scale_and_pack, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - FP8_DTYPE=cfg.FP8_DTYPE, - ) + _quantize_colwise( + rX, + gO_col_thread, + gS_col_thread, + output_store_atom=store_atom, + col_scale_store_atom=col_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + ) class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): @@ -742,7 +795,11 @@ def _kernel_main( gS_row_tile, gO_col_tile, gS_col_tile, - self.cfg, + ROWWISE=self.cfg.ROWWISE, + COLWISE=self.cfg.COLWISE, + DTYPE=self.cfg.DTYPE, + FP8_DTYPE=self.cfg.FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=self.cfg.WITH_GEMM_SWIZZLED_SCALES, CTA_THREADS_Y=self._CTA_THREADS_Y, CTA_THREADS_X=self._CTA_THREADS_X, ELEMENTS_Y_PER_THREAD=self._ELEMENTS_Y_PER_THREAD, 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 fd76e039fa6..1c005c22fdb 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 @@ -39,6 +39,172 @@ ) +@cute.jit +def _flush_scales_to_gmem( + sScale: cute.Tensor, + gS_tile: cute.Tensor, + cta_row: Int32, + cta_col: Int32, + M: Int32, + padded_cols: Int32, + CTA_Y: cutlass.Constexpr[int], + CTA_X: cutlass.Constexpr[int], + width: cutlass.Constexpr[int], +): + """Flush the staged (CTA_Y, CTA_X) scale tile to gmem with vectorized stores.""" + tidx, _, _ = cute.arch.thread_idx() + # Previously each threads has 1 byte, but now we are doing vectorized store, + # which means only a subset of threads will need to issue the store while other threads are not used. + active_threads = CTA_X // width + _, tv_flush = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, active_threads), stride=(active_threads, 1)), + val_layout=cute.make_layout((1, width), stride=(width, 1)), + ) + # We only need to use a subset of threads with shape (CTA_Y, active_threads) to write + # so if the thread is outside of this subset, it will remain inactive + if tidx < CTA_Y * active_threads: + # Absolute position of the scale vector to write in the GMEM buffer + thread_y = cta_row * CTA_Y + tidx // active_threads + thread_x = cta_col * CTA_X + (tidx % active_threads) * width + if thread_y < M and thread_x < padded_cols: + cute.autovec_copy( + cute.composition(sScale, tv_flush)[tidx, None], + cute.composition(gS_tile, tv_flush)[tidx, None], + ) + + +@cute.jit +def quantize_rowwise_mxfp8_g2r_1lane( + gX_tile: cute.Tensor, + gO_tile: cute.Tensor, + gS_tile: cute.Tensor, + M: Int32, + N: Int32, + padded_cols: Int32, + cta_row: Int32, + cta_col: Int32, + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + MAX_NORM_RCP: cutlass.Constexpr[float], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], + TILE_ROWS: cutlass.Constexpr[int], + TILE_COLS: cutlass.Constexpr[int], + STASH_SCALE_TO_SMEM: cutlass.Constexpr[bool], +): + """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile rowwise with TILE_ROWS * TILE_COLS / 32 threads. + + Each thread handles one 32-element MXFP8 chunk (= one scale block). (cta_row, cta_col) is the + tile coordinate, used with M and N for bounds checks; padded_cols is the column count of the + (padded) scale tensor that gS_tile is a tile of. + """ + tidx, _, _ = cute.arch.thread_idx() + + # The CTA threads are grouped as (CTA_Y, CTA_X), so they cover a (TILE_ROWS, TILE_COLS) + # input tile and the matching (CTA_Y, CTA_X) scale tile. + CTA_Y = TILE_ROWS + CTA_X = TILE_COLS // MXFP8_BLOCK_SCALING_SIZE + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), + val_layout=cute.make_layout( + (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) + ), + ) + _, tv_layout_scale = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), + val_layout=cute.make_layout((1, 1), stride=(1, 1)), + ) + + gX_thread = cute.composition(gX_tile, tv_layout)[tidx, None] + gO_thread = cute.composition(gO_tile, tv_layout)[tidx, None] + gS_thread = cute.composition(gS_tile, tv_layout_scale)[tidx, None] + + rX_thread = cute.make_rmem_tensor( + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), + dtype=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 + mul_cvt4 = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, FP8_DTYPE) + rX_i32 = cute.make_tensor( + cute.recast_ptr(rX_thread.iterator, dtype=Int32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), + ) + rO_thread = cute.make_rmem_tensor( + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), + dtype=FP8_DTYPE, + ) + rO_u32 = cute.make_tensor( + cute.recast_ptr(rO_thread.iterator, dtype=Uint32), + cute.make_layout( + (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) + ), # Unit is Uint32, divide by 4 here + ) + + sS_thread = None + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + + @cute.struct + class SharedStorage: + buf: cute.struct.Align[cute.struct.MemRange[Float8E8M0FNU, CTA_Y * CTA_X], 16] + + storage = cutlass.memory.SmemAllocator().allocate(SharedStorage) + sScale = storage.buf.get_tensor(cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1))) + # sScale is (CTA_Y, CTA_X):(CTA_X, 1), which is the same layout as tv_layout_scale + # so sS_thread is really just a 1-element e8m0 buffer for this thread's scale. + sS_thread = cute.composition(sScale, tv_layout_scale)[tidx, None] + # Zero first so padding columns (cols past N/32 in the padded scale + # matrix) flush as 0 and we never read uninitialized smem. + # Raw 0x00: e8m0 has no zero value to convert from. + sS_thread[0] = Uint8(0).bitcast(Float8E8M0FNU) + cute.arch.sync_threads() + + row = cta_row * TILE_ROWS + tidx // CTA_X + col = cta_col * TILE_COLS + (tidx % CTA_X) * MXFP8_BLOCK_SCALING_SIZE + if row < M and col < N: + cute.autovec_copy(gX_thread, rX_thread) + amax_2x = rX_i32[0] + # Skip the first iteration because we assigned rX_i32[0] to amax_2x already + for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE // 2): + amax_2x = abs_max_x2(amax_2x, rX_i32[i]) + amax = cute.arch.fmax(fabs_f32(x2_lo_to_f32(amax_2x)), fabs_f32(x2_hi_to_f32(amax_2x))) + + biased_exp = cvt_f32_to_fp8e8m0fnu(amax * MAX_NORM_RCP) + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + sS_thread[0] = biased_exp + else: + gS_thread[0] = biased_exp + + # Rescale + FP8 cast, 4 elements per fused mul_cvt (one uint32 out), + # then a vectorized store. Mirrors CUDA's _use_cvt_4x path. + inv_scale = exp2f_rcp(biased_exp) + scale_2x = pack_f32x2(inv_scale, inv_scale) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): + rO_u32[i] = mul_cvt4(rX_i32[2 * i], rX_i32[2 * i + 1], scale_2x) + cute.autovec_copy(rO_thread, gO_thread) + + # Cooperative wide flush of the staged scales where padding columns flush as 0. + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + cute.arch.sync_threads() + if cutlass.const_expr(WITH_GEMM_SWIZZLED_SCALES): + # Swizzled rowwise scale layout is ((32, 4, num_tiles_M), (4, num_tiles_SC)):((16, 4, num_tiles_SC * 512), (1, 512)) + # so we have at most 4 elements continuous which is our vectorized store width + _flush_scales_to_gmem( + sScale, gS_tile, cta_row, cta_col, M, padded_cols, CTA_Y, CTA_X, 4 + ) + elif padded_cols % 16 == 0: + # If columns is divisible by 16, use 16 bytes as the vectorized store width + _flush_scales_to_gmem( + sScale, gS_tile, cta_row, cta_col, M, padded_cols, CTA_Y, CTA_X, 16 + ) + else: + # Otherwise use 4 bytes as the vectorized store width. + # Note our fake tensor requires 4 divisibility so this is enforced as long as you can get here + _flush_scales_to_gmem( + sScale, gS_tile, cta_row, cta_col, M, padded_cols, CTA_Y, CTA_X, 4 + ) + + class MXFP8QuantizeG2RRowwise1LaneKernel(MXFP8QuantizeKernelBase): """Specialized cast-only ROWWISE-only MXFP8 kernel. @@ -115,7 +281,6 @@ def __call__( mX, mO_row, mS_row, - mX.element_type, ).launch(grid=grid, block=block, stream=stream) @cute.kernel @@ -124,11 +289,10 @@ def kernel( mX: cute.Tensor, mO_row: cute.Tensor, mS_row: cute.Tensor, - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): """Device entry. Like the CUDA specialized cast kernels it does not consume the cast_noop flag: the dispatcher routes calls with one to the general kernel.""" - self._kernel_main(mX, mO_row, mS_row, DTYPE) + self._kernel_main(mX, mO_row, mS_row) @cute.jit def _kernel_main( @@ -136,151 +300,34 @@ def _kernel_main( mX: cute.Tensor, mO_row: cute.Tensor, mS_row: cute.Tensor, - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], ): """Device entry for the specialized rowwise-only cast kernel (vectorized global loads/stores, no TMA).""" - tidx, _, _ = cute.arch.thread_idx() bidx, bidy, _ = cute.arch.block_idx() M = mX.shape[0] N = mX.shape[1] - # Each thread handles one 32-element MXFP8 chunk (= one scale block). - # The 128 threads in the CTA are grouped as (4, 32), so they cover a - # (4, 1024) input tile and the matching (4, 32) scale tile. - CTA_Y = self._TILE_ROWS - CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE - tiler, tv_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), - val_layout=cute.make_layout( - (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) - ), - ) - tiler_scale, tv_layout_scale = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), - val_layout=cute.make_layout((1, 1), stride=(1, 1)), - ) + TILER = (self._TILE_ROWS, self._TILE_COLS) + SCALE_TILER = (self._TILE_ROWS, self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + gX_tile = cute.local_tile(mX, TILER, (bidy, bidx)) + gO_tile = cute.local_tile(mO_row, TILER, (bidy, bidx)) + gS_tile = cute.local_tile(mS_row, SCALE_TILER, (bidy, bidx)) - # Select the tile that belongs to this CTA, then the fragment per thread. - mX_tile = cute.local_tile(mX, tiler, (bidy, bidx)) - mO_tile = cute.local_tile(mO_row, tiler, (bidy, bidx)) - mS_tile = cute.local_tile(mS_row, tiler_scale, (bidy, bidx)) - mX_thread = cute.composition(mX_tile, tv_layout)[tidx, None] - mO_thread = cute.composition(mO_tile, tv_layout)[tidx, None] - mS_thread = cute.composition(mS_tile, tv_layout_scale)[tidx, None] - - rX_thread = cute.make_rmem_tensor( - cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), - dtype=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 - mul_cvt4 = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, self.cfg.FP8_DTYPE) - rX_i32 = cute.make_tensor( - cute.recast_ptr(rX_thread.iterator, dtype=Int32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), - ) - rO_thread = cute.make_rmem_tensor( - cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), - dtype=self.cfg.FP8_DTYPE, - ) - rO_u32 = cute.make_tensor( - cute.recast_ptr(rO_thread.iterator, dtype=Uint32), - cute.make_layout( - (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) - ), # Unit is Uint32, divide by 4 here - ) - - sS_thread = None - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - - @cute.struct - class SharedStorage: - buf: cute.struct.Align[cute.struct.MemRange[Float8E8M0FNU, CTA_Y * CTA_X], 16] - - storage = cutlass.memory.SmemAllocator().allocate(SharedStorage) - sScale = storage.buf.get_tensor(cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1))) - # sScale is (CTA_Y, CTA_X):(CTA_X, 1), which is the same layout as tv_layout_scale - # so sS_thread is really just a 1-element e8m0 buffer for this thread's scale. - sS_thread = cute.composition(sScale, tv_layout_scale)[tidx, None] - # Zero first so padding columns (cols past N/32 in the padded scale - # matrix) flush as 0 and we never read uninitialized smem. - # Raw 0x00: e8m0 has no zero value to convert from. - sS_thread[0] = Uint8(0).bitcast(Float8E8M0FNU) - cute.arch.sync_threads() - - row = bidy * self._TILE_ROWS + tidx // CTA_X - col = bidx * self._TILE_COLS + (tidx % CTA_X) * MXFP8_BLOCK_SCALING_SIZE - if row < M and col < N: - cute.autovec_copy(mX_thread, rX_thread) - amax_2x = rX_i32[0] - # Skip the first iteration because we assigned rX_i32[0] to amax_2x already - for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE // 2): - amax_2x = abs_max_x2(amax_2x, rX_i32[i]) - amax = cute.arch.fmax(fabs_f32(x2_lo_to_f32(amax_2x)), fabs_f32(x2_hi_to_f32(amax_2x))) - - biased_exp = cvt_f32_to_fp8e8m0fnu(amax * self.cfg.MAX_NORM_RCP) - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - sS_thread[0] = biased_exp - else: - mS_thread[0] = biased_exp - - # Rescale + FP8 cast, 4 elements per fused mul_cvt (one uint32 out), - # then a vectorized store. Mirrors CUDA's _use_cvt_4x path. - inv_scale = exp2f_rcp(biased_exp) - scale_2x = pack_f32x2(inv_scale, inv_scale) - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): - rO_u32[i] = mul_cvt4(rX_i32[2 * i], rX_i32[2 * i + 1], scale_2x) - cute.autovec_copy(rO_thread, mO_thread) - - # Cooperative wide flush of the staged scales where padding columns flush as 0. - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - cute.arch.sync_threads() + quantize_rowwise_mxfp8_g2r_1lane( + gX_tile, + gO_tile, + gS_tile, + M, + N, # Use cute.size instead of .shape[1] because under the swizzled layout mode 1 is # a nested tuple ((4, num_tiles_SC)), not a plain int. - padded_cols = cute.size(mS_row, mode=[1]) - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - # Swizzled rowwise scale layout is ((32, 4, num_tiles_M), (4, num_tiles_SC)):((16, 4, num_tiles_SC * 512), (1, 512)) - # so we have at most 4 elements continuous which is our vectorized store width - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) - elif padded_cols % 16 == 0: - # If columns is divisible by 16, use 16 bytes as the vectorized store width - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 16) - else: - # Otherwise use 4 bytes as the vectorized store width. - # Note our fake tensor requires 4 divisibility so this is enforced as long as you can get here - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) - - @cute.jit - def _flush_scales_to_gmem( - self, - sScale: cute.Tensor, - mS_tile: cute.Tensor, - tidx: Int32, - bidx: Int32, - bidy: Int32, - M: Int32, - padded_cols: Int32, - width: cutlass.Constexpr[int], - ): - """Flush the staged (CTA_Y, CTA_X) scale tile to gmem with vectorized stores.""" - CTA_Y = self._TILE_ROWS - CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE - # Previously each threads has 1 byte, but now we are doing vectorized store, - # which means only a subset of threads will need to issue the store while other threads are not used. - active_threads = CTA_X // width - _, tv_flush = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, active_threads), stride=(active_threads, 1)), - val_layout=cute.make_layout((1, width), stride=(width, 1)), + cute.size(mS_row, mode=[1]), + bidy, + bidx, + DTYPE=self.cfg.DTYPE, + FP8_DTYPE=self.cfg.FP8_DTYPE, + MAX_NORM_RCP=self.cfg.MAX_NORM_RCP, + WITH_GEMM_SWIZZLED_SCALES=self.cfg.WITH_GEMM_SWIZZLED_SCALES, + TILE_ROWS=self._TILE_ROWS, + TILE_COLS=self._TILE_COLS, + STASH_SCALE_TO_SMEM=self._STASH_SCALE_TO_SMEM, ) - # We only need to use a subset of threads with shape (CTA_Y, active_threads) to write - # so if the thread is outside of this subset, it will remain inactive - if tidx < CTA_Y * active_threads: - # Absolute position of the scale vector to write in the GMEM buffer - thread_y = bidy * CTA_Y + tidx // active_threads - thread_x = bidx * CTA_X + (tidx % active_threads) * width - if thread_y < M and thread_x < padded_cols: - cute.autovec_copy( - cute.composition(sScale, tv_flush)[tidx, None], - cute.composition(mS_tile, tv_flush)[tidx, None], - ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py index aca02a253b1..d74d4ca824b 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py @@ -91,9 +91,6 @@ def quantize_rowwise_mxfp8( FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], TILE_X: cutlass.Constexpr[int], TILE_Y: cutlass.Constexpr[int], - WAVES: cutlass.Constexpr[int], - THREADS_PER_BANK: cutlass.Constexpr[int], - PACK_SIZE: cutlass.Constexpr[int], SKIP_INPUT_MASKING: cutlass.Constexpr[bool] = False, SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool] = False, WITH_ACT: cutlass.Constexpr[bool] = False, @@ -104,6 +101,14 @@ def quantize_rowwise_mxfp8( """Quantize one SMEM tile rowwise to MXFP8 (per-row 32-elt block scales); returns the tile amax.""" tidx, _, _ = cute.arch.thread_idx() + # Vectorised SMEM access constants for bank-conflict avoidance. The loads and the fused 4-wide + # mul_cvt below assume PACK_SIZE == 4, so these are not tunable. + PACK_SIZE = 4 # Elements per vector load + # Each thread reads 8 waves with each wave reads 4 packed bf16, so it reads a whole MXFP8 block in total + WAVES = MXFP8_BLOCK_SCALING_SIZE // PACK_SIZE + TOTAL_BANKS_WIDTH = 32 * 4 # 32 banks × 4 bytes, in bytes + THREADS_PER_BANK = TOTAL_BANKS_WIDTH // MXFP8_BLOCK_SCALING_SIZE # 4 threads per bank + CTA_THREADS_Y = TILE_Y # threads per column (rows per tile) CTA_THREADS_X = TILE_X // MXFP8_BLOCK_SCALING_SIZE # threads per row (chunks per row) @@ -534,18 +539,125 @@ def dbias_reduction_colwise( return dbias_partial +@cute.jit +def quantize_mxfp8_tile( + sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA + sActInput_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) act-input smem tile (dact only) + sO_row_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) fp8 smem view (rowwise FP8 output) + mS_row_stage: Optional[cute.Tensor], # rowwise scale tensor (1D swizzled, or 2D linear) + sO_col_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) fp8 smem view (colwise FP8 output) + mS_col_stage: Optional[cute.Tensor], # colwise scale tensor (1D swizzled, or 2D linear) + rowwise_dbias_acc: Optional[cute.Tensor], # rmem Float32[32] (DBIAS_REDUCTION_IN_ROWWISE only) + tile_row_start: Int32, # global row index of this tile's row 0 + tile_col_start: Int32, # global col index of this tile's col 0 + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + WITH_AMAX: cutlass.Constexpr[bool], + WITH_ACT: cutlass.Constexpr[bool], + WITH_DACT: cutlass.Constexpr[bool], + ACTIVATION: cutlass.Constexpr[str | None], + 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_INPUT_MASKING: cutlass.Constexpr[bool], + SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool], + CACHE_ACTIVATION: cutlass.Constexpr[bool], + DBIAS_REDUCTION_IN_COLWISE: cutlass.Constexpr[bool], + COLWISE_DBIAS_REDUCTION_ONLY: cutlass.Constexpr[bool], + DBIAS_REDUCTION_IN_ROWWISE: cutlass.Constexpr[bool], +): + """Quantize one SMEM tile in the directions enabled by ROWWISE / COLWISE; returns (amax, dbias). + + amax is this thread's amax over the tile (None unless WITH_AMAX). dbias is this thread's + column partial sum (None unless DBIAS_REDUCTION_IN_COLWISE or COLWISE_DBIAS_REDUCTION_ONLY); + with DBIAS_REDUCTION_IN_ROWWISE the partial sums are accumulated into rowwise_dbias_acc instead. + """ + amax = None + dbias = None + if cutlass.const_expr(COLWISE): + amax_c, dbias_c = quantize_colwise_mxfp8( + sX_tile, + sActInput_tile, + sO_col_tile, + mS_col_stage, + MAX_NORM_RCP, + tile_row_start, + tile_col_start, + M, + N, + ACTIVATION=ACTIVATION, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + SWIZZLE=WITH_GEMM_SWIZZLED_SCALES, + TILE_X=TILE_X, + TILE_Y=TILE_Y, + SKIP_INPUT_MASKING=SKIP_INPUT_MASKING, + SKIP_SCALE_BOUNDS=SKIP_SCALE_BOUNDS, + WITH_ACT=WITH_ACT, + WITH_DACT=WITH_DACT, + WITH_DBIAS=DBIAS_REDUCTION_IN_COLWISE, + CACHE_ACTIVATION=CACHE_ACTIVATION, + ) + # Either direction's amax already spans every element of the tile, so favor colwise + # and skip the redundant comparison in the rowwise pass. + if cutlass.const_expr(WITH_AMAX): + amax = amax_c + if cutlass.const_expr(DBIAS_REDUCTION_IN_COLWISE): + dbias = dbias_c + # If we don't quantize columnwise but we fuse dbias, do a dbias reduction only pass in the columnwise direction without quantization + if cutlass.const_expr(COLWISE_DBIAS_REDUCTION_ONLY): + dbias = dbias_reduction_colwise( + sX_tile, + sActInput_tile, + ACTIVATION=ACTIVATION, + DTYPE=DTYPE, + TILE_X=TILE_X, + WITH_DACT=WITH_DACT, + CACHE_ACTIVATION=CACHE_ACTIVATION, + ) + # If we cache the activation in shared memory, we need to ensure that all threads have finished writing to the shared memory + # from the columnwise pass before any thread reads from it in the rowwise pass. + if cutlass.const_expr(CACHE_ACTIVATION): + cute.arch.sync_threads() + if cutlass.const_expr(ROWWISE): + amax_r = quantize_rowwise_mxfp8( + sX_tile, + None if CACHE_ACTIVATION else sActInput_tile, + sO_row_tile, + mS_row_stage, + MAX_NORM_RCP, + tile_row_start, + tile_col_start, + M, + N, + ACTIVATION=None if CACHE_ACTIVATION else ACTIVATION, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + TILE_X=TILE_X, + TILE_Y=TILE_Y, + SKIP_INPUT_MASKING=SKIP_INPUT_MASKING, + SKIP_SCALE_BOUNDS=SKIP_SCALE_BOUNDS, + WITH_ACT=WITH_ACT and not CACHE_ACTIVATION, + WITH_DACT=WITH_DACT and not CACHE_ACTIVATION, + WITH_DBIAS=DBIAS_REDUCTION_IN_ROWWISE, + dbias_acc=rowwise_dbias_acc, + ) + if cutlass.const_expr(WITH_AMAX and not COLWISE): + amax = amax_r + return amax, dbias + + class MXFP8QuantizeKernel(MXFP8QuantizeKernelBase): """The MXFP8 quantization kernel that mirrors the standard (non-specialized) MXFP8 CUDA C++ quantization kernel with multiple fusions (activation, dbias, etc.). `__call__` method is the entrypoint which is AOT compiled. `self` will be captured so it's fixed per compiled kernel """ - # Vectorised access constants for bank-conflict avoidance (rowwise pass) - _PACK_SIZE = 4 # Elements per vector load - # Each thread reads 8 waves with each wave reads 4 packed bf16, so it reads a whole MXFP8 block in total - _WAVES = MXFP8_BLOCK_SCALING_SIZE // _PACK_SIZE - _TOTAL_BANKS_WIDTH = (32 * 4) // 1 # 32 banks × 4 bytes, in bytes - _THREADS_PER_BANK = _TOTAL_BANKS_WIDTH // MXFP8_BLOCK_SCALING_SIZE # 4 threads per bank _NUM_STAGES = 2 # The pipeline depth is always 2 def __init__( @@ -602,13 +714,6 @@ def __init__( and (cfg.COLWISE or self.COLWISE_DBIAS_REDUCTION_ONLY) and cfg.ACTIVATION != "relu" ) - # The global tensor amax (mAmax) is the max over ALL elements. Each direction's - # per-block amaxes already span every element, so when both passes run we only - # fold the global amax from one of them — favor colwise (matches the flags - # above). The per-block *scale* amax is still computed in each pass for its own - # scale; this only skips the redundant global comparison in the other pass. - self.AMAX_FROM_COLWISE = cfg.WITH_AMAX and cfg.COLWISE - self.AMAX_FROM_ROWWISE = cfg.WITH_AMAX and not cfg.COLWISE @cute.jit def __call__( @@ -1074,34 +1179,8 @@ class DactStorage: # and tile_idx_y is the tile that this stage handles out of the `NUM_TILES` tiles tile_idx_x = bidx tile_idx_y = bidy * self._NUM_TILES + tile_idx - # Process rowwise and colwise quantization separately - if cutlass.const_expr(cfg.COLWISE): - # The first row that belongs to this CTA. Each CTA handles NUM_TILES of (TILE_Y, TILE_X) tiles stacked vertically, - # and each stage handles one of them. - sO_col_tile = sO_col[(None, stage_idx)] - mS_col_stage = cute.flatten(mS_col[(None, (tile_idx_y, tile_idx_x))]) - - amax_c, dbias_c = self._process_colwise( - sX_tile, - sO_col_tile, - mS_col_stage, - tile_idx_y * self._TILE_ROWS, - bidx * self._TILE_COLS, - M, - N, - sActInput_tile, - ) - if cutlass.const_expr(self.AMAX_FROM_COLWISE): - per_thread_amax = cute.arch.fmax(per_thread_amax, amax_c) - if cutlass.const_expr(self.DBIAS_REDUCTION_IN_COLWISE): - block_dbias += dbias_c - # If we don't quantize columnwise but we fuse dbias, do a dbias reduction only pass in the columnwise direction without quantization - if cutlass.const_expr(self.COLWISE_DBIAS_REDUCTION_ONLY): - block_dbias += self._dbias_only_colwise(sX_tile, sActInput_tile) - # If we cache the activation in shared memory, we need to ensure that all threads have finished writing to the shared memory - # from the columnwise pass before any thread reads from it in the rowwise pass. - if cutlass.const_expr(self.CACHE_ACTIVATION): - cute.arch.sync_threads() + sO_row_tile = None + mS_row_stage = None if cutlass.const_expr(cfg.ROWWISE): sO_row_tile = sO_row[(None, stage_idx)] # mS_row is ((SCALE_TILE), (GRID)) where SCALE_TILE = (32, 2). @@ -1111,20 +1190,49 @@ class DactStorage: # (i_hi, tile_Y) tile-major order for swizzled, and the plain # row-tile order for compact. Same source, both layouts correct. mS_row_stage = cute.flatten(mS_row[(None, (tile_idx_y, tile_idx_x))]) - amax_r = self._process_rowwise( - sX_tile, - sO_row_tile, - mS_row_stage, - tile_idx_y * self._TILE_ROWS, - bidx * self._TILE_COLS, - M, - N, - sActInput_tile, - rowwise_dbias_acc, - ) + sO_col_tile = None + mS_col_stage = None + if cutlass.const_expr(cfg.COLWISE): + sO_col_tile = sO_col[(None, stage_idx)] + mS_col_stage = cute.flatten(mS_col[(None, (tile_idx_y, tile_idx_x))]) - if cutlass.const_expr(self.AMAX_FROM_ROWWISE): - per_thread_amax = cute.arch.fmax(per_thread_amax, amax_r) + amax, dbias = quantize_mxfp8_tile( + sX_tile, + sActInput_tile, + sO_row_tile, + mS_row_stage, + sO_col_tile, + mS_col_stage, + rowwise_dbias_acc, + tile_idx_y * self._TILE_ROWS, + tile_idx_x * self._TILE_COLS, + M, + N, + ROWWISE=self.cfg.ROWWISE, + COLWISE=self.cfg.COLWISE, + WITH_AMAX=self.cfg.WITH_AMAX, + WITH_ACT=self.cfg.WITH_ACT, + WITH_DACT=self.cfg.WITH_DACT, + ACTIVATION=self.cfg.ACTIVATION, + MAX_NORM_RCP=self.cfg.MAX_NORM_RCP, + DTYPE=self.cfg.DTYPE, + FP8_DTYPE=self.cfg.FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=self.cfg.WITH_GEMM_SWIZZLED_SCALES, + TILE_X=self._TILE_COLS, + TILE_Y=self._TILE_ROWS, + SKIP_INPUT_MASKING=self.SKIP_INPUT_MASKING, + SKIP_SCALE_BOUNDS=self.SKIP_SCALE_BOUNDS, + CACHE_ACTIVATION=self.CACHE_ACTIVATION, + 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, + ) + if cutlass.const_expr(cfg.WITH_AMAX): + per_thread_amax = cute.arch.fmax(per_thread_amax, amax) + if cutlass.const_expr( + self.DBIAS_REDUCTION_IN_COLWISE or self.COLWISE_DBIAS_REDUCTION_ONLY + ): + block_dbias += dbias # Make the shared-memory writes visible to the TMA's async proxy before the TMA reads them. cute.arch.fence_proxy( @@ -1199,24 +1307,6 @@ class DactStorage: # before the kernel returns. cute.arch.cp_async_bulk_wait_group(0, read=False) - @cute.jit - def _dbias_only_colwise( - self, - sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA - sActInput_tile: Optional[ - cute.Tensor - ] = None, # (TILE_Y, TILE_X) act_input tile (dact only), - ): - return dbias_reduction_colwise( - sX_tile, - sActInput_tile, - ACTIVATION=self.cfg.ACTIVATION, - DTYPE=self.cfg.DTYPE, - TILE_X=self._TILE_COLS, - WITH_DACT=self.cfg.WITH_DACT, - CACHE_ACTIVATION=self.CACHE_ACTIVATION, - ) - @cute.jit def _dbias_reduction_rowwise_epilouge( self, smem: cutlass.memory.SmemAllocator, tidx: Int32, rowwise_dbias_acc: cute.Tensor @@ -1246,18 +1336,22 @@ class DbiasStorage: ), ) sDbias_write = cute.composition(sDbias, tv_layout_dbias_write) + # rowwise_dbias_acc is in quantize_rowwise_mxfp8's rotated wave order; these must match its constants. + PACK_SIZE = 4 + WAVES = MXFP8_BLOCK_SCALING_SIZE // PACK_SIZE + THREADS_PER_BANK = 32 * 4 // MXFP8_BLOCK_SCALING_SIZE # 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) // self._THREADS_PER_BANK + 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 * self._PACK_SIZE + offset = bank_group * PACK_SIZE for w in cutlass.range_constexpr( - self._WAVES + WAVES ): # Each thread starts from this offset when writing into SMEM to avoid bank conflict - start = (w * self._PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE - for i in cutlass.range_constexpr(self._PACK_SIZE): + start = (w * PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE + for i in cutlass.range_constexpr(PACK_SIZE): # All threads write their per-thread partial sum results to the shared buffer. - sDbias_write[(tidx, start + i)] = rowwise_dbias_acc[w * self._PACK_SIZE + i] + sDbias_write[(tidx, start + i)] = rowwise_dbias_acc[w * PACK_SIZE + i] cute.arch.sync_threads() # All threads reduce the cross-thread partial sums to the per-block partial sum. _, tv_layout_dbias_reduce = cute.make_layout_tv( @@ -1302,82 +1396,3 @@ def _amax_epilogue( amax_i32.iterator, cta_amax.bitcast(Int32), ) - - @cute.jit - def _process_rowwise( - self, - sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 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) - tile_row_start: Int32, # global row of this stage's row 0 - tile_col_start: Int32, # global col of this CTA's col 0 - M: Int32, - N: Int32, # full input extents, for OOB masking - sActInput_tile: Optional[cute.Tensor] = None, # (TILE_Y, TILE_X) act_input tile (dact only) - dbias_acc: Optional[ - cute.Tensor - ] = None, # rmem Float32[32] dbias accumulator (rowwise-only dbias) - ): - cfg = self.cfg - return quantize_rowwise_mxfp8( - sX_tile, - None if self.CACHE_ACTIVATION else sActInput_tile, - sO_row_tile, - mS_row_stage, - self.cfg.MAX_NORM_RCP, - tile_row_start, - tile_col_start, - M, - N, - ACTIVATION=None if self.CACHE_ACTIVATION else cfg.ACTIVATION, - DTYPE=cfg.DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, - TILE_X=self._TILE_COLS, - TILE_Y=self._TILE_ROWS, - WAVES=self._WAVES, - THREADS_PER_BANK=self._THREADS_PER_BANK, - PACK_SIZE=self._PACK_SIZE, - SKIP_INPUT_MASKING=self.SKIP_INPUT_MASKING, - SKIP_SCALE_BOUNDS=self.SKIP_SCALE_BOUNDS, - WITH_ACT=cfg.WITH_ACT and not self.CACHE_ACTIVATION, - WITH_DACT=cfg.WITH_DACT and not self.CACHE_ACTIVATION, - WITH_DBIAS=self.DBIAS_REDUCTION_IN_ROWWISE, - dbias_acc=dbias_acc, - ) - - @cute.jit - def _process_colwise( - self, - sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 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) - tile_row_start: Int32, # global row of this stage's row 0 - tile_col_start: Int32, # global col of this CTA's col 0 - M: Int32, - N: Int32, # full input extents, for OOB masking - sActInput_tile: Optional[cute.Tensor] = None, # (TILE_Y, TILE_X) act_input tile (dact only) - ): - cfg = self.cfg - return quantize_colwise_mxfp8( - sX_tile, - sActInput_tile, - sO_col_tile, - mS_col_stage, - self.cfg.MAX_NORM_RCP, - tile_row_start, - tile_col_start, - M, - N, - ACTIVATION=cfg.ACTIVATION, - DTYPE=cfg.DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, - SWIZZLE=cfg.WITH_GEMM_SWIZZLED_SCALES, - TILE_X=self._TILE_COLS, - TILE_Y=self._TILE_ROWS, - SKIP_INPUT_MASKING=self.SKIP_INPUT_MASKING, - SKIP_SCALE_BOUNDS=self.SKIP_SCALE_BOUNDS, - WITH_ACT=cfg.WITH_ACT, - WITH_DACT=cfg.WITH_DACT, - WITH_DBIAS=self.DBIAS_REDUCTION_IN_COLWISE, - CACHE_ACTIVATION=self.CACHE_ACTIVATION, - ) From fddd076e98610520213d46b2399d783a7d6db53a Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 14 Oct 2026 00:09:31 +0000 Subject: [PATCH 21/29] refactor Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 26 +++++++++++-------- .../mxfp8/quantize_mxfp8_g2s_bidimensional.py | 7 ++--- 2 files changed, 19 insertions(+), 14 deletions(-) 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 68c69040cfc..9db77f4d36e 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 @@ -4,7 +4,7 @@ """Two-lane GMEM-to-RMEM rowwise MXFP8 cast, ported from cast_rowwise.cuh.""" -from typing import Optional +from typing import Optional, Type import cutlass from cutlass import cute @@ -34,8 +34,10 @@ def quantize_rowwise_mxfp8_g2r( gX_tiles: cute.Tensor, gO_tiles: cute.Tensor, mS_tiles: cute.Tensor, - cfg: cutlass.Constexpr[MXFP8QuantizeConfig], N: Int32, + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + MAX_NORM_RCP: cutlass.Constexpr[float], CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], ELEMENTS_PER_THREAD: cutlass.Constexpr[int], @@ -69,21 +71,21 @@ def quantize_rowwise_mxfp8_g2r( ) # We load fragments from all GMEM tiles at once - rX = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, NUM_TILES), cfg.DTYPE) + rX = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, NUM_TILES), DTYPE) # We store a fragment of one tile at a time - rO = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), cfg.FP8_DTYPE) + rO = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), FP8_DTYPE) # RMEM flags to indicate whether each tile is out of bounds or not # With CHECK_BOUNDS=False DCE will remove it so we don't pay extra cost when we don't use it tile_valid_flags = cute.make_rmem_tensor((NUM_TILES,), Boolean) # 16 * 2 bytes = 32 bytes, one 256-bit load where the toolchain supports it LOAD_WIDTH = min( - ELEMENTS_PER_THREAD * cfg.DTYPE.width, + ELEMENTS_PER_THREAD * DTYPE.width, max_num_bits_per_copy_ld_global(), ) load_atom = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), - cfg.DTYPE, + DTYPE, num_bits_per_copy=LOAD_WIDTH, # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using # invariant=True here will drop the cache policy @@ -91,11 +93,11 @@ def quantize_rowwise_mxfp8_g2r( ) store_atom = cute.make_copy_atom( cute.nvgpu.CopyR2GOp(), - cfg.FP8_DTYPE, - num_bits_per_copy=ELEMENTS_PER_THREAD * cfg.FP8_DTYPE.width, + FP8_DTYPE, + num_bits_per_copy=ELEMENTS_PER_THREAD * FP8_DTYPE.width, ) - scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(cfg.FP8_DTYPE) + scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(FP8_DTYPE) # Copy all tiles from GMEM to RMEM and potentially check their bounds for tile_idx in cutlass.range_constexpr(NUM_TILES): @@ -144,7 +146,7 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Float32: rX_tile_i32 = cute.recast_tensor(rX[None, tile_idx], Int32) rO_u32 = cute.recast_tensor(rO, Uint32) amax = extract_mx_block_amax(rX_tile_i32) - exponent = cvt_f32_to_fp8e8m0fnu(amax * cfg.MAX_NORM_RCP) + exponent = cvt_f32_to_fp8e8m0fnu(amax * MAX_NORM_RCP) # The even lane of each pair stores its block's scale byte, as cast_rowwise.cuh does. # A warp's stores coalesce in either layout: one 16-byte run when packed, four @@ -479,8 +481,10 @@ def _kernel_main( gX_tiles, gO_tiles, mS_tiles, - self.cfg, N, + DTYPE=self.cfg.DTYPE, + FP8_DTYPE=self.cfg.FP8_DTYPE, + MAX_NORM_RCP=self.cfg.MAX_NORM_RCP, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py index f6c6d391d99..004d2a64a99 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py @@ -43,8 +43,9 @@ @cute.jit -def quantize_bidimensional_mxfp8_swizzled( - sX_tile: cute.Tensor, # 32x64 input tile (already sliced to this stage), Sw<3,4,3> swizzled +def quantize_bidimensional_mxfp8( + sX_tile: cute.Tensor, # 32x64 input tile (already sliced to this stage); any SMEM layout, + # Sw<3,4,3> avoids bank conflicts sO_row_tile: cute.Tensor, # 32x64 rowwise-output tile sO_col_tile: cute.Tensor, # 32x64 colwise-output tile sS_row_tile: cute.Tensor, # (32, 2) smem rowwise-scale staging tile (flushed in the epilogue) @@ -611,7 +612,7 @@ class SharedStorage: tile_x = tile_idx % num_tiles_x # Process the 32x64 SMEM tile for this stage - quantize_bidimensional_mxfp8_swizzled( + quantize_bidimensional_mxfp8( sX[None, stage_idx], sO_row[None, stage_idx], sO_col[None, stage_idx], From 3a8a8825f9154df770e2bd3d4f218bf961f9b801 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 14 Oct 2026 20:03:01 +0000 Subject: [PATCH 22/29] more refactor Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_common.py | 40 +++++++------------ ...nsional.py => quantize_mxfp8_g2r_bidim.py} | 8 ++-- .../mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py | 10 ++--- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 8 ++-- ...nsional.py => quantize_mxfp8_g2s_bidim.py} | 12 +++--- ...neral.py => quantize_mxfp8_g2s_general.py} | 26 ++++++------ transformer_engine/common/CuTeDSL/utils.py | 11 +++++ 7 files changed, 59 insertions(+), 56 deletions(-) rename transformer_engine/common/CuTeDSL/cast/mxfp8/{quantize_mxfp8_g2r_bidimensional.py => quantize_mxfp8_g2r_bidim.py} (99%) rename transformer_engine/common/CuTeDSL/cast/mxfp8/{quantize_mxfp8_g2s_bidimensional.py => quantize_mxfp8_g2s_bidim.py} (98%) rename transformer_engine/common/CuTeDSL/cast/mxfp8/{quantize_mxfp8_general.py => quantize_mxfp8_g2s_general.py} (98%) 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 748f19756f4..133e7232f48 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -11,13 +11,12 @@ import cutlass from cutlass import cute -from cutlass import Boolean, Float32, Int32, Int64, Uint32 +from cutlass import Boolean, Float32, Int32, Int64 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module import tvm_ffi from transformer_engine.common.CuTeDSL.utils import ( str_to_cutlass_dtype, - abs_max_x2_bf16, device_compute_capability, ) from transformer_engine.common.CuTeDSL.activations import ( @@ -256,17 +255,6 @@ def __call__( """ -@cute.jit -def bf16_pair_magnitude(pair: Int32) -> Uint32: - """Fold both BF16 halves and clear the sign set by max.xorsign.abs.""" - swapped = cute.arch.inline_ptx( - "prmt.b32 {$w0}, {$r0}, {$r0}, 0x1032;", - write_only_types=[Int32], - read_only_args=[pair], - ) - return Uint32(abs_max_x2_bf16(pair, swapped)) & Uint32(0x7FFF) - - class MXFP8QuantizeEntry(MXFP8QuantizeKernelBase): """Select the appropriate MXFP8 quantization kernel based on the configuration and runtime shapes.""" @@ -274,11 +262,11 @@ def __init__(self, cfg: MXFP8QuantizeConfig): # Kernels import shared definitions from this module. Resolve their # implementations only when constructing the compile-time entry object. # pylint: disable=import-outside-toplevel - from .quantize_mxfp8_general import MXFP8QuantizeKernel - from .quantize_mxfp8_g2r_rowwise_1lane import MXFP8QuantizeG2RRowwise1LaneKernel - from .quantize_mxfp8_g2s_bidimensional import MXFP8QuantizeSpecializedBidimensionalKernel - from .quantize_mxfp8_g2r_rowwise_2lane import MXFP8QuantizeG2RRowwise2LaneKernel - from .quantize_mxfp8_g2r_bidimensional import MXFP8QuantizeRegisterBidimensionalKernel + from .quantize_mxfp8_g2s_general import MXFP8QuantizeGeneralG2SKernel + from .quantize_mxfp8_g2r_rowwise_1lane import MXFP8QuantizeRowwise1LaneG2RKernel + from .quantize_mxfp8_g2s_bidim import MXFP8QuantizeBidimensionalG2SKernel + from .quantize_mxfp8_g2r_rowwise_2lane import MXFP8QuantizeRowwise2LaneG2RKernel + from .quantize_mxfp8_g2r_bidim import MXFP8QuantizeBidimensionalG2RKernel # pylint: enable=import-outside-toplevel @@ -296,21 +284,23 @@ def __init__(self, cfg: MXFP8QuantizeConfig): ) # 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 = MXFP8QuantizeKernel(cfg, SKIP_INPUT_MASKING=True) - self.general_skip_input_masking_and_scale_bounds_kernel = MXFP8QuantizeKernel( + 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 ) # We only need to mask when WITH_ACT is enabled because with activations applied zeros filled by TMA affect block statistics - self.general_kernel = MXFP8QuantizeKernel(cfg) if self.ACT_NEED_MASKING else None + self.general_kernel = MXFP8QuantizeGeneralG2SKernel(cfg) if self.ACT_NEED_MASKING else None # Specialized cast-only kernels in a specific dispatch priority order so we prefer to dispatch to kernels that come first in the list. self.specialized_kernels = [ kernel_cls(cfg) for kernel_cls in ( - MXFP8QuantizeG2RRowwise2LaneKernel, - MXFP8QuantizeRegisterBidimensionalKernel, - MXFP8QuantizeG2RRowwise1LaneKernel, - MXFP8QuantizeSpecializedBidimensionalKernel, + MXFP8QuantizeRowwise2LaneG2RKernel, + MXFP8QuantizeBidimensionalG2RKernel, + MXFP8QuantizeRowwise1LaneG2RKernel, + MXFP8QuantizeBidimensionalG2SKernel, ) if kernel_cls.supports(cfg) ] diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py similarity index 99% rename from transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py rename to transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 75ad5799cbf..e1d828147c9 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -13,6 +13,7 @@ from transformer_engine.common.CuTeDSL.utils import ( abs_max_x2_bf16, + bf16_pair_magnitude, bf16_mx_scale_reciprocal, bf16_mx_scale_reciprocal_pair, bf16_mx_scale_bytes_pair, @@ -24,7 +25,6 @@ MXFP8_BLOCK_SCALING_SIZE, MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, - bf16_pair_magnitude, derive_swizzled_scale_layout, ) @@ -388,7 +388,7 @@ def _quantize_colwise( @cute.jit -def quantize_bidimensional_mxfp8_g2r( +def quantize_mxfp8_gmem_tile_bidimensional( gX_tile: cute.Tensor, gO_row_tile: Optional[cute.Tensor], gS_row_tile: Optional[cute.Tensor], @@ -544,7 +544,7 @@ def quantize_bidimensional_mxfp8_g2r( ) -class MXFP8QuantizeRegisterBidimensionalKernel(MXFP8QuantizeKernelBase): +class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): """Load one CTA tile into registers and quantize it along both axes. TILE_ROWS = _CTA_THREADS_Y * _ELEMENTS_Y_PER_THREAD = 32. @@ -789,7 +789,7 @@ def _kernel_main( cache_policy = create_l2_policy(True) - quantize_bidimensional_mxfp8_g2r( + quantize_mxfp8_gmem_tile_bidimensional( gX_tile, gO_row_tile, gS_row_tile, 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 1c005c22fdb..97ea8c0f5f4 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 @@ -74,7 +74,7 @@ def _flush_scales_to_gmem( @cute.jit -def quantize_rowwise_mxfp8_g2r_1lane( +def quantize_mxfp8_gmem_tile_rowwise_1lane( gX_tile: cute.Tensor, gO_tile: cute.Tensor, gS_tile: cute.Tensor, @@ -205,7 +205,7 @@ class SharedStorage: ) -class MXFP8QuantizeG2RRowwise1LaneKernel(MXFP8QuantizeKernelBase): +class MXFP8QuantizeRowwise1LaneG2RKernel(MXFP8QuantizeKernelBase): """Specialized cast-only ROWWISE-only MXFP8 kernel. Requires N % 128 == 0 (full vectorizable column chunks). @@ -236,7 +236,7 @@ def eligible_to_dispatch( # N divisible by 128 for vectorized stores, and grid.y within limits. return ( N % 128 == 0 - and cute.ceil_div(M, MXFP8QuantizeG2RRowwise1LaneKernel._TILE_ROWS) <= 65535 + and cute.ceil_div(M, MXFP8QuantizeRowwise1LaneG2RKernel._TILE_ROWS) <= 65535 and mNoop.toint() == 0 ) @@ -262,7 +262,7 @@ def __call__( ): if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): cute.printf( - f"[CuTeDSL] MXFP8QuantizeG2RRowwise1LaneKernel.__call__() with config: {self.cfg}\n" + f"[CuTeDSL] MXFP8QuantizeRowwise1LaneG2RKernel.__call__() with config: {self.cfg}\n" ) M = mX.shape[0] @@ -312,7 +312,7 @@ def _kernel_main( gO_tile = cute.local_tile(mO_row, TILER, (bidy, bidx)) gS_tile = cute.local_tile(mS_row, SCALE_TILER, (bidy, bidx)) - quantize_rowwise_mxfp8_g2r_1lane( + quantize_mxfp8_gmem_tile_rowwise_1lane( gX_tile, gO_tile, gS_tile, 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 9db77f4d36e..d1eb2cd0753 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 @@ -13,6 +13,7 @@ from transformer_engine.common.CuTeDSL.utils import ( abs_max_x2_bf16, + bf16_pair_magnitude, create_l2_policy, exp2_bf16x2_rcp, max_num_bits_per_copy_ld_global, @@ -24,13 +25,12 @@ from .quantize_mxfp8_common import ( MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, - bf16_pair_magnitude, derive_swizzled_scale_layout, ) @cute.jit -def quantize_rowwise_mxfp8_g2r( +def quantize_mxfp8_gmem_tile_rowwise_2lane( gX_tiles: cute.Tensor, gO_tiles: cute.Tensor, mS_tiles: cute.Tensor, @@ -167,7 +167,7 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Float32: cute.copy(store_atom, rO, gO_thread, cache_policy=output_policy) -class MXFP8QuantizeG2RRowwise2LaneKernel(MXFP8QuantizeKernelBase): +class MXFP8QuantizeRowwise2LaneG2RKernel(MXFP8QuantizeKernelBase): """Two lanes per MXFP8 block, with a CTA sweeping NUM_TILES horizontally. One CTA tile has TILE_ROWS = CTA_THREADS_Y and @@ -477,7 +477,7 @@ def _kernel_main( # The last grid.z slice can overshoot the row tiles; those CTAs have no work. if cta_row * CTA_ROWS < M: - quantize_rowwise_mxfp8_g2r( + quantize_mxfp8_gmem_tile_rowwise_2lane( gX_tiles, gO_tiles, mS_tiles, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py similarity index 98% rename from transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py rename to transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py index 004d2a64a99..65bf548022a 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidimensional.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py @@ -43,7 +43,7 @@ @cute.jit -def quantize_bidimensional_mxfp8( +def quantize_mxfp8_smem_tile_bidimensional( sX_tile: cute.Tensor, # 32x64 input tile (already sliced to this stage); any SMEM layout, # Sw<3,4,3> avoids bank conflicts sO_row_tile: cute.Tensor, # 32x64 rowwise-output tile @@ -229,7 +229,7 @@ def quantize_bidimensional_mxfp8( cute.autovec_copy(rO_col, tXsO_col) -class MXFP8QuantizeSpecializedBidimensionalKernel(MXFP8QuantizeKernelBase): +class MXFP8QuantizeBidimensionalG2SKernel(MXFP8QuantizeKernelBase): """Specialized cast-only rowwise+colwise MXFP8 kernel (swizzled TMA, one 32x32 tile per warp).""" _WARPS_PER_CTA = 2 @@ -271,7 +271,7 @@ def eligible_to_dispatch( mS_col: Optional[cute.Tensor], mNoop: cute.Pointer, ) -> Boolean: - cls = MXFP8QuantizeSpecializedBidimensionalKernel + cls = MXFP8QuantizeBidimensionalG2SKernel rows_per_cta = cls._TILE_ROWS * cls._NUM_TILES_Y # packed scales: one tile row return cute.ceil_div(mX.shape[0], rows_per_cta) <= 65535 and mNoop.toint() == 0 @@ -298,7 +298,7 @@ def __call__( ): if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): cute.printf( - "[CuTeDSL] MXFP8QuantizeSpecializedBidimensionalKernel.__call__() with config:" + "[CuTeDSL] MXFP8QuantizeBidimensionalG2SKernel.__call__() with config:" f" {self.cfg}\n" ) M = mX.shape[0] @@ -611,8 +611,8 @@ class SharedStorage: tile_y = tile_idx // num_tiles_x tile_x = tile_idx % num_tiles_x - # Process the 32x64 SMEM tile for this stage - quantize_bidimensional_mxfp8( + # Process the 32x64 SMEM tile for thiHs stage + quantize_mxfp8_smem_tile_bidimensional( sX[None, stage_idx], sO_row[None, stage_idx], sO_col[None, stage_idx], diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py similarity index 98% rename from transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py rename to transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py index d74d4ca824b..9801d82a259 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py @@ -76,7 +76,7 @@ @cute.jit -def quantize_rowwise_mxfp8( +def _quantize_rowwise_mxfp8( sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA sActInput_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) act-input smem tile (dact only) sO_row_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (rowwise FP8 output) @@ -329,7 +329,7 @@ def quantize_rowwise_mxfp8( @cute.jit -def quantize_colwise_mxfp8( +def _quantize_colwise_mxfp8( sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA sActInput_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) act-input smem tile (dact only) sO_col_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (colwise FP8 output) @@ -505,7 +505,7 @@ def quantize_colwise_mxfp8( @cute.jit -def dbias_reduction_colwise( +def _dbias_reduction_colwise( sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16/fp32 smem view, post-TMA sA_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) activation-input smem tile (dact only) ACTIVATION: cutlass.Constexpr[str | None], @@ -517,7 +517,7 @@ def dbias_reduction_colwise( """Reduce one SMEM tile along rows (the dbias direction) and possibly cache dact values, with no quantization.""" tidx, _, _ = cute.arch.thread_idx() - # Same TV layout as quantize_colwise_mxfp8: thread tidx owns column tidx, all TILE_Y rows. + # Same TV layout as _quantize_colwise_mxfp8: thread tidx owns column tidx, all TILE_Y rows. _, tv_layout = cute.make_layout_tv( thr_layout=cute.make_layout((1, TILE_X), stride=(TILE_X, 1)), val_layout=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE, 1), stride=(1, 1)), @@ -540,7 +540,7 @@ def dbias_reduction_colwise( @cute.jit -def quantize_mxfp8_tile( +def quantize_mxfp8_smem_tile_general( sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA sActInput_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) act-input smem tile (dact only) sO_row_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) fp8 smem view (rowwise FP8 output) @@ -580,7 +580,7 @@ def quantize_mxfp8_tile( amax = None dbias = None if cutlass.const_expr(COLWISE): - amax_c, dbias_c = quantize_colwise_mxfp8( + amax_c, dbias_c = _quantize_colwise_mxfp8( sX_tile, sActInput_tile, sO_col_tile, @@ -611,7 +611,7 @@ def quantize_mxfp8_tile( dbias = dbias_c # If we don't quantize columnwise but we fuse dbias, do a dbias reduction only pass in the columnwise direction without quantization if cutlass.const_expr(COLWISE_DBIAS_REDUCTION_ONLY): - dbias = dbias_reduction_colwise( + dbias = _dbias_reduction_colwise( sX_tile, sActInput_tile, ACTIVATION=ACTIVATION, @@ -625,7 +625,7 @@ def quantize_mxfp8_tile( if cutlass.const_expr(CACHE_ACTIVATION): cute.arch.sync_threads() if cutlass.const_expr(ROWWISE): - amax_r = quantize_rowwise_mxfp8( + amax_r = _quantize_rowwise_mxfp8( sX_tile, None if CACHE_ACTIVATION else sActInput_tile, sO_row_tile, @@ -652,7 +652,7 @@ def quantize_mxfp8_tile( return amax, dbias -class MXFP8QuantizeKernel(MXFP8QuantizeKernelBase): +class MXFP8QuantizeGeneralG2SKernel(MXFP8QuantizeKernelBase): """The MXFP8 quantization kernel that mirrors the standard (non-specialized) MXFP8 CUDA C++ quantization kernel with multiple fusions (activation, dbias, etc.). `__call__` method is the entrypoint which is AOT compiled. `self` will be captured so it's fixed per compiled kernel @@ -734,7 +734,9 @@ def __call__( stream: CUstream, ): if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): - cute.printf(f"[CuTeDSL] MXFP8QuantizeKernel.__call__() with config: {self.cfg}\n") + cute.printf( + f"[CuTeDSL] MXFP8QuantizeGeneralG2SKernel.__call__() with config: {self.cfg}\n" + ) M = mX.shape[0] N = mX.shape[1] @@ -1196,7 +1198,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_tile( + amax, dbias = quantize_mxfp8_smem_tile_general( sX_tile, sActInput_tile, sO_row_tile, @@ -1336,7 +1338,7 @@ class DbiasStorage: ), ) sDbias_write = cute.composition(sDbias, tv_layout_dbias_write) - # rowwise_dbias_acc is in quantize_rowwise_mxfp8's rotated wave order; these must match its constants. + # rowwise_dbias_acc is in _quantize_rowwise_mxfp8's rotated wave order; these must match its constants. PACK_SIZE = 4 WAVES = MXFP8_BLOCK_SCALING_SIZE // PACK_SIZE THREADS_PER_BANK = 32 * 4 // MXFP8_BLOCK_SCALING_SIZE diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index d5ffd4abfe9..caaf8339524 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -102,6 +102,17 @@ def exp2_bf16x2_rcp(scale_e8m0) -> Uint32: return bf16_bits | (bf16_bits << Uint32(16)) +@cute.jit +def bf16_pair_magnitude(pair: Int32) -> Uint32: + """Fold both BF16 halves and clear the sign set by max.xorsign.abs.""" + swapped = cute.arch.inline_ptx( + "prmt.b32 {$w0}, {$r0}, {$r0}, 0x1032;", + write_only_types=[Int32], + read_only_args=[pair], + ) + return Uint32(abs_max_x2_bf16(pair, swapped)) & Uint32(0x7FFF) + + @cute.jit def bf16_mx_scale_reciprocal(magnitude: Uint32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: """Exact BF16 reciprocal scale bits, including E8M0 zero, Inf and NaN cases.""" From 30c92ab3bc274e01fecb36554cfea5c9a46705b4 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 14 Oct 2026 14:16:13 -0700 Subject: [PATCH 23/29] g2r bidim can also support colwise only Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 53 ++++++++++++++----- 1 file changed, 39 insertions(+), 14 deletions(-) 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 e1d828147c9..5e8d724504e 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 @@ -561,7 +561,7 @@ class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. - return cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 and cfg.ROWWISE and cfg.COLWISE + return cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 and cfg.COLWISE @staticmethod @cute.jit @@ -680,14 +680,24 @@ def _launch( if N == 2048: payload_layout = cute.make_layout((M, 2048), stride=(2048, 1)) gX = cute.make_tensor(mX.iterator, payload_layout) - gO_row = cute.make_tensor(mO_row.iterator, payload_layout) + gO_row = ( + cute.make_tensor(mO_row.iterator, payload_layout) + if cutlass.const_expr(self.cfg.ROWWISE) + else None + ) gO_col = cute.make_tensor(mO_col.iterator, payload_layout) if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - gS_row, gS_col = derive_swizzled_scale_layout(M, 2048, True, True, mS_row, mS_col) + gS_row, gS_col = derive_swizzled_scale_layout( + M, 2048, self.cfg.ROWWISE, True, mS_row, mS_col + ) else: - gS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout((cute.size(mS_row, mode=[0]), 64), stride=(64, 1)), + gS_row = ( + cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 64), stride=(64, 1)), + ) + if cutlass.const_expr(self.cfg.ROWWISE) + else None ) gS_col = cute.make_tensor( mS_col.iterator, @@ -704,14 +714,24 @@ def _launch( elif N == 4096: payload_layout = cute.make_layout((M, 4096), stride=(4096, 1)) gX = cute.make_tensor(mX.iterator, payload_layout) - gO_row = cute.make_tensor(mO_row.iterator, payload_layout) + gO_row = ( + cute.make_tensor(mO_row.iterator, payload_layout) + if cutlass.const_expr(self.cfg.ROWWISE) + else None + ) gO_col = cute.make_tensor(mO_col.iterator, payload_layout) if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - gS_row, gS_col = derive_swizzled_scale_layout(M, 4096, True, True, mS_row, mS_col) + gS_row, gS_col = derive_swizzled_scale_layout( + M, 4096, self.cfg.ROWWISE, True, mS_row, mS_col + ) else: - gS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout((cute.size(mS_row, mode=[0]), 128), stride=(128, 1)), + gS_row = ( + cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 128), stride=(128, 1)), + ) + if cutlass.const_expr(self.cfg.ROWWISE) + else None ) gS_col = cute.make_tensor( mS_col.iterator, @@ -727,7 +747,9 @@ def _launch( ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) else: if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - gS_row, gS_col = derive_swizzled_scale_layout(M, N, True, True, mS_row, mS_col) + gS_row, gS_col = derive_swizzled_scale_layout( + M, N, self.cfg.ROWWISE, True, mS_row, mS_col + ) else: gS_row, gS_col = mS_row, mS_col self.kernel( @@ -782,9 +804,12 @@ def _kernel_main( COL_SCALE_TILER = (TILE_ROWS // 32, TILE_COLS) gX_tile = cute.local_tile(mX, TILER, cta_coord) - gO_row_tile = cute.local_tile(mO_row, TILER, cta_coord) + gO_row_tile = None + gS_row_tile = None + if cutlass.const_expr(self.cfg.ROWWISE): + gO_row_tile = cute.local_tile(mO_row, TILER, cta_coord) + gS_row_tile = cute.local_tile(mS_row, ROW_SCALE_TILER, cta_coord) gO_col_tile = cute.local_tile(mO_col, TILER, cta_coord) - gS_row_tile = cute.local_tile(mS_row, ROW_SCALE_TILER, cta_coord) gS_col_tile = cute.local_tile(mS_col, COL_SCALE_TILER, cta_coord) cache_policy = create_l2_policy(True) From 5c3465a56103e1805e4e737ed0334b1139358f42 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 20:29:21 +0000 Subject: [PATCH 24/29] make g2r support activations Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_common.py | 7 +- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 93 +++++++++++++- .../mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py | 1 + .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 114 +++++++++++++++++- .../cast/mxfp8/quantize_mxfp8_g2s_bidim.py | 1 + transformer_engine/common/CuTeDSL/utils.py | 10 ++ 6 files changed, 219 insertions(+), 7 deletions(-) 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 133e7232f48..0df60e0714e 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -230,6 +230,7 @@ def eligible_to_dispatch( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], ) -> Boolean: """Runtime check: whether we can dispatch to the instantiated kernel with this call's shapes and noops.""" raise NotImplementedError @@ -322,7 +323,9 @@ def __call__( dispatched_to_specialized = False for kernel in self.specialized_kernels: if not dispatched_to_specialized: - if kernel.eligible_to_dispatch(self.cfg, mX, mO_row, mS_row, mO_col, mS_col, mNoop): + if kernel.eligible_to_dispatch( + self.cfg, mX, mO_row, mS_row, mO_col, mS_col, mNoop, mDActInput + ): dispatched_to_specialized = True kernel( mX, @@ -503,7 +506,7 @@ def compile_cutedsl_function_from_cfg(cfg): scale_col_fake, # mO_col, mS_col amax_fake, # mAmax noop_fake, # mNoop (pointer to the cast_noop flag, may be null at launch) - act_input_fake, # mDActInput (backward slot, unused) + act_input_fake, # mDActInput (activation input, dact only) workspace_fake, # mWorkspace(backward slot, unused) cute.runtime.make_fake_stream(), # stream (compiled as an explicit tvm-ffi # "handle" arg; C++ passes the CUDA stream 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 5e8d724504e..0e8efdb3bd3 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 @@ -8,7 +8,7 @@ import cutlass from cutlass import cute -from cutlass import Boolean, Int32, Int64, Uint32, Uint8 +from cutlass import Boolean, Float32, Int32, Int64, Uint32, Uint8 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( @@ -19,12 +19,17 @@ bf16_mx_scale_bytes_pair, create_l2_policy, max_num_bits_per_copy_ld_global, + x2_hi_to_f32_bf16, + x2_lo_to_f32_bf16, + cvt_f32x2_to_bf16x2, ) from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 from .quantize_mxfp8_common import ( MXFP8_BLOCK_SCALING_SIZE, MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, + SUPPORTED_ACTIVATIONS, + SUPPORTED_DACTIVATIONS, derive_swizzled_scale_layout, ) @@ -390,12 +395,16 @@ def _quantize_colwise( @cute.jit def quantize_mxfp8_gmem_tile_bidimensional( gX_tile: cute.Tensor, + gA_tile: Optional[cute.Tensor], gO_row_tile: Optional[cute.Tensor], gS_row_tile: Optional[cute.Tensor], gO_col_tile: Optional[cute.Tensor], gS_col_tile: Optional[cute.Tensor], ROWWISE: cutlass.Constexpr[bool], COLWISE: cutlass.Constexpr[bool], + ACTIVATION: cutlass.Constexpr[Optional[str]], + WITH_ACT: cutlass.Constexpr[bool], + WITH_DACT: cutlass.Constexpr[bool], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], @@ -462,6 +471,44 @@ def quantize_mxfp8_gmem_tile_bidimensional( load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy ) + if cutlass.const_expr(WITH_ACT or WITH_DACT): + rA = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), DTYPE) + if cutlass.const_expr(WITH_DACT): + act_load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + DTYPE, + num_bits_per_copy=min( + ELEMENTS_X_PER_THREAD * DTYPE.width, + max_num_bits_per_copy_ld_global(), + ), + invariant=False, + ) + # Activation tile has the same layout as the input tile so we can reuse the tv layout + gA_thread = cute.composition(gA_tile, tv_layout)[tidx, None] + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + cute.copy( + act_load_atom, + gA_thread[(None, row_idx),], + rA[None, row_idx], + cache_policy=cache_policy, + ) + # Apply activation row by row + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + rX_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + rA_i32 = cute.recast_tensor(rA[None, row_idx], Int32) + for i in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): + lo = x2_lo_to_f32_bf16(rX_i32[i]) + hi = x2_hi_to_f32_bf16(rX_i32[i]) + if cutlass.const_expr(WITH_DACT): + dop = SUPPORTED_DACTIVATIONS[ACTIVATION] + lo = lo * dop(x2_lo_to_f32_bf16(rA_i32[i])) + hi = hi * dop(x2_hi_to_f32_bf16(rA_i32[i])) + else: + op = SUPPORTED_ACTIVATIONS[ACTIVATION] + lo = op(lo) + hi = op(hi) + rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) + if cutlass.const_expr(ROWWISE): gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] @@ -561,7 +608,12 @@ class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. - return cfg.is_plain_cast() and cfg.DTYPE is cutlass.BFloat16 and cfg.COLWISE + return ( + not cfg.WITH_AMAX + and not cfg.WITH_DBIAS + and cfg.DTYPE is cutlass.BFloat16 + and cfg.COLWISE + ) @staticmethod @cute.jit @@ -573,8 +625,12 @@ def eligible_to_dispatch( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], ) -> Boolean: M, N = mX.shape + act_input_aligned = Boolean(True) + if cutlass.const_expr(mDActInput is not None): + act_input_aligned = mDActInput.iterator.toint() % 32 == 0 # Whole 32-row by 256-column tiles, grid.y within limits, and aligned input and # colwise-scale pointers. return ( @@ -582,6 +638,7 @@ def eligible_to_dispatch( and N % 256 == 0 and M // 32 <= 65535 and mX.iterator.toint() % 32 == 0 + and act_input_aligned and mS_col.iterator.toint() % 16 == 0 and mNoop.toint() == 0 ) @@ -612,6 +669,7 @@ def __call__( if Int64(M // TILE_ROWS) * (N // WIDE_TILE_COLS) >= 4096: self._launch( mX, + mDActInput, mO_row, mS_row, mO_col, @@ -623,6 +681,7 @@ def __call__( else: self._launch( mX, + mDActInput, mO_row, mS_row, mO_col, @@ -634,6 +693,7 @@ def __call__( else: self._launch( mX, + mDActInput, mO_row, mS_row, mO_col, @@ -647,6 +707,7 @@ def __call__( def _launch( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], mO_row: cute.Tensor, mS_row: cute.Tensor, mO_col: cute.Tensor, @@ -660,6 +721,8 @@ def _launch( # Dispatch checks these alignments before selecting this kernel. mX = cute.make_tensor(mX.iterator.align(32), mX.layout) mS_col = cute.make_tensor(mS_col.iterator.align(16), mS_col.layout) + if cutlass.const_expr(self.cfg.WITH_DACT): + mA = cute.make_tensor(mA.iterator.align(32), mA.layout) M, N = mX.shape TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD TILE_COLS = self._CTA_THREADS_X * ELEMENTS_X_PER_THREAD @@ -676,10 +739,17 @@ def _launch( block = [THREADS_PER_CTA] cluster = [cluster_cols, 1, 1] + gA = None # Hardcode the layout using static N for some of the commonly used shapes if N == 2048: payload_layout = cute.make_layout((M, 2048), stride=(2048, 1)) gX = cute.make_tensor(mX.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT): + gA = ( + cute.make_tensor(mA.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT) + else None + ) gO_row = ( cute.make_tensor(mO_row.iterator, payload_layout) if cutlass.const_expr(self.cfg.ROWWISE) @@ -705,6 +775,7 @@ def _launch( ) self.kernel( gX, + gA, gO_row, gS_row, gO_col, @@ -714,6 +785,12 @@ def _launch( elif N == 4096: payload_layout = cute.make_layout((M, 4096), stride=(4096, 1)) gX = cute.make_tensor(mX.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT): + gA = ( + cute.make_tensor(mA.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT) + else None + ) gO_row = ( cute.make_tensor(mO_row.iterator, payload_layout) if cutlass.const_expr(self.cfg.ROWWISE) @@ -739,6 +816,7 @@ def _launch( ) self.kernel( gX, + gA, gO_row, gS_row, gO_col, @@ -754,6 +832,7 @@ def _launch( gS_row, gS_col = mS_row, mS_col self.kernel( mX, + mA, mO_row, gS_row, mO_col, @@ -765,6 +844,7 @@ def _launch( def kernel( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], mO_row: cute.Tensor, mS_row: cute.Tensor, mO_col: cute.Tensor, @@ -775,6 +855,7 @@ def kernel( cast_noop flag: the dispatcher routes calls with one to the general kernel.""" self._kernel_main( mX, + mA, mO_row, mS_row, mO_col, @@ -786,6 +867,7 @@ def kernel( def _kernel_main( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], mO_row: cute.Tensor, mS_row: cute.Tensor, mO_col: cute.Tensor, @@ -804,6 +886,9 @@ def _kernel_main( COL_SCALE_TILER = (TILE_ROWS // 32, TILE_COLS) gX_tile = cute.local_tile(mX, TILER, cta_coord) + gA_tile = None + if cutlass.const_expr(self.cfg.WITH_DACT): + gA_tile = cute.local_tile(mA, TILER, cta_coord) gO_row_tile = None gS_row_tile = None if cutlass.const_expr(self.cfg.ROWWISE): @@ -816,12 +901,16 @@ def _kernel_main( quantize_mxfp8_gmem_tile_bidimensional( gX_tile, + gA_tile, gO_row_tile, gS_row_tile, gO_col_tile, gS_col_tile, ROWWISE=self.cfg.ROWWISE, COLWISE=self.cfg.COLWISE, + ACTIVATION=self.cfg.ACTIVATION, + WITH_ACT=self.cfg.WITH_ACT, + WITH_DACT=self.cfg.WITH_DACT, DTYPE=self.cfg.DTYPE, FP8_DTYPE=self.cfg.FP8_DTYPE, WITH_GEMM_SWIZZLED_SCALES=self.cfg.WITH_GEMM_SWIZZLED_SCALES, 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 97ea8c0f5f4..682ba02c611 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 @@ -231,6 +231,7 @@ def eligible_to_dispatch( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], ) -> Boolean: M, N = mX.shape # N divisible by 128 for vectorized stores, and grid.y within limits. 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 d1eb2cd0753..ee2f53c1ccb 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 @@ -17,6 +17,9 @@ create_l2_policy, exp2_bf16x2_rcp, max_num_bits_per_copy_ld_global, + x2_hi_to_f32_bf16, + x2_lo_to_f32_bf16, + cvt_f32x2_to_bf16x2, ) from transformer_engine.common.CuTeDSL.utils_fp8 import ( cvt_f32_to_fp8e8m0fnu, @@ -25,6 +28,8 @@ from .quantize_mxfp8_common import ( MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, + SUPPORTED_ACTIVATIONS, + SUPPORTED_DACTIVATIONS, derive_swizzled_scale_layout, ) @@ -32,9 +37,13 @@ @cute.jit def quantize_mxfp8_gmem_tile_rowwise_2lane( gX_tiles: cute.Tensor, + gA_tiles: Optional[cute.Tensor], gO_tiles: cute.Tensor, mS_tiles: cute.Tensor, N: Int32, + ACTIVATION: cutlass.Constexpr[Optional[str]], + WITH_ACT: cutlass.Constexpr[bool], + WITH_DACT: cutlass.Constexpr[bool], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], MAX_NORM_RCP: cutlass.Constexpr[float], @@ -72,6 +81,8 @@ def quantize_mxfp8_gmem_tile_rowwise_2lane( # We load fragments from all GMEM tiles at once rX = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, NUM_TILES), DTYPE) + # Backward: the activation input, loaded alongside the incoming gradient in rX + rA = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, NUM_TILES), DTYPE) # We store a fragment of one tile at a time rO = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), FP8_DTYPE) # RMEM flags to indicate whether each tile is out of bounds or not @@ -91,6 +102,12 @@ def quantize_mxfp8_gmem_tile_rowwise_2lane( # invariant=True here will drop the cache policy invariant=False, ) + act_load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + DTYPE, + num_bits_per_copy=LOAD_WIDTH, + invariant=False, + ) store_atom = cute.make_copy_atom( cute.nvgpu.CopyR2GOp(), FP8_DTYPE, @@ -116,8 +133,33 @@ def quantize_mxfp8_gmem_tile_rowwise_2lane( valid = col < N and Int64(row) * N + col < num_elems tile_valid_flags[tile_idx] = valid rX[None, tile_idx].fill(0) + if cutlass.const_expr(WITH_DACT): + rA[None, tile_idx].fill(0) if valid: cute.copy(load_atom, gX_thread, rX[None, tile_idx], cache_policy=input_policy) + if cutlass.const_expr(WITH_DACT): + gA_thread = cute.composition(gA_tiles[None, None, tile_idx], tv_payload)[tidx, None] + cute.copy(act_load_atom, gA_thread, rA[None, tile_idx], cache_policy=input_policy) + + # Fused activation: widen to FP32, apply it, and round back to BF16 as the general kernel + # does (it truncates the activated value to the input type before the amax). The values are + # then exactly BF16, so the packed BF16 amax and scaling below apply unchanged. + if cutlass.const_expr(WITH_ACT or WITH_DACT): + for tile_idx in cutlass.range_constexpr(NUM_TILES): + rX_i32 = cute.recast_tensor(rX[None, tile_idx], Int32) + rA_i32 = cute.recast_tensor(rA[None, tile_idx], Int32) + for i in cutlass.range_constexpr(ELEMENTS_PER_THREAD // 2): + lo = x2_lo_to_f32_bf16(rX_i32[i]) + hi = x2_hi_to_f32_bf16(rX_i32[i]) + if cutlass.const_expr(WITH_DACT): + dop = SUPPORTED_DACTIVATIONS[ACTIVATION] + lo = lo * dop(x2_lo_to_f32_bf16(rA_i32[i])) + hi = hi * dop(x2_hi_to_f32_bf16(rA_i32[i])) + else: + op = SUPPORTED_ACTIVATIONS[ACTIVATION] + lo = op(lo) + hi = op(hi) + rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) @cute.jit def extract_mx_block_amax(values_i32: cute.Tensor) -> Float32: @@ -189,7 +231,8 @@ class MXFP8QuantizeRowwise2LaneG2RKernel(MXFP8QuantizeKernelBase): def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. return ( - cfg.is_plain_cast() + not cfg.WITH_AMAX + and not cfg.WITH_DBIAS and cfg.DTYPE is cutlass.BFloat16 and cfg.ROWWISE and not cfg.COLWISE @@ -205,8 +248,12 @@ def eligible_to_dispatch( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], ) -> Boolean: N = mX.shape[1] + act_input_aligned = Boolean(True) + if cutlass.const_expr(mDActInput is not None): + act_input_aligned = mDActInput.iterator.toint() % 32 == 0 scale_stride_matches = Boolean(True) if cutlass.const_expr(not cfg.WITH_GEMM_SWIZZLED_SCALES): # The flat view needs the packed scales to be contiguous. @@ -216,6 +263,7 @@ def eligible_to_dispatch( return ( N % 128 == 0 and mX.iterator.toint() % 32 == 0 + and act_input_aligned and scale_stride_matches and mNoop.toint() == 0 ) @@ -246,22 +294,63 @@ def __call__( if num_elems <= 2**31 - 1: rows = cute.ceil_div(num_elems, 1024) flat = cute.make_layout((rows, 1024), stride=(1024, 1)) + mA = None + if cutlass.const_expr(self.cfg.WITH_DACT): + mA = cute.make_tensor(mDActInput.iterator, flat) self._launch_tiers( cute.make_tensor(mX.iterator, flat), + mA, cute.make_tensor(mO_row.iterator, flat), cute.make_tensor(mS_row.iterator, cute.make_layout((rows, 32), stride=(32, 1))), stream, num_elems, ) else: - self._launch_tiers(mX, mO_row, mS_row, stream, num_elems) + self._launch_tiers(mX, mDActInput, mO_row, mS_row, stream, num_elems) else: - self._launch_tiers(mX, mO_row, mS_row, stream, num_elems) + self._launch_tiers(mX, mDActInput, mO_row, mS_row, stream, num_elems) @cute.jit def _launch_tiers( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], + mO_row: cute.Tensor, + mS_row: cute.Tensor, + stream: CUstream, + num_elems: Int64, + ): + """Pick the launch configuration: the size tiers, except for swizzled scales with a fused + activation when N does not fill the tiers' 1024-column CTAs.""" + if cutlass.const_expr( + self.cfg.WITH_GEMM_SWIZZLED_SCALES and (self.cfg.WITH_ACT or self.cfg.WITH_DACT) + ): + if mX.shape[1] % 1024 != 0: + # The swizzled path tiles the real (M, N) shape, so a 1024-column CTA overhangs + # N and its idle lanes still pay for the activation math. 256-column CTAs divide + # every eligible N (N % 128 == 0 and the tile is 2 x 128 columns). + self._launch( + mX, + mA, + mO_row, + mS_row, + stream, + num_elems, + CTA_THREADS_Y=16, + CTA_THREADS_X=8, + NUM_TILES=2, + L2_CACHED_CTA_PERCENT=0, + ) + else: + self._launch_size_tiers(mX, mA, mO_row, mS_row, stream, num_elems) + else: + self._launch_size_tiers(mX, mA, mO_row, mS_row, stream, num_elems) + + @cute.jit + def _launch_size_tiers( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], mO_row: cute.Tensor, mS_row: cute.Tensor, stream: CUstream, @@ -272,6 +361,7 @@ def _launch_tiers( if output_bytes <= 12 << 20: self._launch( mX, + mA, mO_row, mS_row, stream, @@ -284,6 +374,7 @@ def _launch_tiers( elif output_bytes <= 48 << 20: self._launch( mX, + mA, mO_row, mS_row, stream, @@ -296,6 +387,7 @@ def _launch_tiers( elif output_bytes <= 96 << 20: self._launch( mX, + mA, mO_row, mS_row, stream, @@ -308,6 +400,7 @@ def _launch_tiers( else: self._launch( mX, + mA, mO_row, mS_row, stream, @@ -322,6 +415,7 @@ def _launch_tiers( def _launch( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], mO: cute.Tensor, mS: cute.Tensor, stream: CUstream, @@ -361,6 +455,7 @@ def _launch( if check_bounds: self.kernel( mX, + mA, mO, mS, first_streaming_cta, @@ -374,6 +469,7 @@ def _launch( else: self.kernel( mX, + mA, mO, mS, first_streaming_cta, @@ -389,6 +485,7 @@ def _launch( def kernel( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], mO: cute.Tensor, mS: cute.Tensor, first_streaming_cta: Int64, @@ -403,6 +500,7 @@ def kernel( dispatcher only selects this kernel when no cast_noop flag is passed.""" self._kernel_main( mX, + mA, mO, mS, first_streaming_cta, @@ -418,6 +516,7 @@ def kernel( def _kernel_main( self, mX: cute.Tensor, + mA: Optional[cute.Tensor], mO: cute.Tensor, mS: cute.Tensor, first_streaming_cta: Int64, @@ -431,6 +530,8 @@ def _kernel_main( """Construct CTA tile views and run the specialized quantization.""" # Dispatch checks 32-byte input alignment before selecting this kernel. gX = cute.make_tensor(mX.iterator.align(32), mX.layout) + if cutlass.const_expr(self.cfg.WITH_DACT): + mA = cute.make_tensor(mA.iterator.align(32), mA.layout) # grid.x walks column tiles; grid.y and grid.z together walk row tiles. bidx, bidy, bidz = cute.arch.block_idx() cta_row = bidz * self._MAX_GRID_Y + bidy @@ -472,6 +573,9 @@ def _kernel_main( SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) # Mode 1's first index is 0 because CTA_ROWS == TILE_ROWS (there is only 1 tile on the row dimension). gX_tiles = cute.local_tile(gX_cta, TILER, (0, None)) + gA_tiles = None + if cutlass.const_expr(self.cfg.WITH_DACT): + gA_tiles = cute.local_tile(cute.local_tile(mA, CTA_TILER, cta_coord), TILER, (0, None)) gO_tiles = cute.local_tile(gO_cta, TILER, (0, None)) mS_tiles = cute.local_tile(gS_cta, SCALE_TILER, (0, None)) @@ -479,9 +583,13 @@ def _kernel_main( if cta_row * CTA_ROWS < M: quantize_mxfp8_gmem_tile_rowwise_2lane( gX_tiles, + gA_tiles, gO_tiles, mS_tiles, N, + ACTIVATION=self.cfg.ACTIVATION, + WITH_ACT=self.cfg.WITH_ACT, + WITH_DACT=self.cfg.WITH_DACT, DTYPE=self.cfg.DTYPE, FP8_DTYPE=self.cfg.FP8_DTYPE, MAX_NORM_RCP=self.cfg.MAX_NORM_RCP, 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 65bf548022a..1192d98cb3a 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 @@ -270,6 +270,7 @@ def eligible_to_dispatch( mO_col: Optional[cute.Tensor], mS_col: Optional[cute.Tensor], mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], ) -> Boolean: cls = MXFP8QuantizeBidimensionalG2SKernel rows_per_cta = cls._TILE_ROWS * cls._NUM_TILES_Y # packed scales: one tile row diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index caaf8339524..90c61d03052 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -286,6 +286,16 @@ def x2_hi_to_f32_bf16(bits: Int32) -> Float32: return ((bits >> Int32(16)) << Int32(16)).bitcast(Float32) +@cute.jit +def cvt_f32x2_to_bf16x2(lo: Float32, hi: Float32) -> Int32: + """Round two FP32 values to BF16 (round-to-nearest-even) and pack them, lo in the low half.""" + return cute.arch.inline_ptx( + "cvt.rn.bf16x2.f32 {$w0}, {$r0}, {$r1};", + write_only_types=[Int32], + read_only_args=[hi, lo], + ) + + @cute.jit def truncate_f32_bf16(val: Float32) -> Float32: """Round FP32 to BF16 precision (round-to-nearest-even), keeping FP32 storage.""" From b1db277941e6221817f3d86d718206d9e78ab00b Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 13:44:36 -0700 Subject: [PATCH 25/29] fix Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 23 ++++++++----------- 1 file changed, 10 insertions(+), 13 deletions(-) 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 0e8efdb3bd3..776fb1db21e 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 @@ -739,17 +739,15 @@ def _launch( block = [THREADS_PER_CTA] cluster = [cluster_cols, 1, 1] - gA = None # Hardcode the layout using static N for some of the commonly used shapes if N == 2048: payload_layout = cute.make_layout((M, 2048), stride=(2048, 1)) gX = cute.make_tensor(mX.iterator, payload_layout) - if cutlass.const_expr(self.cfg.WITH_DACT): - gA = ( - cute.make_tensor(mA.iterator, payload_layout) - if cutlass.const_expr(self.cfg.WITH_DACT) - else None - ) + gA = ( + cute.make_tensor(mA.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT) + else None + ) gO_row = ( cute.make_tensor(mO_row.iterator, payload_layout) if cutlass.const_expr(self.cfg.ROWWISE) @@ -785,12 +783,11 @@ def _launch( elif N == 4096: payload_layout = cute.make_layout((M, 4096), stride=(4096, 1)) gX = cute.make_tensor(mX.iterator, payload_layout) - if cutlass.const_expr(self.cfg.WITH_DACT): - gA = ( - cute.make_tensor(mA.iterator, payload_layout) - if cutlass.const_expr(self.cfg.WITH_DACT) - else None - ) + gA = ( + cute.make_tensor(mA.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT) + else None + ) gO_row = ( cute.make_tensor(mO_row.iterator, payload_layout) if cutlass.const_expr(self.cfg.ROWWISE) From dcd8b08ed7491d7772863c8919e27218f0d23cfe Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 14:51:01 -0700 Subject: [PATCH 26/29] fix Signed-off-by: Kaining Zhong --- transformer_engine/common/CuTeDSL/utils_fp8.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/utils_fp8.py b/transformer_engine/common/CuTeDSL/utils_fp8.py index 9a972e39469..61355cca1f2 100644 --- a/transformer_engine/common/CuTeDSL/utils_fp8.py +++ b/transformer_engine/common/CuTeDSL/utils_fp8.py @@ -183,11 +183,11 @@ def mul_f32x4_cvt_f32x4_to_fp8x4(fp8_dtype, relu: bool = False) -> Callable[..., def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: """Return a multiply/cast for two BF16 pairs, producing four packed FP8 bytes. - CUDA 13.1+ supports direct BF16-to-FP8 conversion. Older compilers use - FP32 fused multiply-add, matching the corresponding ptx::mul_cvt_4x overload. + Uses a direct BF16-to-FP8 conversion on CUDA 13.1+ and an FP32 detour before that. """ out_op = "e5m2x2" if fp8_dtype is Float8E5M2 else "e4m3x2" if cutlass.target_version(min_version="13.1"): + # Multiply in BF16, then convert BF16 -> FP8. body = ( ".reg.b32 y0, y1;\n\t" "mul.rn.bf16x2 y0, {$r0}, {$r1};\n\t" @@ -196,6 +196,8 @@ def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: f"cvt.rn.satfinite.{out_op}.bf16x2 z1, y1;\n\t" ) else: + # No BF16 -> FP8 conversion before CUDA 13.1: fma.rn.f32.bf16 multiplies into FP32 + # (a*b + -0.0 is exact and keeps -0.0), then convert FP32 -> FP8. body = ( ".reg.b16 x0,x1,x2,x3,s0,s1,s2,s3;\n\t" ".reg.f32 y0,y1,y2,y3;\n\t" @@ -203,10 +205,10 @@ def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: "mov.b32 {s0,s1}, {$r1};\n\t" "mov.b32 {x2,x3}, {$r2};\n\t" "mov.b32 {s2,s3}, {$r3};\n\t" - "fma.rn.f32.bf16 y0, x0, s0, 0f00000000;\n\t" - "fma.rn.f32.bf16 y1, x1, s1, 0f00000000;\n\t" - "fma.rn.f32.bf16 y2, x2, s2, 0f00000000;\n\t" - "fma.rn.f32.bf16 y3, x3, s3, 0f00000000;\n\t" + "fma.rn.f32.bf16 y0, x0, s0, 0f80000000;\n\t" + "fma.rn.f32.bf16 y1, x1, s1, 0f80000000;\n\t" + "fma.rn.f32.bf16 y2, x2, s2, 0f80000000;\n\t" + "fma.rn.f32.bf16 y3, x3, s3, 0f80000000;\n\t" f"cvt.rn.satfinite.{out_op}.f32 z0, y1, y0;\n\t" f"cvt.rn.satfinite.{out_op}.f32 z1, y3, y2;\n\t" ) From 86233c813ea4f8b572b23e501cb1639198319cf9 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 16:51:00 -0700 Subject: [PATCH 27/29] +0 and -0 should both legal Signed-off-by: Kaining Zhong --- tests/jax/test_mxfp8_cutedsl_backend.py | 5 ++++- tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py | 6 ++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/jax/test_mxfp8_cutedsl_backend.py b/tests/jax/test_mxfp8_cutedsl_backend.py index bd492c967a1..2ee82f8c5d5 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -197,12 +197,15 @@ def extract_quantized_output(out, dbias): ScaledTensor1x carries scale_inv already trimmed to the unpadded shape (see its __post_init__), so there is no uninitialized scale padding to exclude. + FP8 -0.0 (0x80) is folded into +0.0: both backends round to zero there and may differ only + in its sign, which training does not depend on. """ tensors = [out] if isinstance(out, ScaledTensor1x) else [out.rowwise_tensor, out.colwise_tensor] parts = {} for t in tensors: name = "colwise" if t.is_colwise else "rowwise" - parts[f"{name} data"] = np.asarray(t.data.view(jnp.uint8)) + data = np.asarray(t.data.view(jnp.uint8)) + parts[f"{name} data"] = np.where(data == 0x80, 0, data).astype(np.uint8) parts[f"{name} scales"] = np.asarray(t.scale_inv.view(jnp.uint8)) return parts, None if dbias is None else np.asarray(dbias) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index a886b63c088..5d7feeb6ce6 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -251,18 +251,20 @@ def extract_quantized_output(out, rowwise, columnwise, swizzled): compared. Swizzled layout: the meaningful scales are scattered by the swizzle, so the top-left slice is meaningless; both backends zero the padding (see zero_scales_kernel), so compare the whole buffer instead -- which also covers the padding-zeroing itself. + FP8 -0.0 (0x80) is folded into +0.0: both backends round to zero there and may differ only + in its sign, which training does not depend on. """ parts = {} if rowwise: d = out._rowwise_data.view(torch.uint8) M, N = d.shape - parts["rowwise data"] = d.clone() + parts["rowwise data"] = d.masked_fill(d == 0x80, 0) s = out._rowwise_scale_inv parts["rowwise scales"] = (s if swizzled else s[:M, : (N + 31) // 32]).clone() if columnwise: d = out._columnwise_data.view(torch.uint8) M, N = d.shape - parts["colwise data"] = d.clone() + parts["colwise data"] = d.masked_fill(d == 0x80, 0) s = out._columnwise_scale_inv parts["colwise scales"] = (s if swizzled else s[: (M + 31) // 32, :N]).clone() return parts From dba8dbf3f408053cbc627873be096c4ba2aee513 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 16:57:58 -0700 Subject: [PATCH 28/29] nit Signed-off-by: Kaining Zhong --- .../CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 2 +- transformer_engine/common/CuTeDSL/utils_fp8.py | 14 ++++++-------- 2 files changed, 7 insertions(+), 9 deletions(-) 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 776fb1db21e..61598a7d37b 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 @@ -8,7 +8,7 @@ import cutlass from cutlass import cute -from cutlass import Boolean, Float32, Int32, Int64, Uint32, Uint8 +from cutlass import Boolean, Int32, Int64, Uint32, Uint8 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( diff --git a/transformer_engine/common/CuTeDSL/utils_fp8.py b/transformer_engine/common/CuTeDSL/utils_fp8.py index 61355cca1f2..9a972e39469 100644 --- a/transformer_engine/common/CuTeDSL/utils_fp8.py +++ b/transformer_engine/common/CuTeDSL/utils_fp8.py @@ -183,11 +183,11 @@ def mul_f32x4_cvt_f32x4_to_fp8x4(fp8_dtype, relu: bool = False) -> Callable[..., def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: """Return a multiply/cast for two BF16 pairs, producing four packed FP8 bytes. - Uses a direct BF16-to-FP8 conversion on CUDA 13.1+ and an FP32 detour before that. + CUDA 13.1+ supports direct BF16-to-FP8 conversion. Older compilers use + FP32 fused multiply-add, matching the corresponding ptx::mul_cvt_4x overload. """ out_op = "e5m2x2" if fp8_dtype is Float8E5M2 else "e4m3x2" if cutlass.target_version(min_version="13.1"): - # Multiply in BF16, then convert BF16 -> FP8. body = ( ".reg.b32 y0, y1;\n\t" "mul.rn.bf16x2 y0, {$r0}, {$r1};\n\t" @@ -196,8 +196,6 @@ def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: f"cvt.rn.satfinite.{out_op}.bf16x2 z1, y1;\n\t" ) else: - # No BF16 -> FP8 conversion before CUDA 13.1: fma.rn.f32.bf16 multiplies into FP32 - # (a*b + -0.0 is exact and keeps -0.0), then convert FP32 -> FP8. body = ( ".reg.b16 x0,x1,x2,x3,s0,s1,s2,s3;\n\t" ".reg.f32 y0,y1,y2,y3;\n\t" @@ -205,10 +203,10 @@ def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: "mov.b32 {s0,s1}, {$r1};\n\t" "mov.b32 {x2,x3}, {$r2};\n\t" "mov.b32 {s2,s3}, {$r3};\n\t" - "fma.rn.f32.bf16 y0, x0, s0, 0f80000000;\n\t" - "fma.rn.f32.bf16 y1, x1, s1, 0f80000000;\n\t" - "fma.rn.f32.bf16 y2, x2, s2, 0f80000000;\n\t" - "fma.rn.f32.bf16 y3, x3, s3, 0f80000000;\n\t" + "fma.rn.f32.bf16 y0, x0, s0, 0f00000000;\n\t" + "fma.rn.f32.bf16 y1, x1, s1, 0f00000000;\n\t" + "fma.rn.f32.bf16 y2, x2, s2, 0f00000000;\n\t" + "fma.rn.f32.bf16 y3, x3, s3, 0f00000000;\n\t" f"cvt.rn.satfinite.{out_op}.f32 z0, y1, y0;\n\t" f"cvt.rn.satfinite.{out_op}.f32 z1, y3, y2;\n\t" ) From 287505f070f1cd9cc8691425aa161d1b0192fe5f Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Sat, 17 Oct 2026 02:12:54 +0000 Subject: [PATCH 29/29] fix Signed-off-by: Kaining Zhong --- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 145 ++++++++++++++---- 1 file changed, 114 insertions(+), 31 deletions(-) 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 ee2f53c1ccb..03d5d510ab8 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 @@ -34,6 +34,37 @@ ) +@cute.jit +def _flush_scales_to_gmem( + sScale: cute.Tensor, + gS_tile: cute.Tensor, + cta_row: Int32, + cta_col: Int32, + M: Int32, + padded_cols: Int32, + CTA_Y: cutlass.Constexpr[int], + CTA_X: cutlass.Constexpr[int], +): + """Flush each 4-row by 4-column scale block with a 16-byte store. + + sScale and gS_tile must have byte strides (4, 1) within each aligned (4, 4) row-column block. + """ + tidx, _, _ = cute.arch.thread_idx() + active_threads = CTA_X // 4 + _, tv_flush = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y // 4, active_threads), stride=(active_threads, 1)), + val_layout=cute.make_layout((4, 4), stride=(4, 1)), + ) + if tidx < (CTA_Y // 4) * active_threads: + thread_y = cta_row * CTA_Y + (tidx // active_threads) * 4 + thread_x = cta_col * CTA_X + (tidx % active_threads) * 4 + if thread_y < M and thread_x < padded_cols: + cute.autovec_copy( + cute.composition(sScale, tv_flush)[tidx, None], + cute.composition(gS_tile, tv_flush)[tidx, None], + ) + + @cute.jit def quantize_mxfp8_gmem_tile_rowwise_2lane( gX_tiles: cute.Tensor, @@ -214,8 +245,8 @@ class MXFP8QuantizeRowwise2LaneG2RKernel(MXFP8QuantizeKernelBase): One CTA tile has TILE_ROWS = CTA_THREADS_Y and TILE_COLS = CTA_THREADS_X * _ELEMENTS_PER_THREAD. The CTA - processes NUM_TILES = 1 or 2 in the original input matrix. Both scale - formats use the same payload tiling; swizzled scale offsets are computed directly. + processes NUM_TILES = 1 or 2 horizontally. With staged swizzled scales, + tensor views group rows 32 apart to collect complete scale vectors. _L2_CACHED_CTA_PERCENT tunes the share of CTAs whose input loads use normal L2 retention: 0 streams all input and 100 caches all input normally. """ @@ -451,35 +482,49 @@ def _launch( grid_y = cutlass.min(grid_rows, self._MAX_GRID_Y) grid_z = cute.ceil_div(grid_rows, self._MAX_GRID_Y) launch_grid = [grid_cols, grid_y, grid_z] + kernel_args = ( + mX, + mA, + mO, + mS, + first_streaming_cta, + num_elems, + CTA_THREADS_Y, + CTA_THREADS_X, + NUM_TILES, + L2_CACHED_CTA_PERCENT, + ) + launch_args = {"grid": launch_grid, "block": [THREADS_PER_CTA], "stream": stream} + stash_scale_to_smem = False + if cutlass.const_expr( + self.cfg.WITH_GEMM_SWIZZLED_SCALES and CTA_THREADS_Y % 4 == 0 and not self.cfg.WITH_DACT + ): + # Staged scale vectors need complete 128-row bands. + stash_scale_to_smem = M % 128 == 0 + if grid > 0: if check_bounds: - self.kernel( - mX, - mA, - mO, - mS, - first_streaming_cta, - num_elems, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - NUM_TILES=NUM_TILES, - L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, - CHECK_BOUNDS=True, - ).launch(grid=launch_grid, block=[THREADS_PER_CTA], stream=stream) + if stash_scale_to_smem: + self._launch_kernel(kernel_args, launch_args, True, True) + else: + self._launch_kernel(kernel_args, launch_args, True, False) else: - self.kernel( - mX, - mA, - mO, - mS, - first_streaming_cta, - num_elems, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - NUM_TILES=NUM_TILES, - L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, - CHECK_BOUNDS=False, - ).launch(grid=launch_grid, block=[THREADS_PER_CTA], stream=stream) + if stash_scale_to_smem: + self._launch_kernel(kernel_args, launch_args, False, True) + else: + self._launch_kernel(kernel_args, launch_args, False, False) + + @cute.jit + def _launch_kernel( + self, + kernel_args, + launch_args, + CHECK_BOUNDS: cutlass.Constexpr[bool], + STASH_SCALE_TO_SMEM: cutlass.Constexpr[bool], + ): + self.kernel( + *kernel_args, CHECK_BOUNDS=CHECK_BOUNDS, STASH_SCALE_TO_SMEM=STASH_SCALE_TO_SMEM + ).launch(**launch_args) @cute.kernel def kernel( @@ -495,6 +540,7 @@ def kernel( NUM_TILES: cutlass.Constexpr[int], L2_CACHED_CTA_PERCENT: cutlass.Constexpr[int], CHECK_BOUNDS: cutlass.Constexpr[bool], + STASH_SCALE_TO_SMEM: cutlass.Constexpr[bool], ): """Run the kernel body. Like cast_rowwise.cuh, there is no noop handling: the dispatcher only selects this kernel when no cast_noop flag is passed.""" @@ -510,6 +556,7 @@ def kernel( NUM_TILES=NUM_TILES, L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, CHECK_BOUNDS=CHECK_BOUNDS, + STASH_SCALE_TO_SMEM=STASH_SCALE_TO_SMEM, ) @cute.jit @@ -526,6 +573,7 @@ def _kernel_main( NUM_TILES: cutlass.Constexpr[int], L2_CACHED_CTA_PERCENT: cutlass.Constexpr[int], CHECK_BOUNDS: cutlass.Constexpr[bool], + STASH_SCALE_TO_SMEM: cutlass.Constexpr[bool], ): """Construct CTA tile views and run the specialized quantization.""" # Dispatch checks 32-byte input alignment before selecting this kernel. @@ -549,7 +597,19 @@ def _kernel_main( if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): gS, _ = derive_swizzled_scale_layout(M, N, True, False, mS, None) - # CTA_TILER is the block that a CTA processes in total, which consists of NUM_TILES horizontal tiles + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + # Apply the same permutation to input, output and scale so it becomes transparent and + # other code is unaware of this permutation pattern + # However, this this permutation, the scale bytes emitted by a CTA are all in this CTA's scale + # tile partition instead of scatter everywhere, so we can vectorize the flush to GMEM + # For example, with a 128x128 input: + # gX & mO: (128,128):(128,1) -> ((4,32),128):((4096,128),1) + # gS: ((32,4,1),(4,1)):((16,4,512),(1,512)) -> ((4,32), 4):((4, 16),1) + row_perm = cute.make_layout((4, 32, M // 128), stride=(32, 1, 128)) + gX = cute.composition(gX, (row_perm, None)) + mO = cute.composition(mO, (row_perm, None)) + gS = cute.composition(gS, (row_perm, None)) + CTA_TILER = (CTA_ROWS, CTA_COLS) CTA_SCALE_TILER = (CTA_ROWS, CTA_COLS // 32) gX_cta = cute.local_tile(gX, CTA_TILER, cta_coord) @@ -571,13 +631,23 @@ def _kernel_main( # Divide CTA_TILER to multiple TILER shape TILES TILER = (TILE_ROWS, TILE_COLS) SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) - # Mode 1's first index is 0 because CTA_ROWS == TILE_ROWS (there is only 1 tile on the row dimension). + # There is only one tile on the row dimension because CTA_ROWS == TILE_ROWS. gX_tiles = cute.local_tile(gX_cta, TILER, (0, None)) gA_tiles = None if cutlass.const_expr(self.cfg.WITH_DACT): gA_tiles = cute.local_tile(cute.local_tile(mA, CTA_TILER, cta_coord), TILER, (0, None)) gO_tiles = cute.local_tile(gO_cta, TILER, (0, None)) - mS_tiles = cute.local_tile(gS_cta, SCALE_TILER, (0, None)) + + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + # Preserve gS_cta's (4 rows, 4 columns):(4, 1) block ordering for the flush. + sScale_scratch = cutlass.memory.SmemAllocator().allocate_tensor( + mS.element_type, cute.make_layout_like(gS_cta), byte_alignment=16 + ) + # The tile quantizer writes scales to SMEM for the cooperative vectorized flush. + mS_tiles = cute.local_tile(sScale_scratch, SCALE_TILER, (0, None)) + else: + # The tile quantizer writes scales directly to GMEM. + mS_tiles = cute.local_tile(gS_cta, SCALE_TILER, (0, None)) # The last grid.z slice can overshoot the row tiles; those CTAs have no work. if cta_row * CTA_ROWS < M: @@ -604,3 +674,16 @@ def _kernel_main( cta_col=cta_coord[1], num_elems=num_elems, ) + + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + cute.arch.sync_threads() + _flush_scales_to_gmem( + sScale_scratch, + cute.make_tensor(gS_cta.iterator.align(16), gS_cta.layout), + cta_row, + bidx, + M, + N // 32, + CTA_ROWS, + CTA_COLS // 32, + )