From 7c3abf6bdc9442b8645e4d639c603c1b8fc98cdc Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 7 Oct 2026 07:22:07 +0000 Subject: [PATCH 01/43] [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/43] 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/43] 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/43] 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/43] 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/43] 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/43] [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/43] 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/43] 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/43] 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/43] 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/43] [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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] +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/43] 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/43] 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, + ) From 2c6f6ec209fded2287b0bea8067767dbac7e5f7d Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 19 Oct 2026 21:23:08 -0700 Subject: [PATCH 30/43] fix Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_common.py | 11 ++++++++- .../cast/mxfp8/quantize_mxfp8_g2s_general.py | 23 ++++++++++++------- 2 files changed, 25 insertions(+), 9 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 0df60e0714e..f003a20df30 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -341,10 +341,19 @@ def __call__( ) # If not using a specialized kernel, fall back to the general kernel if not dispatched_to_specialized: + # With _USE_STATIC_NUM_TILES, each CTA always processes _NUM_TILES tiles, including OOB ones, + # so the scale-bounds check can only be dropped if M divides TILE_ROWS * _NUM_TILES. + num_tiles_per_cta = ( + self.general_skip_input_masking_kernel._NUM_TILES + if self.general_skip_input_masking_kernel._USE_STATIC_NUM_TILES + else 1 + ) # If the input shape can be perfectly tiled by the general kernel's tile size, we can skip some boundary check because # we know we will not touch any out-of-bounds region. shape_is_divisible = ( - mX.shape[0] % self.general_skip_input_masking_kernel._TILE_ROWS == 0 + mX.shape[0] + % (self.general_skip_input_masking_kernel._TILE_ROWS * num_tiles_per_cta) + == 0 and mX.shape[1] % self.general_skip_input_masking_kernel._TILE_COLS == 0 ) # If shape is divisible, then we won't access OOB regions regardless whatever diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py index 9801d82a259..9ffd0a6e4df 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py @@ -296,8 +296,7 @@ def _quantize_rowwise_mxfp8( # For irregular shapes, skip the scale store if this thread's logical row / col-block lies past the input's actual extents. # TMA already zero-fills OOB input reads and drops OOB output writes; only the direct scale-byte gmem store needs an explicit guard. - # If the input shape is divisible by the tile size, then we won't access OOB regions because - # we never we only access tiles we actually need (num_tiles) which are never OOB + # If the input shape is divisible by the tile size * tile num, then we won't access OOB regions for sure if cutlass.const_expr(SKIP_SCALE_BOUNDS): mS_row_stage[(tidx // CTA_THREADS_X, tidx % CTA_THREADS_X)] = biased_exp_r else: @@ -457,8 +456,7 @@ def _quantize_colwise_mxfp8( # column lies past the input extents. TILE_Y == MXFP8_BLOCK_SCALING_SIZE so each stage # is exactly one scale-row; valid iff `tile_row_start < M`. biased_exp_c = cvt_f32_to_fp8e8m0fnu(amax_c * MAX_NORM_RCP) - # If the input shape is divisible by the tile size, then we won't access OOB regions because - # we never we only access tiles we actually need (num_tiles) which are never OOB + # If the input shape is divisible by the tile size * tile num, then we won't access OOB regions for sure if cutlass.const_expr(SKIP_SCALE_BOUNDS): if cutlass.const_expr(SWIZZLE): mS_col_stage[(0, tidx % 32, tidx // 32)] = biased_exp_c @@ -685,6 +683,10 @@ def __init__( else: self._NUM_TILES = 2 # Each CTA handles 2 tiles stacked vertically self._THREADS_PER_CTA = 64 + # If _USE_STATIC_NUM_TILES, then each CTA will always process _NUM_TILES tiles without checking if those tiles + # are OOB (which will be zero padded by TMA). + # Benchmarks show this is in general faster than the dynamic count except when DACT is fused + self._USE_STATIC_NUM_TILES = not cfg.WITH_DACT # Each thread handles a (1, MXFP8_BLOCK_SCALING_SIZE) chunk self._TILE_COLS = self._THREADS_PER_CTA self._TILE_ROWS = MXFP8_BLOCK_SCALING_SIZE @@ -1062,10 +1064,15 @@ class DactStorage: M = mX.shape[0] N = mX.shape[1] - num_tiles = cutlass.min( - self._NUM_TILES, - cute.ceil_div(M - bidy * self._TILE_ROWS * self._NUM_TILES, self._TILE_ROWS), - ) + # With static num tiles, always process _NUM_TILES tiles like the CUDA kernel: TMA zero-fills + # OOB tiles, so the loop trip count is a compile-time constant. + if cutlass.const_expr(self._USE_STATIC_NUM_TILES): + num_tiles = self._NUM_TILES + else: + num_tiles = cutlass.min( + self._NUM_TILES, + cute.ceil_div(M - bidy * self._TILE_ROWS * self._NUM_TILES, self._TILE_ROWS), + ) # Tile the TMA gmem view: ((TILE_Y, TILE_X), (M/TILE_Y, N/TILE_X)). gX_tiled = cute.zipped_divide(tma_src, (self._TILE_ROWS, self._TILE_COLS)) From 53bc64775c252404ac1405a74231e93f1fd5f18c Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 21 Oct 2026 10:54:40 -0700 Subject: [PATCH 31/43] close fp32 perf gap Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_g2s_general.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py index 9ffd0a6e4df..3ef1704f82a 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py @@ -687,6 +687,8 @@ def __init__( # are OOB (which will be zero padded by TMA). # Benchmarks show this is in general faster than the dynamic count except when DACT is fused self._USE_STATIC_NUM_TILES = not cfg.WITH_DACT + # A hint that is tuned based on the benchmark which helps when ACT/DACT is fused + self._MIN_BLOCKS_PER_MP = 1 if (cfg.WITH_ACT or cfg.WITH_DACT) else 0 # Each thread handles a (1, MXFP8_BLOCK_SCALING_SIZE) chunk self._TILE_COLS = self._THREADS_PER_CTA self._TILE_ROWS = MXFP8_BLOCK_SCALING_SIZE @@ -831,6 +833,7 @@ def __call__( ).launch( grid=grid, block=block, + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, stream=stream, ) @@ -1172,10 +1175,12 @@ class DactStorage: 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() + # that is read _NUM_STAGES iterations ago so we must wait for whoever is reading that buffer to finish. + # Note: no sO buffer is reused when _NUM_TILES <= _NUM_STAGES, so we can always skip waiting when this holds + if cutlass.const_expr(self._NUM_TILES > self._NUM_STAGES): + 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 @@ -1312,9 +1317,9 @@ class DactStorage: sAmax = storage.sAmax.get_tensor(cute.make_layout(self._NUM_WARPS)) self._amax_epilogue(sAmax, mAmax, tidx, warp_idx, per_thread_amax) - # Wait for in-flight TMA stores so data is visible to the host - # before the kernel returns. - cute.arch.cp_async_bulk_wait_group(0, read=False) + # Wait only until in-flight TMA stores have read SMEM so the CTA retires (frees its SMEM) early; + # the grid's completion still guarantees the global writes are done. + cute.arch.cp_async_bulk_wait_group(0, read=True) @cute.jit def _dbias_reduction_rowwise_epilouge( From bfdb240d45ffde11aba06a196faed0bf3b3025c9 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 21 Oct 2026 11:30:36 -0700 Subject: [PATCH 32/43] can always use static tile num Signed-off-by: Kaining Zhong --- .../common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py index 3ef1704f82a..2ea9fcc8e9b 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py @@ -684,9 +684,8 @@ def __init__( self._NUM_TILES = 2 # Each CTA handles 2 tiles stacked vertically self._THREADS_PER_CTA = 64 # If _USE_STATIC_NUM_TILES, then each CTA will always process _NUM_TILES tiles without checking if those tiles - # are OOB (which will be zero padded by TMA). - # Benchmarks show this is in general faster than the dynamic count except when DACT is fused - self._USE_STATIC_NUM_TILES = not cfg.WITH_DACT + # are OOB (which will be zero padded by TMA). Hardcoded to true because benchmark shows this is always better + self._USE_STATIC_NUM_TILES = True # A hint that is tuned based on the benchmark which helps when ACT/DACT is fused self._MIN_BLOCKS_PER_MP = 1 if (cfg.WITH_ACT or cfg.WITH_DACT) else 0 # Each thread handles a (1, MXFP8_BLOCK_SCALING_SIZE) chunk From 1cbe8f33a20c5ca4454cd98ece3748ac789623ba Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 29 Oct 2026 18:32:48 +0000 Subject: [PATCH 33/43] draft: gated mxfp8 cutedsl kernel Signed-off-by: Kaining Zhong --- tests/jax/test_mxfp8_cutedsl_backend.py | 119 ++- .../mxfp8/test_mxfp8_cutedsl_backend.py | 128 ++- transformer_engine/common/CuTeDSL/__init__.py | 9 +- .../common/CuTeDSL/activations.py | 54 ++ .../cast/mxfp8/quantize_mxfp8_common.py | 372 ++++---- .../cast/mxfp8/quantize_mxfp8_gated_g2s.py | 861 ++++++++++++++++++ .../common/cast/dispatch/gated.cuh | 28 +- .../common/cast/mxfp8/gated_mxfp8_cutedsl.cuh | 205 +++++ .../common/util/cutedsl_utils.h | 70 ++ 9 files changed, 1673 insertions(+), 173 deletions(-) create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py create mode 100644 transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh diff --git a/tests/jax/test_mxfp8_cutedsl_backend.py b/tests/jax/test_mxfp8_cutedsl_backend.py index 2ee82f8c5d5..2529cbfadfa 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -23,6 +23,7 @@ from transformer_engine.common import _get_shared_object_file from transformer_engine.common.CuTeDSL.utils import device_compute_capability from transformer_engine.jax import cpp_extensions as tex +from transformer_engine.jax.cpp_extensions.activation import ActivationParams from transformer_engine.jax.quantize import ( QuantizerFactory, QuantizeLayout, @@ -68,8 +69,8 @@ # Every activation the CuTeDSL backend implements, as (JAX activation_type, name it goes by # in the config key). The two differ for qgelu/srelu: JAX spells them quick_gelu/squared_relu # (ActivationEnum in jax/cpp_extensions/activation.py), the backend uses the C++ spelling -# (Activation in common/util/cutedsl_utils.h). Gated variants are absent because they go to a -# separate TE/common kernel that the backend does not cover. +# (Activation in common/util/cutedsl_utils.h). Gated variants go to a separate kernel and are +# covered by GATED_ACTIVATIONS below. ACTIVATIONS = [ (("gelu",), "gelu"), (("relu",), "relu"), @@ -316,3 +317,117 @@ def test_dtypes(method, act_type, act_desc, fp8_dtype, in_dtype): ) def test_specialized_kernels(shape, q_layout, in_dtype): run_test_case("CAST_ONLY", *NO_ACT, shape, q_layout, in_dtype, jnp.float8_e4m3fn) + + +# Gated activations (gated_mxfp8.cuh), as (JAX activation_type, name in the config key +# (GatedActivation in common/util/cutedsl_utils.h), activation parameters). The clamped SwiGLU +# parameters are chosen so that its clamps are active on the [-2, 2] inputs generate_inputs draws. +GATED_ACTIVATIONS = [ + (("gelu", "linear"), "geglu", {}), + (("quick_gelu", "linear"), "qgeglu", {}), + (("relu", "linear"), "reglu", {}), + (("squared_relu", "linear"), "sreglu", {}), + (("silu", "linear"), "swiglu", {}), + (("sigmoid", "linear"), "glu", {}), + ( + ("clamped_silu", "clamped_linear"), + "clamped_swiglu", + {"limit": 1.0, "alpha": 1.702, "glu_linear_offset": 0.5}, + ), +] +GATED_IDS = [desc for _, desc, _ in GATED_ACTIVATIONS] +get_direction_id = lambda is_bwd: "bwd" if is_bwd else "fwd" + + +def get_gated_cfg_key(act_desc, in_dtype, fp8_dtype, q_layout, is_bwd): + """Mirror of MXFP8GatedQuantConfig::to_key (gated_mxfp8_cutedsl.cuh).""" + # JAX quantize emits scales in the linear (non-swizzled) layout + flags = (True, q_layout.has_colwise, False, is_bwd) + major, minor = device_compute_capability() + return ( + f"cutedsl_mxfp8_gated_sm{major * 10 + minor}_{DTYPE_TO_STR[in_dtype]}_" + f"{FP8_TO_KEY[fp8_dtype]}_" + "_".join("1" if f else "0" for f in flags) + f"_{act_desc}" + ) + + +def run_gated_test_case( + act_type, act_desc, act_kwargs, is_bwd, shape, q_layout, in_dtype, fp8_dtype +): + """Assert the CuTeDSL and CUDA backends produce bit-identical gated activation outputs. + shape is (M, N): the gated input is (M, 2, N), [act | gate] along axis -2. + """ + M, N = shape + # generate_inputs returns (M, N) and (M, 1, N) draws; stack them into the (M, 2, N) input + act, gate = generate_inputs(M, N, in_dtype) + x = jnp.concatenate([jnp.expand_dims(act, axis=-2), gate], axis=-2) + grad, _ = generate_inputs(M, N, in_dtype, seed=1) + act_params = ActivationParams.create(act_type, **act_kwargs) + + def run(): + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, q_dtype=fp8_dtype, q_layout=q_layout + ) + if is_bwd: + out, _ = tex.quantize_dact_dbias( + grad, x, act_type, is_dbias=False, quantizer=quantizer, act_params=act_params + ) + else: + out = tex.act_lu(x, act_type, quantizer=quantizer, act_params=act_params) + return extract_quantized_output(out, None)[0] + + set_cutedsl_backend(False) + cuda_output = run() + set_cutedsl_backend(True) + try: + cutedsl_output = run() + finally: + set_cutedsl_backend(False) + + key = get_gated_cfg_key(act_desc, in_dtype, fp8_dtype, q_layout, is_bwd) + assert tvm_ffi.get_global_func(key, allow_missing=True) is not None, ( + f"CuTeDSL kernel not registered for {key}; the CuTeDSL backend fell back " + "to CUDA and this case compared CUDA against itself" + ) + + tag = ( + f"{act_desc}/{get_direction_id(is_bwd)}/{get_layout_id(q_layout)}/{M}x{N}/" + f"{DTYPE_TO_STR[in_dtype]}/{FP8_TO_STR[fp8_dtype]}" + ) + for name, cuda_bytes in cuda_output.items(): + assert np.array_equal( + cutedsl_output[name], cuda_bytes + ), f"{tag}: {name} differ between backends" + + +@pytest.mark.parametrize("shape", MATRIX_SIZES, ids=get_shape_id) +@pytest.mark.parametrize("q_layout", Q_LAYOUTS, ids=get_layout_id) +@pytest.mark.parametrize("act_type,act_desc,act_kwargs", GATED_ACTIVATIONS, ids=GATED_IDS) +@pytest.mark.parametrize("is_bwd", [False, True], ids=get_direction_id) +def test_gated_sizes(is_bwd, act_type, act_desc, act_kwargs, q_layout, shape): + run_gated_test_case( + act_type, + act_desc, + act_kwargs, + is_bwd, + shape, + q_layout, + jnp.bfloat16, + jnp.float8_e4m3fn, + ) + + +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id) +@pytest.mark.parametrize("act_type,act_desc,act_kwargs", GATED_ACTIVATIONS, ids=GATED_IDS) +@pytest.mark.parametrize("is_bwd", [False, True], ids=get_direction_id) +def test_gated_dtypes(is_bwd, act_type, act_desc, act_kwargs, fp8_dtype, in_dtype): + run_gated_test_case( + act_type, + act_desc, + act_kwargs, + is_bwd, + (256, 384), + QuantizeLayout.ROWWISE_COLWISE, + in_dtype, + fp8_dtype, + ) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index 5d7feeb6ce6..0c57f9c1be0 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -63,8 +63,8 @@ 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 -# they go to a separate TE/common kernel that the backend does not cover. +# config key (Activation in common/util/cutedsl_utils.h). Gated variants go to a separate kernel +# and are covered by GATED_ACTIVATIONS below. IDENTITY = {"name": "Identity", "act": None, "dact": None, "dbias_dact": None, "desc": "none"} ACTIVATIONS = [ { @@ -431,3 +431,127 @@ def test_specialized_kernels(shape, block_size, in_dtype, swizzled, input_offset swizzled, input_offset, ) + + +# Gated activations (gated_mxfp8.cuh). Input is [M, 2N] = [act | gate]; the forward writes +# [M, N] and the backward, given grad [M, N], writes [M, 2N]. "desc" is the name the activation +# goes by in the config key (GatedActivation in common/util/cutedsl_utils.h). The parameters of +# the parameterized ones are chosen so that their clamps and tanh terms are active on the +# [-2, 2] inputs generate_inputs draws. +GATED_ACTIVATIONS = [ + {"name": "GeGLU", "act": tex.geglu, "dact": tex.dgeglu, "desc": "geglu", "kwargs": {}}, + {"name": "QGeGLU", "act": tex.qgeglu, "dact": tex.dqgeglu, "desc": "qgeglu", "kwargs": {}}, + {"name": "ReGLU", "act": tex.reglu, "dact": tex.dreglu, "desc": "reglu", "kwargs": {}}, + {"name": "SReGLU", "act": tex.sreglu, "dact": tex.dsreglu, "desc": "sreglu", "kwargs": {}}, + {"name": "SwiGLU", "act": tex.swiglu, "dact": tex.dswiglu, "desc": "swiglu", "kwargs": {}}, + {"name": "GLU", "act": tex.glu, "dact": tex.dglu, "desc": "glu", "kwargs": {}}, + { + "name": "ClampedSwiGLU", + "act": tex.clamped_swiglu, + "dact": tex.clamped_dswiglu, + "desc": "clamped_swiglu", + "kwargs": {"limit": 1.0, "alpha": 1.702, "glu_linear_offset": 0.5}, + }, + { + "name": "SiTUGLU", + "act": tex.situglu, + "dact": tex.dsituglu, + "desc": "situglu", + "kwargs": {"beta1": 1.0, "beta2": 2.0}, + }, +] +get_gated_id = lambda a: a["name"] +GATED_DIRECTIONS = [False, True] +get_direction_id = lambda is_bwd: "bwd" if is_bwd else "fwd" + + +def run_on_both_backends(run): + """Return (CUDA result, CuTeDSL result) of run(), with run() responsible for materializing + whatever it returns before the backend is switched.""" + set_cutedsl_backend(False) + cuda_result = run() + set_cutedsl_backend(True) + try: + cutedsl_result = run() + finally: + set_cutedsl_backend(False) + return cuda_result, cutedsl_result + + +def assert_cutedsl_kernel_registered(key): + """Guard against a silent CUDA fallback, which would make a cross-backend comparison compare + CUDA against itself.""" + assert tvm_ffi.get_global_func(key, allow_missing=True) is not None, ( + f"CuTeDSL kernel not registered for {key}; the CuTeDSL backend fell back " + "to CUDA and this case compared CUDA against itself" + ) + + +def get_gated_cfg_key(act, in_dtype, fp8_dtype, rowwise, colwise, swizzled, is_bwd): + """Mirror of MXFP8GatedQuantConfig::to_key (gated_mxfp8_cutedsl.cuh).""" + major, minor = device_compute_capability() + flags = (rowwise, colwise, swizzled, is_bwd) + return ( + f"cutedsl_mxfp8_gated_sm{major * 10 + minor}_{DTYPE_TO_STR[in_dtype]}_" + f"{FP8_TO_KEY[fp8_dtype]}_" + "_".join("1" if f else "0" for f in flags) + f"_{act['desc']}" + ) + + +def run_gated_test_case(act, is_bwd, shape, block_size, in_dtype, fp8_dtype, swizzled=False): + """Assert the CuTeDSL and CUDA backends produce bit-identical gated activation outputs. + shape is (M, N): the gated input is [M, 2N]. + """ + M, N = shape + if is_bwd and swizzled and N % 128 != 0: + pytest.skip( + "CuTeDSL hands backward gated activations with GEMM-swizzled scales and " + "N % 128 != 0 back to CUDA" + ) + rowwise = block_size[1] != 1 + columnwise = block_size[0] != 1 + x, _ = generate_inputs(M, 2 * N, in_dtype) + grad, _ = generate_inputs(M, N, in_dtype, seed=1) + + def run(): + q = MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=rowwise, columnwise=columnwise) + q.optimize_for_gemm = swizzled + if is_bwd: + out = act["dact"](grad, x, q, **act["kwargs"]) + else: + out = act["act"](x, q, **act["kwargs"]) + return extract_quantized_output(out, rowwise, columnwise, swizzled) + + cuda_output, cutedsl_output = run_on_both_backends(run) + assert_cutedsl_kernel_registered( + get_gated_cfg_key(act, in_dtype, fp8_dtype, rowwise, columnwise, swizzled, is_bwd) + ) + + layout = "swizzled" if swizzled else "linear" + tag = ( + f"{act['name']}/{get_direction_id(is_bwd)}/{M}x{N}/{DTYPE_TO_STR[in_dtype]}/" + f"{FP8_TO_STR[fp8_dtype]}/{layout}" + ) + for name, cuda_bytes in cuda_output.items(): + assert torch.equal( + cutedsl_output[name], cuda_bytes + ), f"{tag}: {name} differ between backends" + + +@pytest.mark.parametrize("shape", MATRIX_SIZES, ids=get_shape_id) +@pytest.mark.parametrize("block_size", BLOCK_SIZES, ids=get_block_id) +@pytest.mark.parametrize("act", GATED_ACTIVATIONS, ids=get_gated_id) +@pytest.mark.parametrize("is_bwd", GATED_DIRECTIONS, ids=get_direction_id) +@pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) +def test_gated_sizes(swizzled, is_bwd, act, block_size, shape): + run_gated_test_case( + act, is_bwd, shape, block_size, torch.bfloat16, tex.DType.kFloat8E4M3, swizzled + ) + + +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id) +@pytest.mark.parametrize("act", GATED_ACTIVATIONS, ids=get_gated_id) +@pytest.mark.parametrize("is_bwd", GATED_DIRECTIONS, ids=get_direction_id) +@pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) +def test_gated_dtypes(swizzled, is_bwd, act, fp8_dtype, in_dtype): + run_gated_test_case(act, is_bwd, (256, 384), BOTH, in_dtype, fp8_dtype, swizzled) diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index e3920b13486..f9fabe72e5b 100644 --- a/transformer_engine/common/CuTeDSL/__init__.py +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -12,6 +12,7 @@ from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8_common import ( get_mxfp8_quantization_function, + get_mxfp8_gated_quantization_function, ) @@ -19,6 +20,8 @@ def register_cutedsl_backends(): """Register all available CuTeDSL backends for on-demand compilation via TVM-FFI. The C++ dispatcher retrieves them by the names defined here. """ - tvm_ffi.register_global_func( - "get_mxfp8_quantization_function", get_mxfp8_quantization_function, override=True - ) + for entrypoint in ( + get_mxfp8_quantization_function, + get_mxfp8_gated_quantization_function, + ): + tvm_ffi.register_global_func(entrypoint.__name__, entrypoint, override=True) diff --git a/transformer_engine/common/CuTeDSL/activations.py b/transformer_engine/common/CuTeDSL/activations.py index 34772eb7257..636dafe7f82 100644 --- a/transformer_engine/common/CuTeDSL/activations.py +++ b/transformer_engine/common/CuTeDSL/activations.py @@ -80,3 +80,57 @@ def dact_dgelu(x: Float32) -> Float32: return Float32(0.5) * x * ( (Float32(1.0) - t * t) * (Float32(0.79788456) + Float32(0.1070322243) * x * x) ) + Float32(0.5) * (Float32(1.0) + t) + + +# ---- Gated activations (gated_mxfp8.cuh) ---- +# The parameterized ones take their ClampedSwiGLUParam / SiTUGLUParam fields as +# trailing runtime arguments. + +LOG2E = 1.4426950408889634 # log2(e), what __expf scales its argument by + + +def act_sigmoid(x: Float32) -> Float32: + return sigmoid(x) + + +def dact_dsigmoid(x: Float32) -> Float32: + return dsigmoid(x) + + +def act_clamped_silu(x: Float32, limit: Float32, alpha: Float32) -> Float32: + c = cute.arch.fmin(limit, x) + return c * sigmoid(alpha * c) + + +def act_situ_gate(x: Float32, beta1: Float32) -> Float32: + return beta1 * tanh(x / beta1) * sigmoid(x) + + +def dact_dsitu_gate(x: Float32, beta1: Float32) -> Float32: + t = tanh(x / beta1) + s = sigmoid(x) + return (Float32(1.0) - t * t) * s + beta1 * t * s * (Float32(1.0) - s) + + +def situ_up(x: Float32, beta2: Float32) -> Float32: + return beta2 * tanh(x / beta2) + + +def dsitu_up(x: Float32, beta2: Float32) -> Float32: + t = tanh(x / beta2) + return Float32(1.0) - t * t + + +@cute.jit +def sigmoidf(x: Float32) -> Float32: + """math.h's sigmoidf: __frcp_rn(1.0f + __expf(-x)).""" + e = cute.arch.inline_ptx( + "ex2.approx.f32 {$w0}, {$r0};", + write_only_types=[Float32], + read_only_args=[-x * Float32(LOG2E)], + ) + return cute.arch.inline_ptx( + "rcp.rn.f32 {$w0}, {$r0};", + write_only_types=[Float32], + read_only_args=[Float32(1.0) + e], + ) 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 f003a20df30..85ece5fb0ce 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -212,6 +212,134 @@ def is_plain_cast(self) -> bool: """Whether cfg is a plain cast: no amax, dbias or (d)activation fused in.""" return not (self.WITH_AMAX or self.WITH_DBIAS or self.WITH_DACT or self.WITH_ACT) + def compile(self): + """ + Return the compiled CuTeDSL function object for the given MXFP8 quantization config. + """ + + kernel_obj = MXFP8QuantizeEntry(self) + 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 = self.FP8_DTYPE + scale_dtype = cutlass.Float8E8M0FNU + + in_fake = cute.runtime.make_fake_compact_tensor( + self.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 self.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 self.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 self.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 self.COLWISE + else None + ) + amax_fake = ( + cute.runtime.make_fake_compact_tensor( + Float32, (1,), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4 + ) + if self.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( + self.DTYPE, + in_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if self.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 self.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 (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 + # as void*) + options="--enable-tvm-ffi", + ) + return compiled + class MXFP8QuantizeKernelBase(abc.ABC): """Base class for MXFP8 quantize kernels.""" @@ -404,142 +532,8 @@ def __call__( ) -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 (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 - # 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 +def compile_and_register_function(fn_name: str, cfg: MXFP8QuantizeConfig) -> bool: + """Compile the CuTeDSL 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. @@ -559,33 +553,8 @@ def get_mxfp8_quantization_function( ) 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) + compiled = cfg.compile() # 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 @@ -607,3 +576,82 @@ def get_mxfp8_quantization_function( ) # Unconditionally fallback to CUDA path because we can't tell if this exception is transient or permanent. return False + + +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 and register the MXFP8 quantize kernel for this config under `fn_name`.""" + 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 + return compile_and_register_function(fn_name, cfg) + + +def get_mxfp8_gated_quantization_function( + fn_name: str, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + is_bwd: bool, + activation: str, +) -> bool: + """Compile and register the gated MXFP8 quantize kernel for this config under `fn_name`.""" + # pylint: disable=import-outside-toplevel + from .quantize_mxfp8_gated_g2s import MXFP8GatedQuantizeConfig + + try: + cfg = MXFP8GatedQuantizeConfig( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + is_bwd=is_bwd, + 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 + return compile_and_register_function(fn_name, cfg) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py new file mode 100644 index 00000000000..2baa03dd3bf --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py @@ -0,0 +1,861 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Gated-activation MXFP8 casts in CuTeDSL, ported from gated_mxfp8.cuh. + +The kernel reuses MXFP8QuantizeGeneralG2SKernel's TMA pipeline and its plain-cast SMEM tile +quantizer. What it adds is an elementwise pass that runs on each staged tile before it is +quantized: it combines the act and gate halves of the input (and the incoming gradient, in +backward) and writes the result back into the act (and gate) input buffer in the input dtype. +That is exactly the truncation the CUDA kernel applies before computing block amaxes, so the +plain cast that follows reproduces its output. + +Gated input is [M, 2N] with [act | gate] along the last dim. Every operand is addressed through +(M, N) views (stride 2N, column offset 0 or N), mirroring the CUDA kernel's tensor maps: + forward : out[M, N] = act(x[:, :N]) * gate(x[:, N:]) + backward: out[M, 2N] = [d/d act | d/d gate] given grad[M, N] +""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import pipeline +from cutlass import Boolean, Float32, Int32, Int64 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + fabs_f32, + is_packed16, + truncate_f32_bf16, + truncate_f32_f16, +) +from transformer_engine.common.CuTeDSL.activations import ( + act_gelu, + act_qgelu, + act_relu, + act_srelu, + act_silu, + act_sigmoid, + act_clamped_silu, + act_situ_gate, + dact_dgelu, + dact_dqgelu, + dact_drelu, + dact_dsrelu, + dact_dsigmoid, + dact_dsitu_gate, + situ_up, + dsitu_up, + sigmoidf, +) +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + MXFP8QuantizeConfig, + MXFP8_BLOCK_SCALING_SIZE, + derive_swizzled_scale_layout, +) +from .quantize_mxfp8_g2s_general import ( + MXFP8QuantizeGeneralG2SKernel, + quantize_mxfp8_smem_tile_general, +) + +# Gated activations of gated_mxfp8.cuh, as (ActOP, DActOP) ports of the math.h functions the +# CUDA dispatcher instantiates them with. swiglu and clamped_swiglu backward, and situglu in +# both directions, do not go through this table: like the CUDA kernel they take a special path. +GATED_ACTIVATIONS = { + "geglu": (act_gelu, dact_dgelu), + "qgeglu": (act_qgelu, dact_dqgelu), + "reglu": (act_relu, dact_drelu), + "sreglu": (act_srelu, dact_dsrelu), + "glu": (act_sigmoid, dact_dsigmoid), + "swiglu": (act_silu, None), + "clamped_swiglu": (None, None), + "situglu": (None, None), +} + + +@cute.jit +def _select_f32(pred: Boolean, if_true: Float32, if_false: Float32) -> Float32: + result = if_false + if pred: + result = if_true + return result + + +@cute.jit +def _gated_element( + act_elt: Float32, + gate_elt: Float32, + grad_elt: Optional[Float32], + limit: Float32, + alpha: Float32, + glu_linear_offset: Float32, + beta1: Float32, + beta2: Float32, + ACTIVATION: cutlass.Constexpr[str], + IS_BWD: cutlass.Constexpr[bool], +): + """One element of the gated op in FP32; returns (act output, gate output or None). + + Mirrors the elementwise math of gated_mxfp8.cuh's columnwise pass (forward and backward), + operation by operation. + """ + out_gate = None + # Transform the gate. Gating is ideally the identity, with derivative 1. + gate = gate_elt + dgate = None + if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): + # |g| > limit is gated_mxfp8.cuh's (g > limit || g < -limit), NaN included + dgate = _select_f32(fabs_f32(gate_elt) > limit, Float32(0.0), Float32(1.0)) + gate = cute.arch.fmin(cute.arch.fmax(-limit, gate_elt), limit) + glu_linear_offset + elif cutlass.const_expr(ACTIVATION == "situglu"): + dgate = dsitu_up(gate_elt, beta2) + gate = situ_up(gate_elt, beta2) + + if cutlass.const_expr(not IS_BWD): + if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): + act = act_clamped_silu(act_elt, limit, alpha) + elif cutlass.const_expr(ACTIVATION == "situglu"): + act = act_situ_gate(act_elt, beta1) + else: + act = GATED_ACTIVATIONS[ACTIVATION][0](act_elt) + out_act = act * gate + else: + if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): + x = cute.arch.fmin(act_elt, limit) + s = sigmoidf(alpha * x) + act_x = x * s + dact_x = _select_f32( + act_elt <= limit, s + s * (Float32(1.0) - s) * alpha * x, Float32(0.0) + ) + elif cutlass.const_expr(ACTIVATION == "situglu"): + act_x = act_situ_gate(act_elt, beta1) + dact_x = dact_dsitu_gate(act_elt, beta1) + elif cutlass.const_expr(ACTIVATION == "swiglu"): + s = sigmoidf(act_elt) + act_x = act_elt * s + dact_x = act_elt * s * (Float32(1.0) - s) + s + else: + act_fn, dact_fn = GATED_ACTIVATIONS[ACTIVATION] + act_x = act_fn(act_elt) + dact_x = dact_fn(act_elt) + out_act = dact_x * grad_elt * gate + if cutlass.const_expr(dgate is None): + out_gate = act_x * grad_elt + else: + out_gate = dgate * act_x * grad_elt + return out_act, out_gate + + +@cute.jit +def gated_elementwise_smem_tile( + sAct_tile: cute.Tensor, # (TILE_Y, TILE_X) act half of the input; receives the act output + sGate_tile: cute.Tensor, # (TILE_Y, TILE_X) gate half; receives the gate output (bwd only) + sGrad_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) incoming gradient (bwd only) + limit: Float32, + alpha: Float32, + glu_linear_offset: Float32, + beta1: Float32, + beta2: Float32, + ACTIVATION: cutlass.Constexpr[str], + IS_BWD: cutlass.Constexpr[bool], + DTYPE: cutlass.Constexpr[type], + TILE_X: cutlass.Constexpr[int], +): + """Compute the gated op on one staged tile, in place, rounded to the input dtype. + + Thread tidx owns column tidx of the tile, which is also the element set it owns in the + colwise quantization pass, so a colwise-only consumer needs no barrier after this. + """ + tidx, _, _ = cute.arch.thread_idx() + + _, 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)), + ) + sAct_thread = cute.composition(sAct_tile, tv_layout)[tidx, None] + sGate_thread = cute.composition(sGate_tile, tv_layout)[tidx, None] + if cutlass.const_expr(IS_BWD): + sGrad_thread = cute.composition(sGrad_tile, tv_layout)[tidx, None] + if cutlass.const_expr(is_packed16(DTYPE)): + truncate_f32 = truncate_f32_f16 if DTYPE is cutlass.Float16 else truncate_f32_bf16 + + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + grad_elt = None + if cutlass.const_expr(IS_BWD): + grad_elt = Float32(sGrad_thread[i]) + out_act, out_gate = _gated_element( + Float32(sAct_thread[i]), + Float32(sGate_thread[i]), + grad_elt, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=ACTIVATION, + IS_BWD=IS_BWD, + ) + # Numerical truncation: round to the input dtype, like the CUDA kernels do before + # taking the amax. The cast below is then exact. + if cutlass.const_expr(is_packed16(DTYPE)): + out_act = truncate_f32(out_act) + if cutlass.const_expr(IS_BWD): + out_gate = truncate_f32(out_gate) + sAct_thread[i] = DTYPE(out_act) + if cutlass.const_expr(IS_BWD): + sGate_thread[i] = DTYPE(out_gate) + + +@cute.jit +def quantize_gated_mxfp8_smem_tile( + sAct_tile: cute.Tensor, + sGate_tile: cute.Tensor, + sGrad_tile: Optional[cute.Tensor], + sO_row_act_tile: Optional[cute.Tensor], + mS_row_act_stage: Optional[cute.Tensor], + sO_row_gate_tile: Optional[cute.Tensor], # bwd only + mS_row_gate_stage: Optional[cute.Tensor], # bwd only + sO_col_act_tile: Optional[cute.Tensor], + mS_col_act_stage: Optional[cute.Tensor], + sO_col_gate_tile: Optional[cute.Tensor], # bwd only + mS_col_gate_stage: Optional[cute.Tensor], # bwd only + tile_row_start: Int32, + tile_col_start: Int32, + M: Int32, + N: Int32, + limit: Float32, + alpha: Float32, + glu_linear_offset: Float32, + beta1: Float32, + beta2: Float32, + ACTIVATION: cutlass.Constexpr[str], + IS_BWD: cutlass.Constexpr[bool], + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + MAX_NORM_RCP: cutlass.Constexpr[float], + DTYPE: cutlass.Constexpr[type], + FP8_DTYPE: cutlass.Constexpr[type], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], + TILE_X: cutlass.Constexpr[int], + TILE_Y: cutlass.Constexpr[int], +): + """Apply the gated op to one staged tile, then MXFP8-quantize the result(s) as a plain cast.""" + gated_elementwise_smem_tile( + sAct_tile, + sGate_tile, + sGrad_tile, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=ACTIVATION, + IS_BWD=IS_BWD, + DTYPE=DTYPE, + TILE_X=TILE_X, + ) + # The rowwise pass reads elements other threads wrote. + if cutlass.const_expr(ROWWISE): + cute.arch.sync_threads() + + # The act output and, backward, the gate output. Both share this tile's coordinates; only + # their output buffers and scale views differ. + halves = [(sAct_tile, sO_row_act_tile, mS_row_act_stage, sO_col_act_tile, mS_col_act_stage)] + if cutlass.const_expr(IS_BWD): + halves.append( + (sGate_tile, sO_row_gate_tile, mS_row_gate_stage, sO_col_gate_tile, mS_col_gate_stage) + ) + for sX_tile, sO_row_tile, mS_row_stage, sO_col_tile, mS_col_stage in halves: + quantize_mxfp8_smem_tile_general( + sX_tile, + None, + sO_row_tile, + mS_row_stage, + sO_col_tile, + mS_col_stage, + None, + tile_row_start, + tile_col_start, + M, + N, + ROWWISE=ROWWISE, + COLWISE=COLWISE, + WITH_AMAX=False, + WITH_ACT=False, + WITH_DACT=False, + ACTIVATION=None, + MAX_NORM_RCP=MAX_NORM_RCP, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=WITH_GEMM_SWIZZLED_SCALES, + TILE_X=TILE_X, + TILE_Y=TILE_Y, + # The gated op maps the zeros TMA fills out-of-bounds elements with to zero. + SKIP_INPUT_MASKING=True, + SKIP_SCALE_BOUNDS=False, + CACHE_ACTIVATION=False, + DBIAS_REDUCTION_IN_COLWISE=False, + COLWISE_DBIAS_REDUCTION_ONLY=False, + DBIAS_REDUCTION_IN_ROWWISE=False, + ) + + +@cute.jit +def _column_slice(mT: cute.Tensor, col_offset: Int32, num_cols: Int32) -> cute.Tensor: + """(rows, num_cols) view of the row-major mT starting at column col_offset. + + The pointer is rebuilt rather than offset so that TMA sees it 16-byte aligned, which the C++ + dispatcher guarantees for both the tensor and the column offset of every slice. + """ + ptr = cute.make_ptr( + mT.element_type, + mT.iterator.toint() + Int64(col_offset) * Int64(mT.element_type.width // 8), + cute.AddressSpace.gmem, + assumed_align=16, + ) + return cute.make_tensor( + ptr, cute.make_layout((mT.shape[0], num_cols), stride=(mT.stride[0], 1)) + ) + + +def _fake_gmem_tensor(dtype, shape, align): + """Row-major fake gmem tensor for cute.compile.""" + return cute.runtime.make_fake_compact_tensor( + dtype, + shape, + stride_order=tuple(reversed(range(len(shape)))), + memspace=cute.AddressSpace.gmem, + assumed_align=align, + ) + + +class MXFP8GatedQuantizeConfig(MXFP8QuantizeConfig): + """Compile-time config of gated_mxfp8.cuh's port: a plain MXFP8 cast config plus the gated op.""" + + def __init__( + self, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + is_bwd: bool, + activation: str, + ): + super().__init__( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + with_amax=False, + activation="none", + ) + if activation not in GATED_ACTIVATIONS: + raise ValueError( + f"unknown gated activation {activation!r}; expected one of" + f" {sorted(GATED_ACTIVATIONS)}" + ) + self.IS_BWD = is_bwd + self.GATED_ACTIVATION = activation + + def __str__(self): + return ( + f"MXFP8GatedQuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " + f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, " + f"swizzled={self.WITH_GEMM_SWIZZLED_SCALES}, is_bwd={self.IS_BWD}, " + f"activation={self.GATED_ACTIVATION})" + ) + + __repr__ = __str__ + + def compile(self): + """Return the compiled CuTeDSL function object for this config.""" + sym_M = cute.sym_int32() + # The C++ dispatcher requires N % 32 == 0, like the CUDA kernel. + sym_N = cute.sym_int32(divisibility=32) + sym_2N = cute.sym_int32(divisibility=64) + out_cols = sym_2N if self.IS_BWD else sym_N + # Padded scale shapes, see MXFP8QuantizeConfig.compile + 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)) + params = [Float32(0.0)] * 5 # limit, alpha, glu_linear_offset, beta1, beta2 + return cute.compile( + MXFP8GatedQuantizeG2SKernel(self), + _fake_gmem_tensor(self.DTYPE, (sym_M, sym_2N), 16), # mGatedInput + _fake_gmem_tensor(self.DTYPE, (sym_M, sym_N), 16) if self.IS_BWD else None, # mGrad + # mO_row, mS_row + _fake_gmem_tensor(self.FP8_DTYPE, (sym_M, out_cols), 16) if self.ROWWISE else None, + ( + _fake_gmem_tensor(cutlass.Float8E8M0FNU, scale_rowwise_shape, 4) + if self.ROWWISE + else None + ), + # mO_col, mS_col + _fake_gmem_tensor(self.FP8_DTYPE, (sym_M, out_cols), 16) if self.COLWISE else None, + ( + _fake_gmem_tensor(cutlass.Float8E8M0FNU, scale_colwise_shape, 4) + if self.COLWISE + else None + ), + *params, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + + +class MXFP8GatedQuantizeG2SKernel(MXFP8QuantizeGeneralG2SKernel): + """gated_mxfp8.cuh in CuTeDSL. The general kernel's pipeline, loading the act and gate halves + (and grad, backward) each stage and storing one (forward) or two (backward) quantized halves. + + Inherits the plain-cast tile config, which is gated_mxfp8.cuh's: 64 threads per CTA, each CTA + covering two vertically stacked 32x64 tiles. + """ + + def __init__(self, cfg: MXFP8GatedQuantizeConfig): + super().__init__(cfg) + # Same occupancy hint as the general kernel with a fused activation + self._MIN_BLOCKS_PER_MP = 1 + + @cute.jit + def __call__( + self, + mGatedInput: cute.Tensor, # [M, 2N] = [act | gate] + mGrad: Optional[cute.Tensor], # [M, N], bwd only + mO_row: Optional[cute.Tensor], # [M, N] (fwd) or [M, 2N] (bwd) + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + limit: Float32, # ClampedSwiGLUParam + alpha: Float32, + glu_linear_offset: Float32, + beta1: Float32, # SiTUGLUParam + beta2: Float32, + stream: CUstream, + ): + cfg = self.cfg + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf(f"[CuTeDSL] MXFP8GatedQuantizeG2SKernel.__call__() with config: {cfg}\n") + + mO = mO_row if cfg.ROWWISE else mO_col + M = mGatedInput.shape[0] + N = mGrad.shape[1] if cfg.IS_BWD else mO.shape[1] + + # Scales cover the whole output: [M, 2N] backward, with the gate half's scales starting + # at scale column N / 32 (rowwise) or N (colwise). + if cutlass.const_expr(cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row, mS_col = derive_swizzled_scale_layout( + M, mO.shape[1], cfg.ROWWISE, cfg.COLWISE, mS_row, mS_col + ) + mS_row_gate = None + mS_col_gate = None + mO_row_gate = None + mO_col_gate = None + if cutlass.const_expr(cfg.IS_BWD): + # Swizzled scales are only translation invariant across whole 128-column tiles, + # which the C++ dispatcher guarantees by requiring N % 128 == 0 for them. + if cutlass.const_expr(cfg.ROWWISE): + mS_row_gate = cute.domain_offset((0, N // MXFP8_BLOCK_SCALING_SIZE), mS_row) + mO_row_gate = _column_slice(mO_row, N, N) + mO_row = _column_slice(mO_row, 0, N) + if cutlass.const_expr(cfg.COLWISE): + mS_col_gate = cute.domain_offset((0, N), mS_col) + mO_col_gate = _column_slice(mO_col, N, N) + mO_col = _column_slice(mO_col, 0, N) + + atoms = self._make_tma_atoms( + mGrad, + _column_slice(mGatedInput, 0, N), + _column_slice(mGatedInput, N, N), + mO_row, + mO_row_gate, + mO_col, + mO_col_gate, + ) + # The TMA atoms are spread from a list, which pylint cannot count + self.kernel( # pylint: disable=no-value-for-parameter + M, + N, + mS_row, + mS_row_gate, + mS_col, + mS_col_gate, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + *atoms, + ).launch( + grid=self._grid(M, N), + block=[self._THREADS_PER_CTA], + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, + ) + + def _grid(self, M: Int32, N: Int32): + return [ + cute.ceil_div(N, self._TILE_COLS), + cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES), + ] + + @cute.jit + def _make_tma_atoms( + self, + mGrad: Optional[cute.Tensor], + mAct: cute.Tensor, + mGate: cute.Tensor, + mO_row_act: Optional[cute.Tensor], + mO_row_gate: Optional[cute.Tensor], + mO_col_act: Optional[cute.Tensor], + mO_col_gate: Optional[cute.Tensor], + ): + """TMA atoms and their gmem views for every (TILE_ROWS, TILE_COLS) operand; None if absent.""" + smem_tile_layout = cute.make_ordered_layout( + (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) + ) + cta_tiler = (self._TILE_ROWS, self._TILE_COLS) + op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + + atoms = [] + for op, mT in ( + (op_load, mGrad), + (op_load, mAct), + (op_load, mGate), + (op_store, mO_row_act), + (op_store, mO_row_gate), + (op_store, mO_col_act), + (op_store, mO_col_gate), + ): + if cutlass.const_expr(mT is None): + atoms.extend([None, None]) + else: + atoms.extend( + cute.nvgpu.cpasync.make_tiled_tma_atom( + op, mT, smem_tile_layout, cta_tiler, num_multicast=1 + ) + ) + return atoms + + @cute.jit + def _copy_inputs_to_stage( + self, + tile_y: Int32, + tile_x: Int32, + stage: Int32, + tma_bar_ptr: cute.Pointer, + tma_atom_grad: Optional[cute.CopyAtom], + tXgGrad: Optional[cute.Tensor], + tXsGrad: Optional[cute.Tensor], + tma_atom_act: cute.CopyAtom, + tXgAct: cute.Tensor, + tXsAct: cute.Tensor, + tma_atom_gate: cute.CopyAtom, + tXgGate: cute.Tensor, + tXsGate: cute.Tensor, + ): + """Issue this tile's input TMA loads into the pipeline stage, all on one barrier.""" + if cutlass.const_expr(self.cfg.IS_BWD): + cute.copy( + tma_atom_grad, + tXgGrad[(None, (tile_y, tile_x))], + tXsGrad[(None, stage)], + tma_bar_ptr=tma_bar_ptr, + ) + cute.copy( + tma_atom_act, + tXgAct[(None, (tile_y, tile_x))], + tXsAct[(None, stage)], + tma_bar_ptr=tma_bar_ptr, + ) + cute.copy( + tma_atom_gate, + tXgGate[(None, (tile_y, tile_x))], + tXsGate[(None, stage)], + tma_bar_ptr=tma_bar_ptr, + ) + + @cute.kernel + def kernel( + self, + M: Int32, + N: Int32, + mS_row_act: Optional[cute.Tensor], + mS_row_gate: Optional[cute.Tensor], + mS_col_act: Optional[cute.Tensor], + mS_col_gate: Optional[cute.Tensor], + limit: Float32, + alpha: Float32, + glu_linear_offset: Float32, + beta1: Float32, + beta2: Float32, + tma_atom_grad: Optional[cute.CopyAtom], + tma_src_grad: Optional[cute.Tensor], + tma_atom_act: cute.CopyAtom, + tma_src_act: cute.Tensor, + tma_atom_gate: cute.CopyAtom, + tma_src_gate: cute.Tensor, + tma_atom_row_act: Optional[cute.CopyAtom], + tma_dst_row_act: Optional[cute.Tensor], + tma_atom_row_gate: Optional[cute.CopyAtom], + tma_dst_row_gate: Optional[cute.Tensor], + tma_atom_col_act: Optional[cute.CopyAtom], + tma_dst_col_act: Optional[cute.Tensor], + tma_atom_col_gate: Optional[cute.CopyAtom], + tma_dst_col_gate: Optional[cute.Tensor], + ): + cfg = self.cfg + DTYPE = cfg.DTYPE + TILE_ROWS = self._TILE_ROWS + TILE_COLS = self._TILE_COLS + NUM_STAGES = self._NUM_STAGES + NUM_TILES = self._NUM_TILES + + # Scale views tiled like the general kernel's + row_scale_tile = (TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + col_scale_tile = (TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, TILE_COLS) + if cutlass.const_expr(cfg.ROWWISE): + mS_row_act = cute.zipped_divide(mS_row_act, row_scale_tile) + if cutlass.const_expr(cfg.IS_BWD): + mS_row_gate = cute.zipped_divide(mS_row_gate, row_scale_tile) + if cutlass.const_expr(cfg.COLWISE): + mS_col_act = cute.zipped_divide(mS_col_act, col_scale_tile) + if cutlass.const_expr(cfg.IS_BWD): + mS_col_gate = cute.zipped_divide(mS_col_gate, col_scale_tile) + + # Shared memory: one ((TILE_ROWS, TILE_COLS), NUM_STAGES) ring buffer per operand + smem = cutlass.memory.SmemAllocator() + + @cute.struct + class BarrierStorage: # pylint: disable=missing-class-docstring + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * NUM_STAGES] + + storage = smem.allocate(BarrierStorage) + stage_layout = cute.make_layout( + ((TILE_ROWS, TILE_COLS), NUM_STAGES), + stride=((TILE_COLS, 1), TILE_ROWS * TILE_COLS), + ) + sAct = smem.allocate_tensor(DTYPE, stage_layout, byte_alignment=128) + sGate = smem.allocate_tensor(DTYPE, stage_layout, byte_alignment=128) + sGrad = None + if cutlass.const_expr(cfg.IS_BWD): + sGrad = smem.allocate_tensor(DTYPE, stage_layout, byte_alignment=128) + sO_row_act = None + sO_row_gate = None + sO_col_act = None + sO_col_gate = None + if cutlass.const_expr(cfg.ROWWISE): + sO_row_act = smem.allocate_tensor(cfg.FP8_DTYPE, stage_layout, byte_alignment=128) + if cutlass.const_expr(cfg.IS_BWD): + sO_row_gate = smem.allocate_tensor(cfg.FP8_DTYPE, stage_layout, byte_alignment=128) + if cutlass.const_expr(cfg.COLWISE): + sO_col_act = smem.allocate_tensor(cfg.FP8_DTYPE, stage_layout, byte_alignment=128) + if cutlass.const_expr(cfg.IS_BWD): + sO_col_gate = smem.allocate_tensor(cfg.FP8_DTYPE, stage_layout, byte_alignment=128) + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + bidx, bidy, _ = cute.arch.block_idx() + + if warp_idx == 0: + if cutlass.const_expr(cfg.IS_BWD): + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_grad) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_act) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_gate) + + # Warp 0's elected thread produces (issues TMA), every warp consumes + producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._NUM_WARPS) + num_inputs = 3 if cfg.IS_BWD else 2 + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar_storage.data_ptr(), + num_stages=NUM_STAGES, + producer_group=producer_group, + consumer_group=consumer_group, + # Every input tile of a stage arrives on the same barrier + tx_count=num_inputs * TILE_ROWS * TILE_COLS * DTYPE.width // 8, + cta_layout_vmnk=None, + ) + prod_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, NUM_STAGES) + cons_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, NUM_STAGES) + + # Partition every operand for TMA (single CTA, no multicast) + cta_tiler = (TILE_ROWS, TILE_COLS) + tXsGrad = None + tXgGrad = None + if cutlass.const_expr(cfg.IS_BWD): + tXsGrad, tXgGrad = cute.nvgpu.cpasync.tma_partition( + tma_atom_grad, + 0, + cute.make_layout(1), + sGrad, + cute.zipped_divide(tma_src_grad, cta_tiler), + ) + tXsAct, tXgAct = cute.nvgpu.cpasync.tma_partition( + tma_atom_act, 0, cute.make_layout(1), sAct, cute.zipped_divide(tma_src_act, cta_tiler) + ) + tXsGate, tXgGate = cute.nvgpu.cpasync.tma_partition( + tma_atom_gate, + 0, + cute.make_layout(1), + sGate, + cute.zipped_divide(tma_src_gate, cta_tiler), + ) + # Output (atom, smem partition, gmem partition) triples, act half first + outputs = [] + for atom_out, dst_out, sO_out in ( + (tma_atom_row_act, tma_dst_row_act, sO_row_act), + (tma_atom_row_gate, tma_dst_row_gate, sO_row_gate), + (tma_atom_col_act, tma_dst_col_act, sO_col_act), + (tma_atom_col_gate, tma_dst_col_gate, sO_col_gate), + ): + if cutlass.const_expr(atom_out is not None): + tXsO_out, tXgO_out = cute.nvgpu.cpasync.tma_partition( + atom_out, + 0, + cute.make_layout(1), + sO_out, + cute.zipped_divide(dst_out, cta_tiler), + ) + outputs.append((atom_out, tXsO_out, tXgO_out)) + + # Barrier init must be visible to every thread + cute.arch.sync_threads() + + # Prologue: fill the pipeline + if warp_idx == 0: + for s in cutlass.range_constexpr(min(NUM_STAGES, NUM_TILES)): + mainloop_pipeline.producer_acquire(prod_state) + self._copy_inputs_to_stage( + bidy * NUM_TILES + s, + bidx, + prod_state.index, + mainloop_pipeline.producer_get_barrier(prod_state), + tma_atom_grad, + tXgGrad, + tXsGrad, + tma_atom_act, + tXgAct, + tXsAct, + tma_atom_gate, + tXgGate, + tXsGate, + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + for tile_idx in cutlass.range(NUM_TILES, unroll=1): + mainloop_pipeline.consumer_wait(cons_state) + # This tile reuses the output buffers of the tile NUM_STAGES ago; wait until TMA has + # read them. + if cutlass.const_expr(NUM_TILES > NUM_STAGES): + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(NUM_STAGES - 1, read=True) + cute.arch.sync_threads() + stage_idx = cons_state.index + tile_idx_y = bidy * NUM_TILES + tile_idx + scale_coord = (None, (tile_idx_y, bidx)) + + sGrad_tile = None + if cutlass.const_expr(cfg.IS_BWD): + sGrad_tile = sGrad[(None, stage_idx)] + sO_row_act_tile = None + mS_row_act_stage = None + sO_row_gate_tile = None + mS_row_gate_stage = None + if cutlass.const_expr(cfg.ROWWISE): + sO_row_act_tile = sO_row_act[(None, stage_idx)] + mS_row_act_stage = cute.flatten(mS_row_act[scale_coord]) + if cutlass.const_expr(cfg.IS_BWD): + sO_row_gate_tile = sO_row_gate[(None, stage_idx)] + mS_row_gate_stage = cute.flatten(mS_row_gate[scale_coord]) + sO_col_act_tile = None + mS_col_act_stage = None + sO_col_gate_tile = None + mS_col_gate_stage = None + if cutlass.const_expr(cfg.COLWISE): + sO_col_act_tile = sO_col_act[(None, stage_idx)] + mS_col_act_stage = cute.flatten(mS_col_act[scale_coord]) + if cutlass.const_expr(cfg.IS_BWD): + sO_col_gate_tile = sO_col_gate[(None, stage_idx)] + mS_col_gate_stage = cute.flatten(mS_col_gate[scale_coord]) + + quantize_gated_mxfp8_smem_tile( + sAct[(None, stage_idx)], + sGate[(None, stage_idx)], + sGrad_tile, + sO_row_act_tile, + mS_row_act_stage, + sO_row_gate_tile, + mS_row_gate_stage, + sO_col_act_tile, + mS_col_act_stage, + sO_col_gate_tile, + mS_col_gate_stage, + tile_idx_y * TILE_ROWS, + bidx * TILE_COLS, + M, + N, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=cfg.GATED_ACTIVATION, + IS_BWD=cfg.IS_BWD, + ROWWISE=cfg.ROWWISE, + COLWISE=cfg.COLWISE, + MAX_NORM_RCP=cfg.MAX_NORM_RCP, + DTYPE=DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=cfg.WITH_GEMM_SWIZZLED_SCALES, + TILE_X=TILE_COLS, + TILE_Y=TILE_ROWS, + ) + + # Make the SMEM writes visible to the TMA's async proxy before it reads them + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + mainloop_pipeline.consumer_release(cons_state) + + if warp_idx == 0: + # Output tensors are only read here, inside this dynamic branch, so these loop + # targets are fresh names rather than rebindings of values defined outside it. + for store_atom, store_sO, store_gO in outputs: + cute.copy( + store_atom, + store_sO[(None, stage_idx)], + store_gO[(None, (tile_idx_y, bidx))], + ) + cute.arch.cp_async_bulk_commit_group() + + cons_state.advance() + + # Refill the stage just released with the tile NUM_STAGES ahead + if warp_idx == 0: + next_tile_idx = tile_idx + NUM_STAGES + if next_tile_idx < NUM_TILES: + mainloop_pipeline.producer_acquire(prod_state) + self._copy_inputs_to_stage( + bidy * NUM_TILES + next_tile_idx, + bidx, + prod_state.index, + mainloop_pipeline.producer_get_barrier(prod_state), + tma_atom_grad, + tXgGrad, + tXsGrad, + tma_atom_act, + tXgAct, + tXsAct, + tma_atom_gate, + tXgGate, + tXsGate, + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + # Only wait until the in-flight TMA stores have read SMEM, so the CTA retires early + cute.arch.cp_async_bulk_wait_group(0, read=True) diff --git a/transformer_engine/common/cast/dispatch/gated.cuh b/transformer_engine/common/cast/dispatch/gated.cuh index 06e8f0e306c..aa4a5f56636 100644 --- a/transformer_engine/common/cast/dispatch/gated.cuh +++ b/transformer_engine/common/cast/dispatch/gated.cuh @@ -19,6 +19,10 @@ #include "../fp8/gated_fp8.cuh" #include "../mxfp8/gated_mxfp8.cuh" +#ifdef NVTE_WITH_CUTEDSL +#include "../mxfp8/gated_mxfp8_cutedsl.cuh" +#endif + namespace transformer_engine { namespace dispatch { @@ -86,8 +90,16 @@ void quantize_gated_fwd_helper(const NVTETensor nvte_input, NVTETensor nvte_outp NVTE_CHECK(is_supported_by_CC_100(), "Gated FWD NVTE_MXFP8_1D_SCALING is only supported on SM 10.0+"); Tensor dummy_grad_tensor; - mxfp8::quantize_gated(input, dummy_grad_tensor, - output, p, stream); + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_gated_cutedsl( + input, dummy_grad_tensor, output, p, stream); +#endif + if (!quantized_with_cutedsl) { + mxfp8::quantize_gated(input, dummy_grad_tensor, + output, p, stream); + } break; } default: @@ -176,8 +188,16 @@ void quantize_gated_bwd_helper(const NVTETensor nvte_grad, const NVTETensor nvte NVTE_CHECK(is_supported_by_CC_100(), "Gated BWD NVTE_MXFP8_1D_SCALING is only supported on SM 10.0+"); - mxfp8::quantize_gated(gated_input, grad, output, p, - stream); + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_gated_cutedsl( + gated_input, grad, output, p, stream); +#endif + if (!quantized_with_cutedsl) { + mxfp8::quantize_gated(gated_input, grad, output, p, + stream); + } break; } default: diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh new file mode 100644 index 00000000000..eab31ed01fd --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh @@ -0,0 +1,205 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file gated_mxfp8_cutedsl.cuh + * \brief CuTeDSL backend of gated_mxfp8.cuh. + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GATED_MXFP8_CUTEDSL_CUH_ +#define TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GATED_MXFP8_CUTEDSL_CUH_ + +#include +#include + +#include +#include +#include +#include +#include + +#include "../../common.h" +#include "../../tvm_ffi_bridge.h" +#include "../../util/cutedsl_utils.h" + +namespace transformer_engine { +namespace cutedsl_backend { + +using namespace tvm_ffi_bridge; + +namespace detail { + +// TMA needs 16-byte aligned global addresses. +inline bool is_16B_aligned(const void *ptr) { return reinterpret_cast(ptr) % 16 == 0; } + +} // namespace detail + +struct MXFP8GatedQuantConfig { + static constexpr const char *kEntrypointName = "get_mxfp8_gated_quantization_function"; + + DType dtype; // The input format + DType fp8_dtype; // The fp8 output format + bool rowwise; // If quantize rowwisely + bool colwise; // If quantize columnwisely + bool swizzled; // If the scale output is used for cudnn's swizzled layout + bool is_bwd; // Backward (dgated) rather than forward + GatedActivation activation; + uint32_t sm_arch = static_cast(cuda::sm_arch()); + + // Bit layout: dtype [3:0], fp8_dtype [7:4], flags [11:8], activation [17:12] and + // SM architecture [26:18]. + uint32_t to_id() const { + static_assert(static_cast(DType::kNumTypes) <= 16, + "DType no longer fits in the 4 bits to_id() gives it."); + static_assert(static_cast(GatedActivation::kNumTypes) <= 64, + "GatedActivation no longer fits in the 6 bits to_id() gives it."); + NVTE_CHECK(sm_arch < 512, "SM architecture no longer fits in the 9 bits to_id() gives it."); + return static_cast(dtype) | (static_cast(fp8_dtype) << 4) | + (static_cast(rowwise) << 8) | (static_cast(colwise) << 9) | + (static_cast(swizzled) << 10) | (static_cast(is_bwd) << 11) | + (static_cast(activation) << 12) | (sm_arch << 18); + } + + std::optional get_kernel() const { + static TVMFFIConfigCache &cache = TVMFFIConfigCache::create(); + return cache.get_or_load(*this); + } + + std::string to_key() const { + std::string key; + key.reserve(80); + key.append("cutedsl_mxfp8_gated_sm") + .append(std::to_string(sm_arch)) + .append("_") + .append(to_string(dtype)) + .append("_") + .append(to_string(fp8_dtype)) + .append("_") + .append(rowwise ? "1" : "0") + .append("_") + .append(colwise ? "1" : "0") + .append("_") + .append(swizzled ? "1" : "0") + .append("_") + .append(is_bwd ? "1" : "0") + .append("_") + .append(gated_activation_to_str(activation)); + return key; + } + + bool retrieve_func_from_python(const std::string &fn_name) const { + auto entrypoint = tvm::ffi::Function::GetGlobal(kEntrypointName); + if (!entrypoint.has_value()) { + return false; + } + tvm::ffi::Any result = + (*entrypoint)(tvm::ffi::String(fn_name), tvm::ffi::String(to_string(dtype)), + tvm::ffi::String(to_string(fp8_dtype)), rowwise, colwise, swizzled, is_bwd, + tvm::ffi::String(gated_activation_to_str(activation))); + return result.try_cast().value_or(false); + } +}; + +// Signature mirrors mxfp8::quantize_gated. Returns false to fall back to the CUDA kernel. +template +bool mxfp8_quantize_gated_cutedsl(const Tensor &gated_input, const Tensor &grad, Tensor *output, + const ParamOP &p, cudaStream_t stream) { + if (!TVMFFICentral::getInstance().get_cutedsl_backend_enabled()) { + maybe_warn_cutedsl_not_chosen("the CuTeDSL backend is disabled."); + return false; + } + constexpr GatedActivation activation = gated_activation_func_to_enum(); + if constexpr (activation == GatedActivation::kUnsupported) { + maybe_warn_cutedsl_not_chosen("the gated activation is not supported."); + return false; + } else { + const size_t rows = gated_input.flat_first_dim(); + const size_t cols = gated_input.flat_last_dim() / 2; + const MXFP8GatedQuantConfig config{/*dtype=*/gated_input.dtype(), + /*fp8_dtype=*/output->dtype(), + /*rowwise=*/output->has_data(), + /*colwise=*/output->has_columnwise_data(), + /*swizzled=*/output->with_gemm_swizzled_scales, + /*is_bwd=*/IS_BWD, + /*activation=*/activation}; + // The kernel offsets the gate half's swizzled scales by a fixed amount, which only matches + // the swizzled layout when the half spans whole 128-column scale tiles. + if (IS_BWD && config.swizzled && cols % 128 != 0) { + maybe_warn_cutedsl_not_chosen( + "backward gated activations with GEMM-swizzled scales need the last dimension of the " + "output halves to be divisible by 128."); + return false; + } + bool aligned = detail::is_16B_aligned(gated_input.data.dptr) && + (!config.rowwise || detail::is_16B_aligned(output->data.dptr)) && + (!config.colwise || detail::is_16B_aligned(output->columnwise_data.dptr)); + if constexpr (IS_BWD) { + aligned = aligned && detail::is_16B_aligned(grad.data.dptr); + } + if (!aligned) { + maybe_warn_cutedsl_not_chosen("an input or output is not 16-byte aligned."); + return false; + } + + std::optional fn = config.get_kernel(); + if (!fn.has_value()) { + return false; + } + checkCuDriverContext(stream); + + // Sanity checks and swizzled-scale padding, mirroring mxfp8::quantize_gated + if (config.rowwise) { + NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); + } + if (config.colwise) { + NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); + } + if (config.swizzled && (cols % 128 != 0 || rows % 128 != 0)) { + if (config.rowwise) { + NVTE_CHECK_CUDA(cudaMemsetAsync(output->scale_inv.dptr, 0, + output->scale_inv.buffer_size_bytes(), stream)); + } + if (config.colwise) { + NVTE_CHECK_CUDA(cudaMemsetAsync(output->columnwise_scale_inv.dptr, 0, + output->columnwise_scale_inv.buffer_size_bytes(), stream)); + } + } + + // Activation parameters travel as runtime arguments so that one compiled kernel serves them all. + float limit = 0.0f, alpha = 0.0f, glu_linear_offset = 0.0f, beta1 = 0.0f, beta2 = 0.0f; + if constexpr (std::is_same_v) { + limit = p.limit; + alpha = p.alpha; + glu_linear_offset = p.glu_linear_offset; + } else if constexpr (std::is_same_v) { + beta1 = p.beta1; + beta2 = p.beta2; + } + + const int32_t device_index = transformer_engine::cuda::current_device(); + DLTensorWrapper mGatedInput(gated_input.data, true, device_index); + DLTensorWrapper mGrad, mO_row, mS_row, mO_col, mS_col; + if constexpr (IS_BWD) { + mGrad = DLTensorWrapper(grad.data, true, device_index); + } + if (config.rowwise) { + mO_row = DLTensorWrapper(output->data, true, device_index); + mS_row = DLTensorWrapper(output->scale_inv, true, device_index); + } + if (config.colwise) { + mO_col = DLTensorWrapper(output->columnwise_data, true, device_index); + mS_col = DLTensorWrapper(output->columnwise_scale_inv, true, device_index); + } + (*fn)(&mGatedInput, &mGrad, &mO_row, &mS_row, &mO_col, &mS_col, limit, alpha, glu_linear_offset, + beta1, beta2, static_cast(stream)); + return true; + } +} + +} // namespace cutedsl_backend +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GATED_MXFP8_CUTEDSL_CUH_ diff --git a/transformer_engine/common/util/cutedsl_utils.h b/transformer_engine/common/util/cutedsl_utils.h index 4fd4ee56f69..4af0d78e906 100644 --- a/transformer_engine/common/util/cutedsl_utils.h +++ b/transformer_engine/common/util/cutedsl_utils.h @@ -90,6 +90,76 @@ NVTE_SPECIALIZE_CUTEDSL_ACTIVATION_FUNC_TO_ENUM(dsrelu, kDSReLU) #undef NVTE_SPECIALIZE_CUTEDSL_ACTIVATION_FUNC_TO_ENUM +// Encodes the gated activation selected by the gated dispatcher (gated_mxfp8.cuh), identified by +// its forward ActOP and, for backward, its DActOP. +enum class GatedActivation { + kGeGLU, + kQGeGLU, + kReGLU, + kSReGLU, + kSwiGLU, + kGLU, + kSiTUGLU, + kClampedSwiGLU, + kUnsupported, + kNumTypes +}; + +inline const char *gated_activation_to_str(GatedActivation act) { + switch (act) { + case GatedActivation::kGeGLU: + return "geglu"; + case GatedActivation::kQGeGLU: + return "qgeglu"; + case GatedActivation::kReGLU: + return "reglu"; + case GatedActivation::kSReGLU: + return "sreglu"; + case GatedActivation::kSwiGLU: + return "swiglu"; + case GatedActivation::kGLU: + return "glu"; + case GatedActivation::kSiTUGLU: + return "situglu"; + case GatedActivation::kClampedSwiGLU: + return "clamped_swiglu"; + case GatedActivation::kUnsupported: + case GatedActivation::kNumTypes: + return "unsupported"; + } + return "unsupported"; +} + +// DActOP is nullptr for the forward. +template +constexpr GatedActivation gated_activation_func_to_enum() { + return GatedActivation::kUnsupported; +} + +#define NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(PARAM, ACT, DACT, ACTIVATION) \ + template <> \ + constexpr GatedActivation gated_activation_func_to_enum, nullptr>() { \ + return GatedActivation::ACTIVATION; \ + } \ + template <> \ + constexpr GatedActivation \ + gated_activation_func_to_enum, DACT>() { \ + return GatedActivation::ACTIVATION; \ + } + +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, gelu, dgelu, kGeGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, qgelu, dqgelu, kQGeGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, relu, drelu, kReGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, srelu, dsrelu, kSReGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, silu, dsilu, kSwiGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, sigmoid, dsigmoid, kGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(SiTUGLUParam, situ_gate, dsitu_gate, kSiTUGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(ClampedSwiGLUParam, clamped_silu, + clamped_dsilu, kClampedSwiGLU) + +#undef NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM + } // namespace cutedsl_backend } // namespace transformer_engine From 82f59274149ea63ec945698266d68f95b832e315 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 30 Oct 2026 17:48:54 +0000 Subject: [PATCH 34/43] nit Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_gated_g2s.py | 184 +++++++++++------- .../common/cast/mxfp8/gated_mxfp8_cutedsl.cuh | 6 +- 2 files changed, 123 insertions(+), 67 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py index 2baa03dd3bf..c35d5728dc1 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py @@ -22,7 +22,7 @@ import cutlass from cutlass import cute from cutlass import pipeline -from cutlass import Boolean, Float32, Int32, Int64 +from cutlass import Boolean, Float32, Int32 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( @@ -89,11 +89,11 @@ def _gated_element( act_elt: Float32, gate_elt: Float32, grad_elt: Optional[Float32], - limit: Float32, - alpha: Float32, - glu_linear_offset: Float32, - beta1: Float32, - beta2: Float32, + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], ACTIVATION: cutlass.Constexpr[str], IS_BWD: cutlass.Constexpr[bool], ): @@ -154,11 +154,11 @@ def gated_elementwise_smem_tile( sAct_tile: cute.Tensor, # (TILE_Y, TILE_X) act half of the input; receives the act output sGate_tile: cute.Tensor, # (TILE_Y, TILE_X) gate half; receives the gate output (bwd only) sGrad_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) incoming gradient (bwd only) - limit: Float32, - alpha: Float32, - glu_linear_offset: Float32, - beta1: Float32, - beta2: Float32, + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], ACTIVATION: cutlass.Constexpr[str], IS_BWD: cutlass.Constexpr[bool], DTYPE: cutlass.Constexpr[type], @@ -226,11 +226,11 @@ def quantize_gated_mxfp8_smem_tile( tile_col_start: Int32, M: Int32, N: Int32, - limit: Float32, - alpha: Float32, - glu_linear_offset: Float32, - beta1: Float32, - beta2: Float32, + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], ACTIVATION: cutlass.Constexpr[str], IS_BWD: cutlass.Constexpr[bool], ROWWISE: cutlass.Constexpr[bool], @@ -304,21 +304,54 @@ def quantize_gated_mxfp8_smem_tile( @cute.jit -def _column_slice(mT: cute.Tensor, col_offset: Int32, num_cols: Int32) -> cute.Tensor: - """(rows, num_cols) view of the row-major mT starting at column col_offset. +def _split_gated_scales( + mS_row: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + SWIZZLED: cutlass.Constexpr[bool], +): + """Act and gate views of the backward output's scales; returns + (act_row, gate_row, act_col, gate_col). - The pointer is rebuilt rather than offset so that TMA sees it 16-byte aligned, which the C++ - dispatcher guarantees for both the tensor and the column offset of every slice. + The [M, 2N] output is [d act | d gate] and its scales are laid out for the whole output. Each + view has the shape of an [M, N] tensor's scales and the strides of the whole output's, and + starts where the whole output's layout places that half's first scale. For swizzled scales + this needs the halves to be whole 128x4 tiles (N % 128 == 0, required by the C++ dispatcher). """ - ptr = cute.make_ptr( - mT.element_type, - mT.iterator.toint() + Int64(col_offset) * Int64(mT.element_type.width // 8), - cute.AddressSpace.gmem, - assumed_align=16, - ) - return cute.make_tensor( - ptr, cute.make_layout((mT.shape[0], num_cols), stride=(mT.stride[0], 1)) - ) + act_row = gate_row = act_col = gate_col = None + if cutlass.const_expr(SWIZZLED): + mS_row, mS_col = derive_swizzled_scale_layout( + M, 2 * N, ROWWISE, COLWISE, mS_row, mS_col + ) + half_tiles = N // 128 + if cutlass.const_expr(ROWWISE): + row_tiler = ( + cute.make_layout((32, 4, cute.ceil_div(M, 128))), # all rows + cute.make_layout((4, half_tiles)), # one half's scale columns + ) + act_row = cute.local_tile(mS_row, row_tiler, (0, 0)) + gate_row = cute.local_tile(mS_row, row_tiler, (0, 1)) + if cutlass.const_expr(COLWISE): + col_tiler = ( + cute.make_layout((4, cute.ceil_div(cute.ceil_div(M, 32), 4))), # all scale rows + cute.make_layout((32, 4, half_tiles)), # one half's columns + ) + act_col = cute.local_tile(mS_col, col_tiler, (0, 0)) + gate_col = cute.local_tile(mS_col, col_tiler, (0, 1)) + else: + # Compact scales are row-major, so each half is a block of scale columns + if cutlass.const_expr(ROWWISE): + row_tiler = (mS_row.shape[0], N // MXFP8_BLOCK_SCALING_SIZE) + act_row = cute.local_tile(mS_row, row_tiler, (0, 0)) + gate_row = cute.local_tile(mS_row, row_tiler, (0, 1)) + if cutlass.const_expr(COLWISE): + col_tiler = (mS_col.shape[0], N) + act_col = cute.local_tile(mS_col, col_tiler, (0, 0)) + gate_col = cute.local_tile(mS_col, col_tiler, (0, 1)) + return act_row, gate_row, act_col, gate_col def _fake_gmem_tensor(dtype, shape, align): @@ -382,7 +415,12 @@ def compile(self): # Padded scale shapes, see MXFP8QuantizeConfig.compile 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)) - params = [Float32(0.0)] * 5 # limit, alpha, glu_linear_offset, beta1, beta2 + # Only the activation's own params are runtime arguments; the others are compiled as None + # and the C++ launcher passes None for them. + clamped = self.GATED_ACTIVATION == "clamped_swiglu" + situ = self.GATED_ACTIVATION == "situglu" + params = [Float32(0.0) if clamped else None] * 3 # limit, alpha, glu_linear_offset + params += [Float32(0.0) if situ else None] * 2 # beta1, beta2 return cute.compile( MXFP8GatedQuantizeG2SKernel(self), _fake_gmem_tensor(self.DTYPE, (sym_M, sym_2N), 16), # mGatedInput @@ -425,51 +463,67 @@ def __call__( self, mGatedInput: cute.Tensor, # [M, 2N] = [act | gate] mGrad: Optional[cute.Tensor], # [M, N], bwd only - mO_row: Optional[cute.Tensor], # [M, N] (fwd) or [M, 2N] (bwd) - mS_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - limit: Float32, # ClampedSwiGLUParam - alpha: Float32, - glu_linear_offset: Float32, - beta1: Float32, # SiTUGLUParam - beta2: Float32, + # Outputs, with C = N (FWD) or 2N (BWD) + mO_row: Optional[cute.Tensor], # [M, C] + mS_row: Optional[cute.Tensor], # [roundup(M, 128), roundup(C / 32, 4)] + mO_col: Optional[cute.Tensor], # [M, C], not transposed + mS_col: Optional[cute.Tensor], # [roundup(ceil(M / 32), 4), roundup(C, 128)] + # Special activation params, None if not used by the activation + limit: Optional[Float32], # Clamped SwiGLU only + alpha: Optional[Float32], # Clamped SwiGLU only + glu_linear_offset: Optional[Float32], # Clamped SwiGLU only + beta1: Optional[Float32], # SiTUGLU only + beta2: Optional[Float32], # SiTUGLU only stream: CUstream, ): - cfg = self.cfg + """ + mGatedInput: [a | g] with shape M x 2N, where a is the activation half and g is the gate half. + mGrad: None for FWD, [dy] with shape [M, N] for BWD + FWD mode: + out = quantize(act(a) * gate(g)) with shape [M, N] + BWD mode: + out = quantize([ dy * gate(g) * dact(a) | dy * act(a) * dgate(g) ]) with shape [M, 2N] + """ if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): - cute.printf(f"[CuTeDSL] MXFP8GatedQuantizeG2SKernel.__call__() with config: {cfg}\n") + cute.printf(f"[CuTeDSL] MXFP8GatedQuantizeG2SKernel.__call__() with config: {self.cfg}\n") - mO = mO_row if cfg.ROWWISE else mO_col M = mGatedInput.shape[0] - N = mGrad.shape[1] if cfg.IS_BWD else mO.shape[1] + # Pick any mO that is used just to get N + mO = mO_row if mO_row is not None else mO_col + # FWD: mO is [M, N]; BWD: mGrad is [M, N] + N = mGrad.shape[1] if self.cfg.IS_BWD else mO.shape[1] - # Scales cover the whole output: [M, 2N] backward, with the gate half's scales starting - # at scale column N / 32 (rowwise) or N (colwise). - if cutlass.const_expr(cfg.WITH_GEMM_SWIZZLED_SCALES): + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES and not self.cfg.IS_BWD): mS_row, mS_col = derive_swizzled_scale_layout( - M, mO.shape[1], cfg.ROWWISE, cfg.COLWISE, mS_row, mS_col + M, N, self.cfg.ROWWISE, self.cfg.COLWISE, mS_row, mS_col ) + mS_row_gate = None mS_col_gate = None mO_row_gate = None mO_col_gate = None - if cutlass.const_expr(cfg.IS_BWD): - # Swizzled scales are only translation invariant across whole 128-column tiles, - # which the C++ dispatcher guarantees by requiring N % 128 == 0 for them. - if cutlass.const_expr(cfg.ROWWISE): - mS_row_gate = cute.domain_offset((0, N // MXFP8_BLOCK_SCALING_SIZE), mS_row) - mO_row_gate = _column_slice(mO_row, N, N) - mO_row = _column_slice(mO_row, 0, N) - if cutlass.const_expr(cfg.COLWISE): - mS_col_gate = cute.domain_offset((0, N), mS_col) - mO_col_gate = _column_slice(mO_col, N, N) - mO_col = _column_slice(mO_col, 0, N) + if cutlass.const_expr(self.cfg.IS_BWD): + # Split the [M, 2N] output and its scales into the act and gate halves + mS_row, mS_row_gate, mS_col, mS_col_gate = _split_gated_scales( + mS_row, + mS_col, + M, + N, + self.cfg.ROWWISE, + self.cfg.COLWISE, + self.cfg.WITH_GEMM_SWIZZLED_SCALES, + ) + if cutlass.const_expr(self.cfg.ROWWISE): + mO_row_gate = cute.local_tile(mO_row, (M, N), (0, 1)) + mO_row = cute.local_tile(mO_row, (M, N), (0, 0)) + if cutlass.const_expr(self.cfg.COLWISE): + mO_col_gate = cute.local_tile(mO_col, (M, N), (0, 1)) + mO_col = cute.local_tile(mO_col, (M, N), (0, 0)) atoms = self._make_tma_atoms( mGrad, - _column_slice(mGatedInput, 0, N), - _column_slice(mGatedInput, N, N), + cute.local_tile(mGatedInput, (M, N), (0, 0)), # act half + cute.local_tile(mGatedInput, (M, N), (0, 1)), # gate half mO_row, mO_row_gate, mO_col, @@ -588,11 +642,11 @@ def kernel( mS_row_gate: Optional[cute.Tensor], mS_col_act: Optional[cute.Tensor], mS_col_gate: Optional[cute.Tensor], - limit: Float32, - alpha: Float32, - glu_linear_offset: Float32, - beta1: Float32, - beta2: Float32, + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], tma_atom_grad: Optional[cute.CopyAtom], tma_src_grad: Optional[cute.Tensor], tma_atom_act: cute.CopyAtom, diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh index eab31ed01fd..cb6a3a44d07 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh @@ -168,8 +168,10 @@ bool mxfp8_quantize_gated_cutedsl(const Tensor &gated_input, const Tensor &grad, } } - // Activation parameters travel as runtime arguments so that one compiled kernel serves them all. - float limit = 0.0f, alpha = 0.0f, glu_linear_offset = 0.0f, beta1 = 0.0f, beta2 = 0.0f; + // The activation's own parameters are runtime arguments; the kernel is compiled with None in + // the slots of the others, so they are passed as None (nullptr). + tvm::ffi::Any limit = nullptr, alpha = nullptr, glu_linear_offset = nullptr, beta1 = nullptr, + beta2 = nullptr; if constexpr (std::is_same_v) { limit = p.limit; alpha = p.alpha; From 4a0145dd0c3105f8a97ab42a014f82891ec4857a Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 30 Oct 2026 18:08:17 +0000 Subject: [PATCH 35/43] move gated config to common Signed-off-by: Kaining Zhong --- .../mxfp8/test_mxfp8_cutedsl_backend.py | 7 +- .../cast/mxfp8/quantize_mxfp8_common.py | 119 +++++++++++++++++- .../cast/mxfp8/quantize_mxfp8_gated_g2s.py | 119 +----------------- 3 files changed, 120 insertions(+), 125 deletions(-) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index 0c57f9c1be0..97712afed84 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -502,11 +502,8 @@ def run_gated_test_case(act, is_bwd, shape, block_size, in_dtype, fp8_dtype, swi shape is (M, N): the gated input is [M, 2N]. """ M, N = shape - if is_bwd and swizzled and N % 128 != 0: - pytest.skip( - "CuTeDSL hands backward gated activations with GEMM-swizzled scales and " - "N % 128 != 0 back to CUDA" - ) + if is_bwd and swizzled: + pytest.skip("CuTeDSL hands backward gated activations with GEMM-swizzled scales to CUDA") rowwise = block_size[1] != 1 columnwise = block_size[0] != 1 x, _ = generate_inputs(M, 2 * N, in_dtype) 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 85ece5fb0ce..6cc00b8bab0 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -30,6 +30,8 @@ dact_dsilu, dact_dqgelu, dact_dgelu, + act_sigmoid, + dact_dsigmoid, ) CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" @@ -68,6 +70,21 @@ } +# Gated activations of gated_mxfp8.cuh, as (ActOP, DActOP) ports of the math.h functions the +# CUDA dispatcher instantiates them with. swiglu and clamped_swiglu backward, and situglu in +# both directions, do not go through this table: like the CUDA kernel they take a special path. +GATED_ACTIVATIONS = { + "geglu": (act_gelu, dact_dgelu), + "qgeglu": (act_qgelu, dact_dqgelu), + "reglu": (act_relu, dact_drelu), + "sreglu": (act_srelu, dact_dsrelu), + "glu": (act_sigmoid, dact_dsigmoid), + "swiglu": (act_silu, None), + "clamped_swiglu": (None, None), + "situglu": (None, None), +} + + @cute.jit def derive_swizzled_scale_layout( M: Int32, @@ -341,6 +358,105 @@ def compile(self): return compiled +def _fake_gmem_tensor(dtype, shape, align): + """Row-major fake gmem tensor for cute.compile.""" + return cute.runtime.make_fake_compact_tensor( + dtype, + shape, + stride_order=tuple(reversed(range(len(shape)))), + memspace=cute.AddressSpace.gmem, + assumed_align=align, + ) + + +class MXFP8GatedQuantizeConfig(MXFP8QuantizeConfig): + """Compile-time config of gated_mxfp8.cuh's port: a plain MXFP8 cast config plus the gated op.""" + + def __init__( + self, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + is_bwd: bool, + activation: str, + ): + super().__init__( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + with_amax=False, + activation="none", + ) + if activation not in GATED_ACTIVATIONS: + raise ValueError( + f"unknown gated activation {activation!r}; expected one of" + f" {sorted(GATED_ACTIVATIONS)}" + ) + if is_bwd and with_gemm_swizzled_scales: + # TODO(kainingz): technically this is supported but I need to come up with a less disgusting way + # to implement it in CuTe + raise ValueError("backward gated activations with GEMM-swizzled scales") + self.IS_BWD = is_bwd + self.GATED_ACTIVATION = activation + + def __str__(self): + return ( + f"MXFP8GatedQuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " + f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, " + f"swizzled={self.WITH_GEMM_SWIZZLED_SCALES}, is_bwd={self.IS_BWD}, " + f"activation={self.GATED_ACTIVATION})" + ) + + __repr__ = __str__ + + def compile(self): + """Return the compiled CuTeDSL function object for this config.""" + # The kernel module imports this one + # pylint: disable-next=import-outside-toplevel + from .quantize_mxfp8_gated_g2s import MXFP8GatedQuantizeG2SKernel + + sym_M = cute.sym_int32() + # The C++ dispatcher requires N % 32 == 0, like the CUDA kernel. + sym_N = cute.sym_int32(divisibility=32) + sym_2N = cute.sym_int32(divisibility=64) + out_cols = sym_2N if self.IS_BWD else sym_N + # Padded scale shapes, see MXFP8QuantizeConfig.compile + 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)) + # Only the activation's own params are runtime arguments; the others are compiled as None + # and the C++ launcher passes None for them. + clamped = self.GATED_ACTIVATION == "clamped_swiglu" + situ = self.GATED_ACTIVATION == "situglu" + params = [Float32(0.0) if clamped else None] * 3 # limit, alpha, glu_linear_offset + params += [Float32(0.0) if situ else None] * 2 # beta1, beta2 + return cute.compile( + MXFP8GatedQuantizeG2SKernel(self), + _fake_gmem_tensor(self.DTYPE, (sym_M, sym_2N), 16), # mGatedInput + _fake_gmem_tensor(self.DTYPE, (sym_M, sym_N), 16) if self.IS_BWD else None, # mGrad + # mO_row, mS_row + _fake_gmem_tensor(self.FP8_DTYPE, (sym_M, out_cols), 16) if self.ROWWISE else None, + ( + _fake_gmem_tensor(cutlass.Float8E8M0FNU, scale_rowwise_shape, 4) + if self.ROWWISE + else None + ), + # mO_col, mS_col + _fake_gmem_tensor(self.FP8_DTYPE, (sym_M, out_cols), 16) if self.COLWISE else None, + ( + _fake_gmem_tensor(cutlass.Float8E8M0FNU, scale_colwise_shape, 4) + if self.COLWISE + else None + ), + *params, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + + class MXFP8QuantizeKernelBase(abc.ABC): """Base class for MXFP8 quantize kernels.""" @@ -631,9 +747,6 @@ def get_mxfp8_gated_quantization_function( activation: str, ) -> bool: """Compile and register the gated MXFP8 quantize kernel for this config under `fn_name`.""" - # pylint: disable=import-outside-toplevel - from .quantize_mxfp8_gated_g2s import MXFP8GatedQuantizeConfig - try: cfg = MXFP8GatedQuantizeConfig( dtype=dtype, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py index c35d5728dc1..fa0716da8ae 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py @@ -32,19 +32,8 @@ truncate_f32_f16, ) from transformer_engine.common.CuTeDSL.activations import ( - act_gelu, - act_qgelu, - act_relu, - act_srelu, - act_silu, - act_sigmoid, act_clamped_silu, act_situ_gate, - dact_dgelu, - dact_dqgelu, - dact_drelu, - dact_dsrelu, - dact_dsigmoid, dact_dsitu_gate, situ_up, dsitu_up, @@ -52,7 +41,8 @@ ) from .quantize_mxfp8_common import ( CUTEDSL_DEBUG_LOGGING, - MXFP8QuantizeConfig, + GATED_ACTIVATIONS, + MXFP8GatedQuantizeConfig, MXFP8_BLOCK_SCALING_SIZE, derive_swizzled_scale_layout, ) @@ -61,20 +51,6 @@ quantize_mxfp8_smem_tile_general, ) -# Gated activations of gated_mxfp8.cuh, as (ActOP, DActOP) ports of the math.h functions the -# CUDA dispatcher instantiates them with. swiglu and clamped_swiglu backward, and situglu in -# both directions, do not go through this table: like the CUDA kernel they take a special path. -GATED_ACTIVATIONS = { - "geglu": (act_gelu, dact_dgelu), - "qgeglu": (act_qgelu, dact_dqgelu), - "reglu": (act_relu, dact_drelu), - "sreglu": (act_srelu, dact_dsrelu), - "glu": (act_sigmoid, dact_dsigmoid), - "swiglu": (act_silu, None), - "clamped_swiglu": (None, None), - "situglu": (None, None), -} - @cute.jit def _select_f32(pred: Boolean, if_true: Float32, if_false: Float32) -> Float32: @@ -354,97 +330,6 @@ def _split_gated_scales( return act_row, gate_row, act_col, gate_col -def _fake_gmem_tensor(dtype, shape, align): - """Row-major fake gmem tensor for cute.compile.""" - return cute.runtime.make_fake_compact_tensor( - dtype, - shape, - stride_order=tuple(reversed(range(len(shape)))), - memspace=cute.AddressSpace.gmem, - assumed_align=align, - ) - - -class MXFP8GatedQuantizeConfig(MXFP8QuantizeConfig): - """Compile-time config of gated_mxfp8.cuh's port: a plain MXFP8 cast config plus the gated op.""" - - def __init__( - self, - dtype: str, - fp8_dtype: str, - rowwise: bool, - colwise: bool, - with_gemm_swizzled_scales: bool, - is_bwd: bool, - activation: str, - ): - super().__init__( - dtype=dtype, - fp8_dtype=fp8_dtype, - rowwise=rowwise, - colwise=colwise, - with_gemm_swizzled_scales=with_gemm_swizzled_scales, - with_amax=False, - activation="none", - ) - if activation not in GATED_ACTIVATIONS: - raise ValueError( - f"unknown gated activation {activation!r}; expected one of" - f" {sorted(GATED_ACTIVATIONS)}" - ) - self.IS_BWD = is_bwd - self.GATED_ACTIVATION = activation - - def __str__(self): - return ( - f"MXFP8GatedQuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " - f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, " - f"swizzled={self.WITH_GEMM_SWIZZLED_SCALES}, is_bwd={self.IS_BWD}, " - f"activation={self.GATED_ACTIVATION})" - ) - - __repr__ = __str__ - - def compile(self): - """Return the compiled CuTeDSL function object for this config.""" - sym_M = cute.sym_int32() - # The C++ dispatcher requires N % 32 == 0, like the CUDA kernel. - sym_N = cute.sym_int32(divisibility=32) - sym_2N = cute.sym_int32(divisibility=64) - out_cols = sym_2N if self.IS_BWD else sym_N - # Padded scale shapes, see MXFP8QuantizeConfig.compile - 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)) - # Only the activation's own params are runtime arguments; the others are compiled as None - # and the C++ launcher passes None for them. - clamped = self.GATED_ACTIVATION == "clamped_swiglu" - situ = self.GATED_ACTIVATION == "situglu" - params = [Float32(0.0) if clamped else None] * 3 # limit, alpha, glu_linear_offset - params += [Float32(0.0) if situ else None] * 2 # beta1, beta2 - return cute.compile( - MXFP8GatedQuantizeG2SKernel(self), - _fake_gmem_tensor(self.DTYPE, (sym_M, sym_2N), 16), # mGatedInput - _fake_gmem_tensor(self.DTYPE, (sym_M, sym_N), 16) if self.IS_BWD else None, # mGrad - # mO_row, mS_row - _fake_gmem_tensor(self.FP8_DTYPE, (sym_M, out_cols), 16) if self.ROWWISE else None, - ( - _fake_gmem_tensor(cutlass.Float8E8M0FNU, scale_rowwise_shape, 4) - if self.ROWWISE - else None - ), - # mO_col, mS_col - _fake_gmem_tensor(self.FP8_DTYPE, (sym_M, out_cols), 16) if self.COLWISE else None, - ( - _fake_gmem_tensor(cutlass.Float8E8M0FNU, scale_colwise_shape, 4) - if self.COLWISE - else None - ), - *params, - cute.runtime.make_fake_stream(), - options="--enable-tvm-ffi", - ) - - class MXFP8GatedQuantizeG2SKernel(MXFP8QuantizeGeneralG2SKernel): """gated_mxfp8.cuh in CuTeDSL. The general kernel's pipeline, loading the act and gate halves (and grad, backward) each stage and storing one (forward) or two (backward) quantized halves. From 6c37b244201c1d52ada449b19d8e18b3e988b327 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Sat, 31 Oct 2026 01:15:30 +0000 Subject: [PATCH 36/43] overhaul Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_gated_g2s.py | 612 +++++++++--------- 1 file changed, 308 insertions(+), 304 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py index fa0716da8ae..3035967c217 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py @@ -127,9 +127,9 @@ def _gated_element( @cute.jit def gated_elementwise_smem_tile( - sAct_tile: cute.Tensor, # (TILE_Y, TILE_X) act half of the input; receives the act output - sGate_tile: cute.Tensor, # (TILE_Y, TILE_X) gate half; receives the gate output (bwd only) - sGrad_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) incoming gradient (bwd only) + sAct_tile: cute.Tensor, + sGate_tile: cute.Tensor, + sGrad_tile: Optional[cute.Tensor], limit: Optional[Float32], alpha: Optional[Float32], glu_linear_offset: Optional[Float32], @@ -186,29 +186,16 @@ def gated_elementwise_smem_tile( @cute.jit -def quantize_gated_mxfp8_smem_tile( - sAct_tile: cute.Tensor, - sGate_tile: cute.Tensor, - sGrad_tile: Optional[cute.Tensor], - sO_row_act_tile: Optional[cute.Tensor], - mS_row_act_stage: Optional[cute.Tensor], - sO_row_gate_tile: Optional[cute.Tensor], # bwd only - mS_row_gate_stage: Optional[cute.Tensor], # bwd only - sO_col_act_tile: Optional[cute.Tensor], - mS_col_act_stage: Optional[cute.Tensor], - sO_col_gate_tile: Optional[cute.Tensor], # bwd only - mS_col_gate_stage: Optional[cute.Tensor], # bwd only +def quantize_mxfp8_smem_tile_general_partial( + sX_tile: cute.Tensor, + sO_row_tile: Optional[cute.Tensor], + mS_row_stage: Optional[cute.Tensor], + sO_col_tile: Optional[cute.Tensor], + mS_col_stage: Optional[cute.Tensor], tile_row_start: Int32, tile_col_start: Int32, M: Int32, N: Int32, - limit: Optional[Float32], - alpha: Optional[Float32], - glu_linear_offset: Optional[Float32], - beta1: Optional[Float32], - beta2: Optional[Float32], - ACTIVATION: cutlass.Constexpr[str], - IS_BWD: cutlass.Constexpr[bool], ROWWISE: cutlass.Constexpr[bool], COLWISE: cutlass.Constexpr[bool], MAX_NORM_RCP: cutlass.Constexpr[float], @@ -218,65 +205,39 @@ def quantize_gated_mxfp8_smem_tile( TILE_X: cutlass.Constexpr[int], TILE_Y: cutlass.Constexpr[int], ): - """Apply the gated op to one staged tile, then MXFP8-quantize the result(s) as a plain cast.""" - gated_elementwise_smem_tile( - sAct_tile, - sGate_tile, - sGrad_tile, - limit, - alpha, - glu_linear_offset, - beta1, - beta2, - ACTIVATION=ACTIVATION, - IS_BWD=IS_BWD, + """A partial quantize_mxfp8_smem_tile_general function with some args fixed.""" + quantize_mxfp8_smem_tile_general( + sX_tile, + None, + sO_row_tile, + mS_row_stage, + sO_col_tile, + mS_col_stage, + None, + tile_row_start, + tile_col_start, + M, + N, + ROWWISE=ROWWISE, + COLWISE=COLWISE, + WITH_AMAX=False, + WITH_ACT=False, + WITH_DACT=False, + ACTIVATION=None, + MAX_NORM_RCP=MAX_NORM_RCP, DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=WITH_GEMM_SWIZZLED_SCALES, TILE_X=TILE_X, + TILE_Y=TILE_Y, + # The gated op maps the zeros TMA fills out-of-bounds elements with to zero. + SKIP_INPUT_MASKING=True, + SKIP_SCALE_BOUNDS=False, + CACHE_ACTIVATION=False, + DBIAS_REDUCTION_IN_COLWISE=False, + COLWISE_DBIAS_REDUCTION_ONLY=False, + DBIAS_REDUCTION_IN_ROWWISE=False, ) - # The rowwise pass reads elements other threads wrote. - if cutlass.const_expr(ROWWISE): - cute.arch.sync_threads() - - # The act output and, backward, the gate output. Both share this tile's coordinates; only - # their output buffers and scale views differ. - halves = [(sAct_tile, sO_row_act_tile, mS_row_act_stage, sO_col_act_tile, mS_col_act_stage)] - if cutlass.const_expr(IS_BWD): - halves.append( - (sGate_tile, sO_row_gate_tile, mS_row_gate_stage, sO_col_gate_tile, mS_col_gate_stage) - ) - for sX_tile, sO_row_tile, mS_row_stage, sO_col_tile, mS_col_stage in halves: - quantize_mxfp8_smem_tile_general( - sX_tile, - None, - sO_row_tile, - mS_row_stage, - sO_col_tile, - mS_col_stage, - None, - tile_row_start, - tile_col_start, - M, - N, - ROWWISE=ROWWISE, - COLWISE=COLWISE, - WITH_AMAX=False, - WITH_ACT=False, - WITH_DACT=False, - ACTIVATION=None, - MAX_NORM_RCP=MAX_NORM_RCP, - DTYPE=DTYPE, - FP8_DTYPE=FP8_DTYPE, - WITH_GEMM_SWIZZLED_SCALES=WITH_GEMM_SWIZZLED_SCALES, - TILE_X=TILE_X, - TILE_Y=TILE_Y, - # The gated op maps the zeros TMA fills out-of-bounds elements with to zero. - SKIP_INPUT_MASKING=True, - SKIP_SCALE_BOUNDS=False, - CACHE_ACTIVATION=False, - DBIAS_REDUCTION_IN_COLWISE=False, - COLWISE_DBIAS_REDUCTION_ONLY=False, - DBIAS_REDUCTION_IN_ROWWISE=False, - ) @cute.jit @@ -342,6 +303,11 @@ def __init__(self, cfg: MXFP8GatedQuantizeConfig): super().__init__(cfg) # Same occupancy hint as the general kernel with a fused activation self._MIN_BLOCKS_PER_MP = 1 + # Every pipeline stage has its own SMEM buffers, so a CTA with fewer tiles than stages + # would allocate buffers it never fills + assert ( + self._NUM_STAGES <= self._NUM_TILES + ), f"{self._NUM_STAGES} pipeline stages for {self._NUM_TILES} tiles per CTA" @cute.jit def __call__( @@ -376,20 +342,18 @@ def __call__( # Pick any mO that is used just to get N mO = mO_row if mO_row is not None else mO_col # FWD: mO is [M, N]; BWD: mGrad is [M, N] - N = mGrad.shape[1] if self.cfg.IS_BWD else mO.shape[1] + N = mO.shape[1] if not self.cfg.IS_BWD else mGrad.shape[1] - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES and not self.cfg.IS_BWD): - mS_row, mS_col = derive_swizzled_scale_layout( - M, N, self.cfg.ROWWISE, self.cfg.COLWISE, mS_row, mS_col - ) + # The act and gate halves of the [M, 2N] input + mAct = cute.local_tile(mGatedInput, (M, N), (0, 0)) + mGate = cute.local_tile(mGatedInput, (M, N), (0, 1)) - mS_row_gate = None - mS_col_gate = None - mO_row_gate = None - mO_col_gate = None + mO_row_dact = mO_row_dgate = mO_col_dact = mO_col_dgate = None # BWD only + mS_row_dact = mS_row_dgate = mS_col_dact = mS_col_dgate = None # BWD only if cutlass.const_expr(self.cfg.IS_BWD): - # Split the [M, 2N] output and its scales into the act and gate halves - mS_row, mS_row_gate, mS_col, mS_col_gate = _split_gated_scales( + # The [M, 2N] output is [d act | d gate]; the kernel takes its two halves + # TODO(kainingz): how to implement this in a more elegant way + mS_row_dact, mS_row_dgate, mS_col_dact, mS_col_dgate = _split_gated_scales( mS_row, mS_col, M, @@ -399,29 +363,39 @@ def __call__( self.cfg.WITH_GEMM_SWIZZLED_SCALES, ) if cutlass.const_expr(self.cfg.ROWWISE): - mO_row_gate = cute.local_tile(mO_row, (M, N), (0, 1)) - mO_row = cute.local_tile(mO_row, (M, N), (0, 0)) + mO_row_dact = cute.local_tile(mO_row, (M, N), (0, 0)) + mO_row_dgate = cute.local_tile(mO_row, (M, N), (0, 1)) if cutlass.const_expr(self.cfg.COLWISE): - mO_col_gate = cute.local_tile(mO_col, (M, N), (0, 1)) - mO_col = cute.local_tile(mO_col, (M, N), (0, 0)) + mO_col_dact = cute.local_tile(mO_col, (M, N), (0, 0)) + mO_col_dgate = cute.local_tile(mO_col, (M, N), (0, 1)) + # In BWD mode we won't use these, and we should set them to None so we don't build TMA atoms for them + mO_row = mO_col = mS_row = mS_col = None + elif 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 + ) atoms = self._make_tma_atoms( mGrad, - cute.local_tile(mGatedInput, (M, N), (0, 0)), # act half - cute.local_tile(mGatedInput, (M, N), (0, 1)), # gate half + mAct, + mGate, mO_row, - mO_row_gate, mO_col, - mO_col_gate, + mO_row_dact, + mO_row_dgate, + mO_col_dact, + mO_col_dgate, ) # The TMA atoms are spread from a list, which pylint cannot count self.kernel( # pylint: disable=no-value-for-parameter M, N, mS_row, - mS_row_gate, mS_col, - mS_col_gate, + mS_row_dact, + mS_row_dgate, + mS_col_dact, + mS_col_dgate, limit, alpha, glu_linear_offset, @@ -429,30 +403,29 @@ def __call__( beta2, *atoms, ).launch( - grid=self._grid(M, N), + grid=[ + cute.ceil_div(N, self._TILE_COLS), + cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES), + ], block=[self._THREADS_PER_CTA], min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, stream=stream, ) - def _grid(self, M: Int32, N: Int32): - return [ - cute.ceil_div(N, self._TILE_COLS), - cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES), - ] - @cute.jit def _make_tma_atoms( self, mGrad: Optional[cute.Tensor], mAct: cute.Tensor, mGate: cute.Tensor, - mO_row_act: Optional[cute.Tensor], - mO_row_gate: Optional[cute.Tensor], - mO_col_act: Optional[cute.Tensor], - mO_col_gate: Optional[cute.Tensor], + mO_row: Optional[cute.Tensor], # FWD only + mO_col: Optional[cute.Tensor], # FWD only + mO_row_dact: Optional[cute.Tensor], # BWD only + mO_row_dgate: Optional[cute.Tensor], # BWD only + mO_col_dact: Optional[cute.Tensor], # BWD only + mO_col_dgate: Optional[cute.Tensor], # BWD only ): - """TMA atoms and their gmem views for every (TILE_ROWS, TILE_COLS) operand; None if absent.""" + """Build TMA atoms for input / quantized output; scales will not be loaded / stored by TMA.""" smem_tile_layout = cute.make_ordered_layout( (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) ) @@ -465,10 +438,12 @@ def _make_tma_atoms( (op_load, mGrad), (op_load, mAct), (op_load, mGate), - (op_store, mO_row_act), - (op_store, mO_row_gate), - (op_store, mO_col_act), - (op_store, mO_col_gate), + (op_store, mO_row), + (op_store, mO_col), + (op_store, mO_row_dact), + (op_store, mO_row_dgate), + (op_store, mO_col_dact), + (op_store, mO_col_dgate), ): if cutlass.const_expr(mT is None): atoms.extend([None, None]) @@ -480,73 +455,74 @@ def _make_tma_atoms( ) return atoms - @cute.jit - def _copy_inputs_to_stage( - self, - tile_y: Int32, - tile_x: Int32, - stage: Int32, - tma_bar_ptr: cute.Pointer, - tma_atom_grad: Optional[cute.CopyAtom], - tXgGrad: Optional[cute.Tensor], - tXsGrad: Optional[cute.Tensor], - tma_atom_act: cute.CopyAtom, - tXgAct: cute.Tensor, - tXsAct: cute.Tensor, - tma_atom_gate: cute.CopyAtom, - tXgGate: cute.Tensor, - tXsGate: cute.Tensor, - ): - """Issue this tile's input TMA loads into the pipeline stage, all on one barrier.""" - if cutlass.const_expr(self.cfg.IS_BWD): - cute.copy( - tma_atom_grad, - tXgGrad[(None, (tile_y, tile_x))], - tXsGrad[(None, stage)], - tma_bar_ptr=tma_bar_ptr, - ) - cute.copy( - tma_atom_act, - tXgAct[(None, (tile_y, tile_x))], - tXsAct[(None, stage)], - tma_bar_ptr=tma_bar_ptr, - ) - cute.copy( - tma_atom_gate, - tXgGate[(None, (tile_y, tile_x))], - tXsGate[(None, stage)], - tma_bar_ptr=tma_bar_ptr, - ) - @cute.kernel def kernel( self, M: Int32, N: Int32, - mS_row_act: Optional[cute.Tensor], - mS_row_gate: Optional[cute.Tensor], - mS_col_act: Optional[cute.Tensor], - mS_col_gate: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], # FWD only + mS_col: Optional[cute.Tensor], # FWD only + mS_row_dact: Optional[cute.Tensor], # BWD only + mS_row_dgate: Optional[cute.Tensor], # BWD only + mS_col_dact: Optional[cute.Tensor], # BWD only + mS_col_dgate: Optional[cute.Tensor], # BWD only limit: Optional[Float32], alpha: Optional[Float32], glu_linear_offset: Optional[Float32], beta1: Optional[Float32], beta2: Optional[Float32], + # TMA for mGrad, only used in BWD tma_atom_grad: Optional[cute.CopyAtom], tma_src_grad: Optional[cute.Tensor], + # TMA for mAct and mGate, used in both FWD and BWD tma_atom_act: cute.CopyAtom, tma_src_act: cute.Tensor, tma_atom_gate: cute.CopyAtom, tma_src_gate: cute.Tensor, - tma_atom_row_act: Optional[cute.CopyAtom], - tma_dst_row_act: Optional[cute.Tensor], - tma_atom_row_gate: Optional[cute.CopyAtom], - tma_dst_row_gate: Optional[cute.Tensor], - tma_atom_col_act: Optional[cute.CopyAtom], - tma_dst_col_act: Optional[cute.Tensor], - tma_atom_col_gate: Optional[cute.CopyAtom], - tma_dst_col_gate: Optional[cute.Tensor], + # TMA for mO_row/col, only used in FWD + tma_atom_row: Optional[cute.CopyAtom], + tma_dst_row: Optional[cute.Tensor], + tma_atom_col: Optional[cute.CopyAtom], + tma_dst_col: Optional[cute.Tensor], + # TMA for mO_row/col_dact/dgate, only used in BWD + tma_atom_row_dact: Optional[cute.CopyAtom], + tma_dst_row_dact: Optional[cute.Tensor], + tma_atom_row_dgate: Optional[cute.CopyAtom], + tma_dst_row_dgate: Optional[cute.Tensor], + tma_atom_col_dact: Optional[cute.CopyAtom], + tma_dst_col_dact: Optional[cute.Tensor], + tma_atom_col_dgate: Optional[cute.CopyAtom], + tma_dst_col_dgate: Optional[cute.Tensor], ): + """ + FWD mode: + mAct: activation half input + mGate: gate half input + mO_row/col: quantized output for row / column directions + mS_row/col: output scales for row / column directions + + BWD mode: + mGrad: incoming gradient + mAct: activation half input from FWD pass + mGate: gate half input from FWD pass + mO_row/col_dact: quantized activation gradient for row / column directions + mO_row/col_dgate: quantized gate gradient for row / column directions + mS_row/col_dact: output scales for activation gradient row / column directions + mS_row/col_dgate: output scales for gate gradient row / column directions + mO_row/col: quantized gradient (both activation and gate) for row / column directions + mS_row/col: output scales for gradient (both activation and gate) row / column directions + + mAct: both FWD and BWD + mGate: both FWD and BWD + mGrad: BWD only + mO_row/col: both FWD (N) and BWD (2N) + mS_row/col: both FWD (N) and BWD (2N) + mO_row/col_dact: BWD only + mO_row/col_dgate: BWD only + mS_row/col_dact: BWD only + mS_row/col_dgate: BWD only + """ + cfg = self.cfg DTYPE = cfg.DTYPE TILE_ROWS = self._TILE_ROWS @@ -554,17 +530,9 @@ def kernel( NUM_STAGES = self._NUM_STAGES NUM_TILES = self._NUM_TILES - # Scale views tiled like the general kernel's + # The scales of one (TILE_ROWS, TILE_COLS) tile row_scale_tile = (TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) col_scale_tile = (TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, TILE_COLS) - if cutlass.const_expr(cfg.ROWWISE): - mS_row_act = cute.zipped_divide(mS_row_act, row_scale_tile) - if cutlass.const_expr(cfg.IS_BWD): - mS_row_gate = cute.zipped_divide(mS_row_gate, row_scale_tile) - if cutlass.const_expr(cfg.COLWISE): - mS_col_act = cute.zipped_divide(mS_col_act, col_scale_tile) - if cutlass.const_expr(cfg.IS_BWD): - mS_col_gate = cute.zipped_divide(mS_col_gate, col_scale_tile) # Shared memory: one ((TILE_ROWS, TILE_COLS), NUM_STAGES) ring buffer per operand smem = cutlass.memory.SmemAllocator() @@ -574,40 +542,47 @@ class BarrierStorage: # pylint: disable=missing-class-docstring mbar_storage: cute.struct.MemRange[cute.Int64, 2 * NUM_STAGES] storage = smem.allocate(BarrierStorage) - stage_layout = cute.make_layout( + stages_layout = cute.make_layout( ((TILE_ROWS, TILE_COLS), NUM_STAGES), stride=((TILE_COLS, 1), TILE_ROWS * TILE_COLS), ) - sAct = smem.allocate_tensor(DTYPE, stage_layout, byte_alignment=128) - sGate = smem.allocate_tensor(DTYPE, stage_layout, byte_alignment=128) + + # Allocate SMEM buffers for inputs + sAct = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) + sGate = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) sGrad = None if cutlass.const_expr(cfg.IS_BWD): - sGrad = smem.allocate_tensor(DTYPE, stage_layout, byte_alignment=128) - sO_row_act = None - sO_row_gate = None - sO_col_act = None - sO_col_gate = None + sGrad = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) + + # Allocate SMEM buffers for outputs + sO_row = sO_col = None + sO_row_dact = sO_row_dgate = sO_col_dact = sO_col_dgate = None if cutlass.const_expr(cfg.ROWWISE): - sO_row_act = smem.allocate_tensor(cfg.FP8_DTYPE, stage_layout, byte_alignment=128) if cutlass.const_expr(cfg.IS_BWD): - sO_row_gate = smem.allocate_tensor(cfg.FP8_DTYPE, stage_layout, byte_alignment=128) + sO_row_dact = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + sO_row_dgate = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + else: + sO_row = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) if cutlass.const_expr(cfg.COLWISE): - sO_col_act = smem.allocate_tensor(cfg.FP8_DTYPE, stage_layout, byte_alignment=128) if cutlass.const_expr(cfg.IS_BWD): - sO_col_gate = smem.allocate_tensor(cfg.FP8_DTYPE, stage_layout, byte_alignment=128) + sO_col_dact = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + sO_col_dgate = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + else: + sO_col = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) bidx, bidy, _ = cute.arch.block_idx() + # Issue TMA descriptor prefetch using the TMA warp (warp 0) if warp_idx == 0: - if cutlass.const_expr(cfg.IS_BWD): - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_grad) - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_act) - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_gate) + for load_atom in (tma_atom_act, tma_atom_gate, tma_atom_grad): + if cutlass.const_expr(load_atom is not None): + cute.nvgpu.cpasync.prefetch_descriptor(load_atom) - # Warp 0's elected thread produces (issues TMA), every warp consumes + # TMA warp (warp 0)'s elected thread produces (issues TMA), every warp consumes producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._NUM_WARPS) + # Input has 2 tiles (act + gate) in FWD, 3 tiles (grad + act + gate) in BWD num_inputs = 3 if cfg.IS_BWD else 2 mainloop_pipeline = pipeline.PipelineTmaAsync.create( barrier_storage=storage.mbar_storage.data_ptr(), @@ -621,35 +596,64 @@ class BarrierStorage: # pylint: disable=missing-class-docstring prod_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, NUM_STAGES) cons_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, NUM_STAGES) - # Partition every operand for TMA (single CTA, no multicast) - cta_tiler = (TILE_ROWS, TILE_COLS) - tXsGrad = None - tXgGrad = None + cta_coord = (bidy, bidx) + # Shape of each individual tile. A CTA processes one tile at a time + tile_shape = (TILE_ROWS, TILE_COLS) + # The block processed by a CTA, consisting of NUM_TILES vertically stacked tiles + cta_block = (NUM_TILES * TILE_ROWS, TILE_COLS) + + # The scales block processed by a CTA, consisting of NUM_TILES vertically stacked scale tiles + cta_block_row_scales = (NUM_TILES * TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + cta_block_col_scales = (NUM_TILES * TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, TILE_COLS) + + # Prepare parameters to pass to the SMEM tile quantize function by dact / dgate / output, where + # for each group, input, quantized output are SMEM tiles, and scales are GMEM tiles if cutlass.const_expr(cfg.IS_BWD): - tXsGrad, tXgGrad = cute.nvgpu.cpasync.tma_partition( - tma_atom_grad, - 0, - cute.make_layout(1), - sGrad, - cute.zipped_divide(tma_src_grad, cta_tiler), - ) - tXsAct, tXgAct = cute.nvgpu.cpasync.tma_partition( - tma_atom_act, 0, cute.make_layout(1), sAct, cute.zipped_divide(tma_src_act, cta_tiler) - ) - tXsGate, tXgGate = cute.nvgpu.cpasync.tma_partition( - tma_atom_gate, - 0, - cute.make_layout(1), - sGate, - cute.zipped_divide(tma_src_gate, cta_tiler), - ) - # Output (atom, smem partition, gmem partition) triples, act half first - outputs = [] + if cutlass.const_expr(cfg.ROWWISE): + mS_row_dact = cute.local_tile(mS_row_dact, cta_block_row_scales, cta_coord) + mS_row_dgate = cute.local_tile(mS_row_dgate, cta_block_row_scales, cta_coord) + if cutlass.const_expr(cfg.COLWISE): + mS_col_dact = cute.local_tile(mS_col_dact, cta_block_col_scales, cta_coord) + mS_col_dgate = cute.local_tile(mS_col_dgate, cta_block_col_scales, cta_coord) + # In BWD mode we output gradient for both the act and gate halves, which are stored in sAct and sGate respectively + tile_quantizer_params = [ + (sAct, sO_row_dact, mS_row_dact, sO_col_dact, mS_col_dact), # We will write unquantized dact to sAct in-place + (sGate, sO_row_dgate, mS_row_dgate, sO_col_dgate, mS_col_dgate), # We will write unquantized dgate to sGate in-place + ] + else: + if cutlass.const_expr(cfg.ROWWISE): + mS_row = cute.local_tile(mS_row, cta_block_row_scales, cta_coord) + if cutlass.const_expr(cfg.COLWISE): + mS_col = cute.local_tile(mS_col, cta_block_col_scales, cta_coord) + # In FWD mode we only output the result after the gated op, which is stored in sAct + tile_quantizer_params = [(sAct, sO_row, mS_row, sO_col, mS_col)] # We will write unquantized output to sAct in-place + + # Prepare input TMA (atom, smem partition, gmem partition) triples: act and gate, and BWD's grad + inputs = [] + for atom_in, src_in, sX_in in ( + (tma_atom_act, tma_src_act, sAct), + (tma_atom_gate, tma_src_gate, sGate), + (tma_atom_grad, tma_src_grad, sGrad), + ): + if cutlass.const_expr(atom_in is not None): + tXsX_in, tXgX_in = cute.nvgpu.cpasync.tma_partition( + atom_in, + 0, + cute.make_layout(1), + sX_in, + cute.zipped_divide(cute.local_tile(src_in, cta_block, cta_coord), tile_shape), + ) + inputs.append((atom_in, tXsX_in, tXgX_in)) + + # Prepare output TMA (atom, smem partition, gmem partition) triples of the outputs this config has + TMA_atom_and_partitions = [] for atom_out, dst_out, sO_out in ( - (tma_atom_row_act, tma_dst_row_act, sO_row_act), - (tma_atom_row_gate, tma_dst_row_gate, sO_row_gate), - (tma_atom_col_act, tma_dst_col_act, sO_col_act), - (tma_atom_col_gate, tma_dst_col_gate, sO_col_gate), + (tma_atom_row, tma_dst_row, sO_row), # FWD only + (tma_atom_col, tma_dst_col, sO_col), # FWD only + (tma_atom_row_dact, tma_dst_row_dact, sO_row_dact), # BWD only + (tma_atom_row_dgate, tma_dst_row_dgate, sO_row_dgate), # BWD only + (tma_atom_col_dact, tma_dst_col_dact, sO_col_dact), # BWD only + (tma_atom_col_dgate, tma_dst_col_dgate, sO_col_dgate), # BWD only ): if cutlass.const_expr(atom_out is not None): tXsO_out, tXgO_out = cute.nvgpu.cpasync.tma_partition( @@ -657,87 +661,58 @@ class BarrierStorage: # pylint: disable=missing-class-docstring 0, cute.make_layout(1), sO_out, - cute.zipped_divide(dst_out, cta_tiler), + cute.zipped_divide(cute.local_tile(dst_out, cta_block, cta_coord), tile_shape), ) - outputs.append((atom_out, tXsO_out, tXgO_out)) + TMA_atom_and_partitions.append((atom_out, tXsO_out, tXgO_out)) # Barrier init must be visible to every thread cute.arch.sync_threads() # Prologue: fill the pipeline if warp_idx == 0: - for s in cutlass.range_constexpr(min(NUM_STAGES, NUM_TILES)): + for s in cutlass.range_constexpr(NUM_STAGES): mainloop_pipeline.producer_acquire(prod_state) - self._copy_inputs_to_stage( - bidy * NUM_TILES + s, - bidx, - prod_state.index, - mainloop_pipeline.producer_get_barrier(prod_state), - tma_atom_grad, - tXgGrad, - tXsGrad, - tma_atom_act, - tXgAct, - tXsAct, - tma_atom_gate, - tXgGate, - tXsGate, - ) + for load_atom, load_sX, load_gX in inputs: + cute.copy( + load_atom, + load_gX[(None, (s, 0))], + load_sX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) mainloop_pipeline.producer_commit(prod_state) prod_state.advance() for tile_idx in cutlass.range(NUM_TILES, unroll=1): mainloop_pipeline.consumer_wait(cons_state) - # This tile reuses the output buffers of the tile NUM_STAGES ago; wait until TMA has - # read them. + # We only reuse previously used SMEM buffers when we have more tiles to process than pipeline stages, and only + # after we've already finished one round and wrap around to the start of the ring buffer if cutlass.const_expr(NUM_TILES > NUM_STAGES): - if warp_idx == 0: - cute.arch.cp_async_bulk_wait_group(NUM_STAGES - 1, read=True) - cute.arch.sync_threads() + if tile_idx >= NUM_STAGES: + # Only allow at most NUM_STAGES - 1 store groups reading SMEM on the fly, so the tile NUM_STAGES + # ago must finish before we reuse its SMEM buffer + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(NUM_STAGES - 1, read=True) + cute.arch.sync_threads() + # The tile should be ready to use now + stage_idx = cons_state.index - tile_idx_y = bidy * NUM_TILES + tile_idx - scale_coord = (None, (tile_idx_y, bidx)) + tile_coord = (tile_idx, 0) # within this CTA's block + # Select current stage's SMEM tile + sAct_tile = sAct[(None, stage_idx)] + sGate_tile = sGate[(None, stage_idx)] sGrad_tile = None if cutlass.const_expr(cfg.IS_BWD): sGrad_tile = sGrad[(None, stage_idx)] - sO_row_act_tile = None - mS_row_act_stage = None - sO_row_gate_tile = None - mS_row_gate_stage = None - if cutlass.const_expr(cfg.ROWWISE): - sO_row_act_tile = sO_row_act[(None, stage_idx)] - mS_row_act_stage = cute.flatten(mS_row_act[scale_coord]) - if cutlass.const_expr(cfg.IS_BWD): - sO_row_gate_tile = sO_row_gate[(None, stage_idx)] - mS_row_gate_stage = cute.flatten(mS_row_gate[scale_coord]) - sO_col_act_tile = None - mS_col_act_stage = None - sO_col_gate_tile = None - mS_col_gate_stage = None - if cutlass.const_expr(cfg.COLWISE): - sO_col_act_tile = sO_col_act[(None, stage_idx)] - mS_col_act_stage = cute.flatten(mS_col_act[scale_coord]) - if cutlass.const_expr(cfg.IS_BWD): - sO_col_gate_tile = sO_col_gate[(None, stage_idx)] - mS_col_gate_stage = cute.flatten(mS_col_gate[scale_coord]) - - quantize_gated_mxfp8_smem_tile( - sAct[(None, stage_idx)], - sGate[(None, stage_idx)], + + # Compute the gated op on this tile in fp32, cast back to the original dtype, + # and store the result in-place in the SMEM tile without quantization. + # FWD: write the output into sAct_tile + # BWD: write the act half's gradient to sAct_tile, and write the gate half's gradient to sGate_tile + gated_elementwise_smem_tile( + sAct_tile, # we also store the FWD output / BWD act half's gradient here + sGate_tile, # we also store the BWD gate half's gradient here sGrad_tile, - sO_row_act_tile, - mS_row_act_stage, - sO_row_gate_tile, - mS_row_gate_stage, - sO_col_act_tile, - mS_col_act_stage, - sO_col_gate_tile, - mS_col_gate_stage, - tile_idx_y * TILE_ROWS, - bidx * TILE_COLS, - M, - N, limit, alpha, glu_linear_offset, @@ -745,29 +720,66 @@ class BarrierStorage: # pylint: disable=missing-class-docstring beta2, ACTIVATION=cfg.GATED_ACTIVATION, IS_BWD=cfg.IS_BWD, - ROWWISE=cfg.ROWWISE, - COLWISE=cfg.COLWISE, - MAX_NORM_RCP=cfg.MAX_NORM_RCP, DTYPE=DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, - WITH_GEMM_SWIZZLED_SCALES=cfg.WITH_GEMM_SWIZZLED_SCALES, TILE_X=TILE_COLS, - TILE_Y=TILE_ROWS, ) + # The rowwise quantization pass reads elements other threads wrote so we need to sync here + if cutlass.const_expr(cfg.ROWWISE): + cute.arch.sync_threads() + + # Quantize every output of this tile into its SMEM buffers and scales + for sX_out, sO_row_out, mS_row_cta, sO_col_out, mS_col_cta in tile_quantizer_params: + # Slice the SMEM / GMEM tile for the current iteration + sO_row_tile = None + mS_row_stage = None + if cutlass.const_expr(cfg.ROWWISE): + sO_row_tile = sO_row_out[(None, stage_idx)] + mS_row_stage = cute.flatten( + cute.local_tile(mS_row_cta, row_scale_tile, tile_coord) + ) + sO_col_tile = None + mS_col_stage = None + if cutlass.const_expr(cfg.COLWISE): + sO_col_tile = sO_col_out[(None, stage_idx)] + mS_col_stage = cute.flatten( + cute.local_tile(mS_col_cta, col_scale_tile, tile_coord) + ) + # Use the general MXFP8 SMEM tile quantize function to quantize it + quantize_mxfp8_smem_tile_general_partial( + sX_out[(None, stage_idx)], # The SMEM tile to quantize + sO_row_tile, # SMEM tile for rowwise quantized output + mS_row_stage, # GMEM tile for rowwise quantized scales + sO_col_tile, # SMEM tile for columnwise quantized output + mS_col_stage, # GMEM tile for columnwise quantized scales + # The tile's first row and column in the whole tensor, for bounds checks + (bidy * NUM_TILES + tile_idx) * TILE_ROWS, + bidx * TILE_COLS, + M, + N, + ROWWISE=cfg.ROWWISE, + COLWISE=cfg.COLWISE, + MAX_NORM_RCP=cfg.MAX_NORM_RCP, + DTYPE=DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=cfg.WITH_GEMM_SWIZZLED_SCALES, + TILE_X=TILE_COLS, + TILE_Y=TILE_ROWS, + ) + # Make the SMEM writes visible to the TMA's async proxy before it reads them cute.arch.fence_proxy("async.shared", space="cta") cute.arch.sync_threads() mainloop_pipeline.consumer_release(cons_state) if warp_idx == 0: - # Output tensors are only read here, inside this dynamic branch, so these loop - # targets are fresh names rather than rebindings of values defined outside it. - for store_atom, store_sO, store_gO in outputs: + # The tile quantize function has already written the quantized outputs to SMEM, + # so we can issue TMA stores from SMEM to GMEM + for store_atom, store_sO, store_gO in TMA_atom_and_partitions: cute.copy( store_atom, store_sO[(None, stage_idx)], - store_gO[(None, (tile_idx_y, bidx))], + store_gO[(None, (tile_idx, 0))], ) cute.arch.cp_async_bulk_commit_group() @@ -778,21 +790,13 @@ class BarrierStorage: # pylint: disable=missing-class-docstring next_tile_idx = tile_idx + NUM_STAGES if next_tile_idx < NUM_TILES: mainloop_pipeline.producer_acquire(prod_state) - self._copy_inputs_to_stage( - bidy * NUM_TILES + next_tile_idx, - bidx, - prod_state.index, - mainloop_pipeline.producer_get_barrier(prod_state), - tma_atom_grad, - tXgGrad, - tXsGrad, - tma_atom_act, - tXgAct, - tXsAct, - tma_atom_gate, - tXgGate, - tXsGate, - ) + for load_atom, load_sX, load_gX in inputs: + cute.copy( + load_atom, + load_gX[(None, (next_tile_idx, 0))], + load_sX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) mainloop_pipeline.producer_commit(prod_state) prod_state.advance() From 4687cc186c8ec51ae0699ae84a3558679ac93c43 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Sat, 31 Oct 2026 01:17:35 +0000 Subject: [PATCH 37/43] nit Signed-off-by: Kaining Zhong --- .../common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py index 3035967c217..e781b6350b2 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py @@ -602,12 +602,10 @@ class BarrierStorage: # pylint: disable=missing-class-docstring # The block processed by a CTA, consisting of NUM_TILES vertically stacked tiles cta_block = (NUM_TILES * TILE_ROWS, TILE_COLS) - # The scales block processed by a CTA, consisting of NUM_TILES vertically stacked scale tiles - cta_block_row_scales = (NUM_TILES * TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) - cta_block_col_scales = (NUM_TILES * TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, TILE_COLS) - # Prepare parameters to pass to the SMEM tile quantize function by dact / dgate / output, where # for each group, input, quantized output are SMEM tiles, and scales are GMEM tiles + cta_block_row_scales = (NUM_TILES * TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + cta_block_col_scales = (NUM_TILES * TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, TILE_COLS) if cutlass.const_expr(cfg.IS_BWD): if cutlass.const_expr(cfg.ROWWISE): mS_row_dact = cute.local_tile(mS_row_dact, cta_block_row_scales, cta_coord) From 1b6ba658d3e2b59fd9ebee37c981371df6ab0823 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 2 Nov 2026 23:18:27 +0000 Subject: [PATCH 38/43] rename & refactor Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_common.py | 2 +- .../cast/mxfp8/quantize_mxfp8_gated_g2s.py | 365 ++++++++++-------- 2 files changed, 199 insertions(+), 168 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 6cc00b8bab0..8395063ca6a 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -436,7 +436,7 @@ def compile(self): return cute.compile( MXFP8GatedQuantizeG2SKernel(self), _fake_gmem_tensor(self.DTYPE, (sym_M, sym_2N), 16), # mGatedInput - _fake_gmem_tensor(self.DTYPE, (sym_M, sym_N), 16) if self.IS_BWD else None, # mGrad + _fake_gmem_tensor(self.DTYPE, (sym_M, sym_N), 16) if self.IS_BWD else None, # mYGrad # mO_row, mS_row _fake_gmem_tensor(self.FP8_DTYPE, (sym_M, out_cols), 16) if self.ROWWISE else None, ( diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py index e781b6350b2..3819fa368dd 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py @@ -6,15 +6,16 @@ The kernel reuses MXFP8QuantizeGeneralG2SKernel's TMA pipeline and its plain-cast SMEM tile quantizer. What it adds is an elementwise pass that runs on each staged tile before it is -quantized: it combines the act and gate halves of the input (and the incoming gradient, in -backward) and writes the result back into the act (and gate) input buffer in the input dtype. +quantized: it combines the act and gate halves of the input (and the incoming gradient y_grad, +in backward) and writes the result back into the act (and gate) input buffer in the input dtype. That is exactly the truncation the CUDA kernel applies before computing block amaxes, so the plain cast that follows reproduces its output. Gated input is [M, 2N] with [act | gate] along the last dim. Every operand is addressed through (M, N) views (stride 2N, column offset 0 or N), mirroring the CUDA kernel's tensor maps: - forward : out[M, N] = act(x[:, :N]) * gate(x[:, N:]) - backward: out[M, 2N] = [d/d act | d/d gate] given grad[M, N] + forward : y[M, N] = act(x[:, :N]) * gate(x[:, N:]) + backward: [act_grad | gate_grad][M, 2N], the gradients of the act and gate halves given + y_grad[M, N] """ from typing import Optional @@ -62,9 +63,9 @@ def _select_f32(pred: Boolean, if_true: Float32, if_false: Float32) -> Float32: @cute.jit def _gated_element( - act_elt: Float32, - gate_elt: Float32, - grad_elt: Optional[Float32], + act_element: Float32, + gate_element: Float32, + y_grad_element: Optional[Float32], limit: Optional[Float32], alpha: Optional[Float32], glu_linear_offset: Optional[Float32], @@ -73,63 +74,64 @@ def _gated_element( ACTIVATION: cutlass.Constexpr[str], IS_BWD: cutlass.Constexpr[bool], ): - """One element of the gated op in FP32; returns (act output, gate output or None). - - Mirrors the elementwise math of gated_mxfp8.cuh's columnwise pass (forward and backward), - operation by operation. """ - out_gate = None + FWD: (y, None) y = act(a) * gate(g) + BWD: (act_grad, gate_grad) act_grad = y_grad * gate(g) * dact(a), + gate_grad = y_grad * act(a) * dgate(g) + """ # Transform the gate. Gating is ideally the identity, with derivative 1. - gate = gate_elt + gate = gate_element dgate = None if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): # |g| > limit is gated_mxfp8.cuh's (g > limit || g < -limit), NaN included - dgate = _select_f32(fabs_f32(gate_elt) > limit, Float32(0.0), Float32(1.0)) - gate = cute.arch.fmin(cute.arch.fmax(-limit, gate_elt), limit) + glu_linear_offset + dgate = _select_f32(fabs_f32(gate_element) > limit, Float32(0.0), Float32(1.0)) + gate = cute.arch.fmin(cute.arch.fmax(-limit, gate_element), limit) + glu_linear_offset elif cutlass.const_expr(ACTIVATION == "situglu"): - dgate = dsitu_up(gate_elt, beta2) - gate = situ_up(gate_elt, beta2) + dgate = dsitu_up(gate_element, beta2) + gate = situ_up(gate_element, beta2) if cutlass.const_expr(not IS_BWD): if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): - act = act_clamped_silu(act_elt, limit, alpha) + act = act_clamped_silu(act_element, limit, alpha) elif cutlass.const_expr(ACTIVATION == "situglu"): - act = act_situ_gate(act_elt, beta1) + act = act_situ_gate(act_element, beta1) else: - act = GATED_ACTIVATIONS[ACTIVATION][0](act_elt) - out_act = act * gate + act = GATED_ACTIVATIONS[ACTIVATION][0](act_element) + y = act * gate + outputs = (y, None) else: if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): - x = cute.arch.fmin(act_elt, limit) + x = cute.arch.fmin(act_element, limit) s = sigmoidf(alpha * x) - act_x = x * s - dact_x = _select_f32( - act_elt <= limit, s + s * (Float32(1.0) - s) * alpha * x, Float32(0.0) + act = x * s + dact = _select_f32( + act_element <= limit, s + s * (Float32(1.0) - s) * alpha * x, Float32(0.0) ) elif cutlass.const_expr(ACTIVATION == "situglu"): - act_x = act_situ_gate(act_elt, beta1) - dact_x = dact_dsitu_gate(act_elt, beta1) + act = act_situ_gate(act_element, beta1) + dact = dact_dsitu_gate(act_element, beta1) elif cutlass.const_expr(ACTIVATION == "swiglu"): - s = sigmoidf(act_elt) - act_x = act_elt * s - dact_x = act_elt * s * (Float32(1.0) - s) + s + s = sigmoidf(act_element) + act = act_element * s + dact = act_element * s * (Float32(1.0) - s) + s else: act_fn, dact_fn = GATED_ACTIVATIONS[ACTIVATION] - act_x = act_fn(act_elt) - dact_x = dact_fn(act_elt) - out_act = dact_x * grad_elt * gate + act = act_fn(act_element) + dact = dact_fn(act_element) + act_grad = dact * y_grad_element * gate if cutlass.const_expr(dgate is None): - out_gate = act_x * grad_elt + gate_grad = act * y_grad_element else: - out_gate = dgate * act_x * grad_elt - return out_act, out_gate + gate_grad = dgate * act * y_grad_element + outputs = (act_grad, gate_grad) + return outputs @cute.jit def gated_elementwise_smem_tile( - sAct_tile: cute.Tensor, - sGate_tile: cute.Tensor, - sGrad_tile: Optional[cute.Tensor], + sAct_tile: cute.Tensor, # also the y tile in FWD / the act_grad tile in BWD + sGate_tile: cute.Tensor, # also the gate_grad tile in BWD + sYGrad_tile: Optional[cute.Tensor], limit: Optional[Float32], alpha: Optional[Float32], glu_linear_offset: Optional[Float32], @@ -137,52 +139,70 @@ def gated_elementwise_smem_tile( beta2: Optional[Float32], ACTIVATION: cutlass.Constexpr[str], IS_BWD: cutlass.Constexpr[bool], + ROWWISE: cutlass.Constexpr[bool], DTYPE: cutlass.Constexpr[type], TILE_X: cutlass.Constexpr[int], ): - """Compute the gated op on one staged tile, in place, rounded to the input dtype. - - Thread tidx owns column tidx of the tile, which is also the element set it owns in the - colwise quantization pass, so a colwise-only consumer needs no barrier after this. - """ + """Compute the gated op on one staged tile, in place, rounded to the input dtype.""" tidx, _, _ = cute.arch.thread_idx() _, 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)), ) - sAct_thread = cute.composition(sAct_tile, tv_layout)[tidx, None] - sGate_thread = cute.composition(sGate_tile, tv_layout)[tidx, None] + + sAct_tile = cute.composition(sAct_tile, tv_layout)[tidx, None] + sGate_tile = cute.composition(sGate_tile, tv_layout)[tidx, None] if cutlass.const_expr(IS_BWD): - sGrad_thread = cute.composition(sGrad_tile, tv_layout)[tidx, None] + sYGrad_tile = cute.composition(sYGrad_tile, tv_layout)[tidx, None] + if cutlass.const_expr(is_packed16(DTYPE)): truncate_f32 = truncate_f32_f16 if DTYPE is cutlass.Float16 else truncate_f32_bf16 for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - grad_elt = None - if cutlass.const_expr(IS_BWD): - grad_elt = Float32(sGrad_thread[i]) - out_act, out_gate = _gated_element( - Float32(sAct_thread[i]), - Float32(sGate_thread[i]), - grad_elt, - limit, - alpha, - glu_linear_offset, - beta1, - beta2, - ACTIVATION=ACTIVATION, - IS_BWD=IS_BWD, - ) - # Numerical truncation: round to the input dtype, like the CUDA kernels do before - # taking the amax. The cast below is then exact. - if cutlass.const_expr(is_packed16(DTYPE)): - out_act = truncate_f32(out_act) - if cutlass.const_expr(IS_BWD): - out_gate = truncate_f32(out_gate) - sAct_thread[i] = DTYPE(out_act) + act_element = Float32(sAct_tile[i]) + gate_element = Float32(sGate_tile[i]) + # Round to the input dtype, like the CUDA kernels do before taking the amax. The cast to + # DTYPE below is then exact. if cutlass.const_expr(IS_BWD): - sGate_thread[i] = DTYPE(out_gate) + y_grad_element = Float32(sYGrad_tile[i]) + act_grad_element, gate_grad_element = _gated_element( + act_element, + gate_element, + y_grad_element, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=ACTIVATION, + IS_BWD=IS_BWD, + ) + if cutlass.const_expr(is_packed16(DTYPE)): + act_grad_element = truncate_f32(act_grad_element) + gate_grad_element = truncate_f32(gate_grad_element) + sAct_tile[i] = DTYPE(act_grad_element) + sGate_tile[i] = DTYPE(gate_grad_element) + else: + y_element, _ = _gated_element( + act_element, + gate_element, + None, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=ACTIVATION, + IS_BWD=IS_BWD, + ) + if cutlass.const_expr(is_packed16(DTYPE)): + y_element = truncate_f32(y_element) + sAct_tile[i] = DTYPE(y_element) + + # The rowwise quantization pass reads elements other threads wrote so we need to sync here + if cutlass.const_expr(ROWWISE): + cute.arch.sync_threads() @cute.jit @@ -253,7 +273,7 @@ def _split_gated_scales( """Act and gate views of the backward output's scales; returns (act_row, gate_row, act_col, gate_col). - The [M, 2N] output is [d act | d gate] and its scales are laid out for the whole output. Each + The [M, 2N] output is [act_grad | gate_grad] and its scales are laid out for the whole output. Each view has the shape of an [M, N] tensor's scales and the strides of the whole output's, and starts where the whole output's layout places that half's first scale. For swizzled scales this needs the halves to be whole 128x4 tiles (N % 128 == 0, required by the C++ dispatcher). @@ -293,7 +313,7 @@ def _split_gated_scales( class MXFP8GatedQuantizeG2SKernel(MXFP8QuantizeGeneralG2SKernel): """gated_mxfp8.cuh in CuTeDSL. The general kernel's pipeline, loading the act and gate halves - (and grad, backward) each stage and storing one (forward) or two (backward) quantized halves. + (and y_grad in BWD) each stage and storing one (forward) or two (backward) quantized halves. Inherits the plain-cast tile config, which is gated_mxfp8.cuh's: 64 threads per CTA, each CTA covering two vertically stacked 32x64 tiles. @@ -313,7 +333,7 @@ def __init__(self, cfg: MXFP8GatedQuantizeConfig): def __call__( self, mGatedInput: cute.Tensor, # [M, 2N] = [act | gate] - mGrad: Optional[cute.Tensor], # [M, N], bwd only + mYGrad: Optional[cute.Tensor], # [M, N], bwd only # Outputs, with C = N (FWD) or 2N (BWD) mO_row: Optional[cute.Tensor], # [M, C] mS_row: Optional[cute.Tensor], # [roundup(M, 128), roundup(C / 32, 4)] @@ -329,11 +349,11 @@ def __call__( ): """ mGatedInput: [a | g] with shape M x 2N, where a is the activation half and g is the gate half. - mGrad: None for FWD, [dy] with shape [M, N] for BWD + mYGrad: None for FWD, y_grad with shape [M, N] for BWD FWD mode: - out = quantize(act(a) * gate(g)) with shape [M, N] + y = quantize(act(a) * gate(g)) with shape [M, N] BWD mode: - out = quantize([ dy * gate(g) * dact(a) | dy * act(a) * dgate(g) ]) with shape [M, 2N] + [act_grad | gate_grad] = quantize([ y_grad * gate(g) * dact(a) | y_grad * act(a) * dgate(g) ]) with shape [M, 2N] """ if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): cute.printf(f"[CuTeDSL] MXFP8GatedQuantizeG2SKernel.__call__() with config: {self.cfg}\n") @@ -341,33 +361,35 @@ def __call__( M = mGatedInput.shape[0] # Pick any mO that is used just to get N mO = mO_row if mO_row is not None else mO_col - # FWD: mO is [M, N]; BWD: mGrad is [M, N] - N = mO.shape[1] if not self.cfg.IS_BWD else mGrad.shape[1] + # FWD: mO is [M, N]; BWD: mYGrad is [M, N] + N = mO.shape[1] if not self.cfg.IS_BWD else mYGrad.shape[1] # The act and gate halves of the [M, 2N] input mAct = cute.local_tile(mGatedInput, (M, N), (0, 0)) mGate = cute.local_tile(mGatedInput, (M, N), (0, 1)) - mO_row_dact = mO_row_dgate = mO_col_dact = mO_col_dgate = None # BWD only - mS_row_dact = mS_row_dgate = mS_col_dact = mS_col_dgate = None # BWD only + mO_row_act_grad = mO_row_gate_grad = mO_col_act_grad = mO_col_gate_grad = None # BWD only + mS_row_act_grad = mS_row_gate_grad = mS_col_act_grad = mS_col_gate_grad = None # BWD only if cutlass.const_expr(self.cfg.IS_BWD): - # The [M, 2N] output is [d act | d gate]; the kernel takes its two halves + # The [M, 2N] output is [act_grad | gate_grad]; the kernel takes its two halves # TODO(kainingz): how to implement this in a more elegant way - mS_row_dact, mS_row_dgate, mS_col_dact, mS_col_dgate = _split_gated_scales( - mS_row, - mS_col, - M, - N, - self.cfg.ROWWISE, - self.cfg.COLWISE, - self.cfg.WITH_GEMM_SWIZZLED_SCALES, + mS_row_act_grad, mS_row_gate_grad, mS_col_act_grad, mS_col_gate_grad = ( + _split_gated_scales( + mS_row, + mS_col, + M, + N, + self.cfg.ROWWISE, + self.cfg.COLWISE, + self.cfg.WITH_GEMM_SWIZZLED_SCALES, + ) ) if cutlass.const_expr(self.cfg.ROWWISE): - mO_row_dact = cute.local_tile(mO_row, (M, N), (0, 0)) - mO_row_dgate = cute.local_tile(mO_row, (M, N), (0, 1)) + mO_row_act_grad = cute.local_tile(mO_row, (M, N), (0, 0)) + mO_row_gate_grad = cute.local_tile(mO_row, (M, N), (0, 1)) if cutlass.const_expr(self.cfg.COLWISE): - mO_col_dact = cute.local_tile(mO_col, (M, N), (0, 0)) - mO_col_dgate = cute.local_tile(mO_col, (M, N), (0, 1)) + mO_col_act_grad = cute.local_tile(mO_col, (M, N), (0, 0)) + mO_col_gate_grad = cute.local_tile(mO_col, (M, N), (0, 1)) # In BWD mode we won't use these, and we should set them to None so we don't build TMA atoms for them mO_row = mO_col = mS_row = mS_col = None elif cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): @@ -376,15 +398,15 @@ def __call__( ) atoms = self._make_tma_atoms( - mGrad, + mYGrad, mAct, mGate, mO_row, mO_col, - mO_row_dact, - mO_row_dgate, - mO_col_dact, - mO_col_dgate, + mO_row_act_grad, + mO_row_gate_grad, + mO_col_act_grad, + mO_col_gate_grad, ) # The TMA atoms are spread from a list, which pylint cannot count self.kernel( # pylint: disable=no-value-for-parameter @@ -392,10 +414,10 @@ def __call__( N, mS_row, mS_col, - mS_row_dact, - mS_row_dgate, - mS_col_dact, - mS_col_dgate, + mS_row_act_grad, + mS_row_gate_grad, + mS_col_act_grad, + mS_col_gate_grad, limit, alpha, glu_linear_offset, @@ -415,15 +437,15 @@ def __call__( @cute.jit def _make_tma_atoms( self, - mGrad: Optional[cute.Tensor], + mYGrad: Optional[cute.Tensor], mAct: cute.Tensor, mGate: cute.Tensor, mO_row: Optional[cute.Tensor], # FWD only mO_col: Optional[cute.Tensor], # FWD only - mO_row_dact: Optional[cute.Tensor], # BWD only - mO_row_dgate: Optional[cute.Tensor], # BWD only - mO_col_dact: Optional[cute.Tensor], # BWD only - mO_col_dgate: Optional[cute.Tensor], # BWD only + mO_row_act_grad: Optional[cute.Tensor], # BWD only + mO_row_gate_grad: Optional[cute.Tensor], # BWD only + mO_col_act_grad: Optional[cute.Tensor], # BWD only + mO_col_gate_grad: Optional[cute.Tensor], # BWD only ): """Build TMA atoms for input / quantized output; scales will not be loaded / stored by TMA.""" smem_tile_layout = cute.make_ordered_layout( @@ -435,15 +457,15 @@ def _make_tma_atoms( atoms = [] for op, mT in ( - (op_load, mGrad), + (op_load, mYGrad), (op_load, mAct), (op_load, mGate), (op_store, mO_row), (op_store, mO_col), - (op_store, mO_row_dact), - (op_store, mO_row_dgate), - (op_store, mO_col_dact), - (op_store, mO_col_dgate), + (op_store, mO_row_act_grad), + (op_store, mO_row_gate_grad), + (op_store, mO_col_act_grad), + (op_store, mO_col_gate_grad), ): if cutlass.const_expr(mT is None): atoms.extend([None, None]) @@ -462,18 +484,18 @@ def kernel( N: Int32, mS_row: Optional[cute.Tensor], # FWD only mS_col: Optional[cute.Tensor], # FWD only - mS_row_dact: Optional[cute.Tensor], # BWD only - mS_row_dgate: Optional[cute.Tensor], # BWD only - mS_col_dact: Optional[cute.Tensor], # BWD only - mS_col_dgate: Optional[cute.Tensor], # BWD only + mS_row_act_grad: Optional[cute.Tensor], # BWD only + mS_row_gate_grad: Optional[cute.Tensor], # BWD only + mS_col_act_grad: Optional[cute.Tensor], # BWD only + mS_col_gate_grad: Optional[cute.Tensor], # BWD only limit: Optional[Float32], alpha: Optional[Float32], glu_linear_offset: Optional[Float32], beta1: Optional[Float32], beta2: Optional[Float32], - # TMA for mGrad, only used in BWD - tma_atom_grad: Optional[cute.CopyAtom], - tma_src_grad: Optional[cute.Tensor], + # TMA for mYGrad, only used in BWD + tma_atom_y_grad: Optional[cute.CopyAtom], + tma_src_y_grad: Optional[cute.Tensor], # TMA for mAct and mGate, used in both FWD and BWD tma_atom_act: cute.CopyAtom, tma_src_act: cute.Tensor, @@ -484,15 +506,15 @@ def kernel( tma_dst_row: Optional[cute.Tensor], tma_atom_col: Optional[cute.CopyAtom], tma_dst_col: Optional[cute.Tensor], - # TMA for mO_row/col_dact/dgate, only used in BWD - tma_atom_row_dact: Optional[cute.CopyAtom], - tma_dst_row_dact: Optional[cute.Tensor], - tma_atom_row_dgate: Optional[cute.CopyAtom], - tma_dst_row_dgate: Optional[cute.Tensor], - tma_atom_col_dact: Optional[cute.CopyAtom], - tma_dst_col_dact: Optional[cute.Tensor], - tma_atom_col_dgate: Optional[cute.CopyAtom], - tma_dst_col_dgate: Optional[cute.Tensor], + # TMA for mO_row/col_act_grad/gate_grad, only used in BWD + tma_atom_row_act_grad: Optional[cute.CopyAtom], + tma_dst_row_act_grad: Optional[cute.Tensor], + tma_atom_row_gate_grad: Optional[cute.CopyAtom], + tma_dst_row_gate_grad: Optional[cute.Tensor], + tma_atom_col_act_grad: Optional[cute.CopyAtom], + tma_dst_col_act_grad: Optional[cute.Tensor], + tma_atom_col_gate_grad: Optional[cute.CopyAtom], + tma_dst_col_gate_grad: Optional[cute.Tensor], ): """ FWD mode: @@ -502,25 +524,25 @@ def kernel( mS_row/col: output scales for row / column directions BWD mode: - mGrad: incoming gradient + mYGrad: y_grad, the incoming gradient mAct: activation half input from FWD pass mGate: gate half input from FWD pass - mO_row/col_dact: quantized activation gradient for row / column directions - mO_row/col_dgate: quantized gate gradient for row / column directions - mS_row/col_dact: output scales for activation gradient row / column directions - mS_row/col_dgate: output scales for gate gradient row / column directions + mO_row/col_act_grad: quantized activation gradient for row / column directions + mO_row/col_gate_grad: quantized gate gradient for row / column directions + mS_row/col_act_grad: output scales for activation gradient row / column directions + mS_row/col_gate_grad: output scales for gate gradient row / column directions mO_row/col: quantized gradient (both activation and gate) for row / column directions mS_row/col: output scales for gradient (both activation and gate) row / column directions mAct: both FWD and BWD mGate: both FWD and BWD - mGrad: BWD only + mYGrad: BWD only mO_row/col: both FWD (N) and BWD (2N) mS_row/col: both FWD (N) and BWD (2N) - mO_row/col_dact: BWD only - mO_row/col_dgate: BWD only - mS_row/col_dact: BWD only - mS_row/col_dgate: BWD only + mO_row/col_act_grad: BWD only + mO_row/col_gate_grad: BWD only + mS_row/col_act_grad: BWD only + mS_row/col_gate_grad: BWD only """ cfg = self.cfg @@ -550,23 +572,31 @@ class BarrierStorage: # pylint: disable=missing-class-docstring # Allocate SMEM buffers for inputs sAct = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) sGate = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) - sGrad = None + sYGrad = None if cutlass.const_expr(cfg.IS_BWD): - sGrad = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) + sYGrad = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) # Allocate SMEM buffers for outputs sO_row = sO_col = None - sO_row_dact = sO_row_dgate = sO_col_dact = sO_col_dgate = None + sO_row_act_grad = sO_row_gate_grad = sO_col_act_grad = sO_col_gate_grad = None if cutlass.const_expr(cfg.ROWWISE): if cutlass.const_expr(cfg.IS_BWD): - sO_row_dact = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) - sO_row_dgate = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + sO_row_act_grad = smem.allocate_tensor( + cfg.FP8_DTYPE, stages_layout, byte_alignment=128 + ) + sO_row_gate_grad = smem.allocate_tensor( + cfg.FP8_DTYPE, stages_layout, byte_alignment=128 + ) else: sO_row = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) if cutlass.const_expr(cfg.COLWISE): if cutlass.const_expr(cfg.IS_BWD): - sO_col_dact = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) - sO_col_dgate = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + sO_col_act_grad = smem.allocate_tensor( + cfg.FP8_DTYPE, stages_layout, byte_alignment=128 + ) + sO_col_gate_grad = smem.allocate_tensor( + cfg.FP8_DTYPE, stages_layout, byte_alignment=128 + ) else: sO_col = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) @@ -575,14 +605,14 @@ class BarrierStorage: # pylint: disable=missing-class-docstring # Issue TMA descriptor prefetch using the TMA warp (warp 0) if warp_idx == 0: - for load_atom in (tma_atom_act, tma_atom_gate, tma_atom_grad): + for load_atom in (tma_atom_act, tma_atom_gate, tma_atom_y_grad): if cutlass.const_expr(load_atom is not None): cute.nvgpu.cpasync.prefetch_descriptor(load_atom) # TMA warp (warp 0)'s elected thread produces (issues TMA), every warp consumes producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._NUM_WARPS) - # Input has 2 tiles (act + gate) in FWD, 3 tiles (grad + act + gate) in BWD + # Input has 2 tiles (act + gate) in FWD, 3 tiles (y_grad + act + gate) in BWD num_inputs = 3 if cfg.IS_BWD else 2 mainloop_pipeline = pipeline.PipelineTmaAsync.create( barrier_storage=storage.mbar_storage.data_ptr(), @@ -602,21 +632,25 @@ class BarrierStorage: # pylint: disable=missing-class-docstring # The block processed by a CTA, consisting of NUM_TILES vertically stacked tiles cta_block = (NUM_TILES * TILE_ROWS, TILE_COLS) - # Prepare parameters to pass to the SMEM tile quantize function by dact / dgate / output, where + # Prepare parameters to pass to the SMEM tile quantize function by output (FWD: y; BWD: act_grad, gate_grad), where # for each group, input, quantized output are SMEM tiles, and scales are GMEM tiles cta_block_row_scales = (NUM_TILES * TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) cta_block_col_scales = (NUM_TILES * TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, TILE_COLS) if cutlass.const_expr(cfg.IS_BWD): if cutlass.const_expr(cfg.ROWWISE): - mS_row_dact = cute.local_tile(mS_row_dact, cta_block_row_scales, cta_coord) - mS_row_dgate = cute.local_tile(mS_row_dgate, cta_block_row_scales, cta_coord) + mS_row_act_grad = cute.local_tile(mS_row_act_grad, cta_block_row_scales, cta_coord) + mS_row_gate_grad = cute.local_tile( + mS_row_gate_grad, cta_block_row_scales, cta_coord + ) if cutlass.const_expr(cfg.COLWISE): - mS_col_dact = cute.local_tile(mS_col_dact, cta_block_col_scales, cta_coord) - mS_col_dgate = cute.local_tile(mS_col_dgate, cta_block_col_scales, cta_coord) + mS_col_act_grad = cute.local_tile(mS_col_act_grad, cta_block_col_scales, cta_coord) + mS_col_gate_grad = cute.local_tile( + mS_col_gate_grad, cta_block_col_scales, cta_coord + ) # In BWD mode we output gradient for both the act and gate halves, which are stored in sAct and sGate respectively tile_quantizer_params = [ - (sAct, sO_row_dact, mS_row_dact, sO_col_dact, mS_col_dact), # We will write unquantized dact to sAct in-place - (sGate, sO_row_dgate, mS_row_dgate, sO_col_dgate, mS_col_dgate), # We will write unquantized dgate to sGate in-place + (sAct, sO_row_act_grad, mS_row_act_grad, sO_col_act_grad, mS_col_act_grad), # We will write unquantized act_grad to sAct in-place + (sGate, sO_row_gate_grad, mS_row_gate_grad, sO_col_gate_grad, mS_col_gate_grad), # We will write unquantized gate_grad to sGate in-place ] else: if cutlass.const_expr(cfg.ROWWISE): @@ -626,12 +660,12 @@ class BarrierStorage: # pylint: disable=missing-class-docstring # In FWD mode we only output the result after the gated op, which is stored in sAct tile_quantizer_params = [(sAct, sO_row, mS_row, sO_col, mS_col)] # We will write unquantized output to sAct in-place - # Prepare input TMA (atom, smem partition, gmem partition) triples: act and gate, and BWD's grad + # Prepare input TMA (atom, smem partition, gmem partition) triples: act and gate, and BWD's y_grad inputs = [] for atom_in, src_in, sX_in in ( (tma_atom_act, tma_src_act, sAct), (tma_atom_gate, tma_src_gate, sGate), - (tma_atom_grad, tma_src_grad, sGrad), + (tma_atom_y_grad, tma_src_y_grad, sYGrad), ): if cutlass.const_expr(atom_in is not None): tXsX_in, tXgX_in = cute.nvgpu.cpasync.tma_partition( @@ -648,10 +682,10 @@ class BarrierStorage: # pylint: disable=missing-class-docstring for atom_out, dst_out, sO_out in ( (tma_atom_row, tma_dst_row, sO_row), # FWD only (tma_atom_col, tma_dst_col, sO_col), # FWD only - (tma_atom_row_dact, tma_dst_row_dact, sO_row_dact), # BWD only - (tma_atom_row_dgate, tma_dst_row_dgate, sO_row_dgate), # BWD only - (tma_atom_col_dact, tma_dst_col_dact, sO_col_dact), # BWD only - (tma_atom_col_dgate, tma_dst_col_dgate, sO_col_dgate), # BWD only + (tma_atom_row_act_grad, tma_dst_row_act_grad, sO_row_act_grad), # BWD only + (tma_atom_row_gate_grad, tma_dst_row_gate_grad, sO_row_gate_grad), # BWD only + (tma_atom_col_act_grad, tma_dst_col_act_grad, sO_col_act_grad), # BWD only + (tma_atom_col_gate_grad, tma_dst_col_gate_grad, sO_col_gate_grad), # BWD only ): if cutlass.const_expr(atom_out is not None): tXsO_out, tXgO_out = cute.nvgpu.cpasync.tma_partition( @@ -699,9 +733,9 @@ class BarrierStorage: # pylint: disable=missing-class-docstring # Select current stage's SMEM tile sAct_tile = sAct[(None, stage_idx)] sGate_tile = sGate[(None, stage_idx)] - sGrad_tile = None + sYGrad_tile = None if cutlass.const_expr(cfg.IS_BWD): - sGrad_tile = sGrad[(None, stage_idx)] + sYGrad_tile = sYGrad[(None, stage_idx)] # Compute the gated op on this tile in fp32, cast back to the original dtype, # and store the result in-place in the SMEM tile without quantization. @@ -710,7 +744,7 @@ class BarrierStorage: # pylint: disable=missing-class-docstring gated_elementwise_smem_tile( sAct_tile, # we also store the FWD output / BWD act half's gradient here sGate_tile, # we also store the BWD gate half's gradient here - sGrad_tile, + sYGrad_tile, limit, alpha, glu_linear_offset, @@ -718,14 +752,11 @@ class BarrierStorage: # pylint: disable=missing-class-docstring beta2, ACTIVATION=cfg.GATED_ACTIVATION, IS_BWD=cfg.IS_BWD, + ROWWISE=cfg.ROWWISE, DTYPE=DTYPE, TILE_X=TILE_COLS, ) - # The rowwise quantization pass reads elements other threads wrote so we need to sync here - if cutlass.const_expr(cfg.ROWWISE): - cute.arch.sync_threads() - # Quantize every output of this tile into its SMEM buffers and scales for sX_out, sO_row_out, mS_row_cta, sO_col_out, mS_col_cta in tile_quantizer_params: # Slice the SMEM / GMEM tile for the current iteration From 7875b37e07bde1273ebff214b85d42b367b181c8 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 2 Nov 2026 23:37:43 +0000 Subject: [PATCH 39/43] rename Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_gated_g2s.py | 69 ++++++++++--------- 1 file changed, 37 insertions(+), 32 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py index 3819fa368dd..70a1b65cf93 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py @@ -270,15 +270,16 @@ def _split_gated_scales( COLWISE: cutlass.Constexpr[bool], SWIZZLED: cutlass.Constexpr[bool], ): - """Act and gate views of the backward output's scales; returns - (act_row, gate_row, act_col, gate_col). - - The [M, 2N] output is [act_grad | gate_grad] and its scales are laid out for the whole output. Each - view has the shape of an [M, N] tensor's scales and the strides of the whole output's, and - starts where the whole output's layout places that half's first scale. For swizzled scales - this needs the halves to be whole 128x4 tiles (N % 128 == 0, required by the C++ dispatcher). + """Views of the backward output's scales split into its act_grad and gate_grad halves; returns + (mS_row_act, mS_row_gate, mS_col_act, mS_col_gate). + + The [M, 2N] output is [act_grad | gate_grad] and its scales are laid out for the whole + output. Each view has the shape of an [M, N] tensor's scales and the strides of the whole + output's, and starts where the whole output's layout places that half's first scale. For + swizzled scales this needs the halves to be whole 128x4 tiles (N % 128 == 0, required by the + C++ dispatcher). """ - act_row = gate_row = act_col = gate_col = None + mS_row_act = mS_row_gate = mS_col_act = mS_col_gate = None if cutlass.const_expr(SWIZZLED): mS_row, mS_col = derive_swizzled_scale_layout( M, 2 * N, ROWWISE, COLWISE, mS_row, mS_col @@ -289,26 +290,26 @@ def _split_gated_scales( cute.make_layout((32, 4, cute.ceil_div(M, 128))), # all rows cute.make_layout((4, half_tiles)), # one half's scale columns ) - act_row = cute.local_tile(mS_row, row_tiler, (0, 0)) - gate_row = cute.local_tile(mS_row, row_tiler, (0, 1)) + mS_row_act = cute.local_tile(mS_row, row_tiler, (0, 0)) + mS_row_gate = cute.local_tile(mS_row, row_tiler, (0, 1)) if cutlass.const_expr(COLWISE): col_tiler = ( cute.make_layout((4, cute.ceil_div(cute.ceil_div(M, 32), 4))), # all scale rows cute.make_layout((32, 4, half_tiles)), # one half's columns ) - act_col = cute.local_tile(mS_col, col_tiler, (0, 0)) - gate_col = cute.local_tile(mS_col, col_tiler, (0, 1)) + mS_col_act = cute.local_tile(mS_col, col_tiler, (0, 0)) + mS_col_gate = cute.local_tile(mS_col, col_tiler, (0, 1)) else: # Compact scales are row-major, so each half is a block of scale columns if cutlass.const_expr(ROWWISE): row_tiler = (mS_row.shape[0], N // MXFP8_BLOCK_SCALING_SIZE) - act_row = cute.local_tile(mS_row, row_tiler, (0, 0)) - gate_row = cute.local_tile(mS_row, row_tiler, (0, 1)) + mS_row_act = cute.local_tile(mS_row, row_tiler, (0, 0)) + mS_row_gate = cute.local_tile(mS_row, row_tiler, (0, 1)) if cutlass.const_expr(COLWISE): col_tiler = (mS_col.shape[0], N) - act_col = cute.local_tile(mS_col, col_tiler, (0, 0)) - gate_col = cute.local_tile(mS_col, col_tiler, (0, 1)) - return act_row, gate_row, act_col, gate_col + mS_col_act = cute.local_tile(mS_col, col_tiler, (0, 0)) + mS_col_gate = cute.local_tile(mS_col, col_tiler, (0, 1)) + return mS_row_act, mS_row_gate, mS_col_act, mS_col_gate class MXFP8GatedQuantizeG2SKernel(MXFP8QuantizeGeneralG2SKernel): @@ -333,7 +334,7 @@ def __init__(self, cfg: MXFP8GatedQuantizeConfig): def __call__( self, mGatedInput: cute.Tensor, # [M, 2N] = [act | gate] - mYGrad: Optional[cute.Tensor], # [M, N], bwd only + mYGrad: Optional[cute.Tensor], # [M, N], BWD only # Outputs, with C = N (FWD) or 2N (BWD) mO_row: Optional[cute.Tensor], # [M, C] mS_row: Optional[cute.Tensor], # [roundup(M, 128), roundup(C / 32, 4)] @@ -451,7 +452,7 @@ def _make_tma_atoms( smem_tile_layout = cute.make_ordered_layout( (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) ) - cta_tiler = (self._TILE_ROWS, self._TILE_COLS) + tiler = (self._TILE_ROWS, self._TILE_COLS) op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() @@ -472,7 +473,7 @@ def _make_tma_atoms( else: atoms.extend( cute.nvgpu.cpasync.make_tiled_tma_atom( - op, mT, smem_tile_layout, cta_tiler, num_multicast=1 + op, mT, smem_tile_layout, tiler, num_multicast=1 ) ) return atoms @@ -628,12 +629,16 @@ class BarrierStorage: # pylint: disable=missing-class-docstring cta_coord = (bidy, bidx) # Shape of each individual tile. A CTA processes one tile at a time - tile_shape = (TILE_ROWS, TILE_COLS) + tiler = (TILE_ROWS, TILE_COLS) # The block processed by a CTA, consisting of NUM_TILES vertically stacked tiles cta_block = (NUM_TILES * TILE_ROWS, TILE_COLS) - # Prepare parameters to pass to the SMEM tile quantize function by output (FWD: y; BWD: act_grad, gate_grad), where - # for each group, input, quantized output are SMEM tiles, and scales are GMEM tiles + # Prepare parameters to pass to the SMEM tile quantize function per output (FWD: y; BWD: + # act_grad, gate_grad), where input and quantized output are SMEM tiles and scales are GMEM + # tiles + # Prepare parameters to pass to the SMEM tile quantize function for all outputs + # (FWD: y; BWD: act_grad, gate_grad) + # Note: input, quantized output are SMEM tiles, and scales are GMEM tiles cta_block_row_scales = (NUM_TILES * TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) cta_block_col_scales = (NUM_TILES * TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, TILE_COLS) if cutlass.const_expr(cfg.IS_BWD): @@ -661,7 +666,7 @@ class BarrierStorage: # pylint: disable=missing-class-docstring tile_quantizer_params = [(sAct, sO_row, mS_row, sO_col, mS_col)] # We will write unquantized output to sAct in-place # Prepare input TMA (atom, smem partition, gmem partition) triples: act and gate, and BWD's y_grad - inputs = [] + TMA_load_atom_and_partitions = [] for atom_in, src_in, sX_in in ( (tma_atom_act, tma_src_act, sAct), (tma_atom_gate, tma_src_gate, sGate), @@ -673,12 +678,12 @@ class BarrierStorage: # pylint: disable=missing-class-docstring 0, cute.make_layout(1), sX_in, - cute.zipped_divide(cute.local_tile(src_in, cta_block, cta_coord), tile_shape), + cute.zipped_divide(cute.local_tile(src_in, cta_block, cta_coord), tiler), ) - inputs.append((atom_in, tXsX_in, tXgX_in)) + TMA_load_atom_and_partitions.append((atom_in, tXsX_in, tXgX_in)) # Prepare output TMA (atom, smem partition, gmem partition) triples of the outputs this config has - TMA_atom_and_partitions = [] + TMA_store_atom_and_partitions = [] for atom_out, dst_out, sO_out in ( (tma_atom_row, tma_dst_row, sO_row), # FWD only (tma_atom_col, tma_dst_col, sO_col), # FWD only @@ -693,9 +698,9 @@ class BarrierStorage: # pylint: disable=missing-class-docstring 0, cute.make_layout(1), sO_out, - cute.zipped_divide(cute.local_tile(dst_out, cta_block, cta_coord), tile_shape), + cute.zipped_divide(cute.local_tile(dst_out, cta_block, cta_coord), tiler), ) - TMA_atom_and_partitions.append((atom_out, tXsO_out, tXgO_out)) + TMA_store_atom_and_partitions.append((atom_out, tXsO_out, tXgO_out)) # Barrier init must be visible to every thread cute.arch.sync_threads() @@ -704,7 +709,7 @@ class BarrierStorage: # pylint: disable=missing-class-docstring if warp_idx == 0: for s in cutlass.range_constexpr(NUM_STAGES): mainloop_pipeline.producer_acquire(prod_state) - for load_atom, load_sX, load_gX in inputs: + for load_atom, load_sX, load_gX in TMA_load_atom_and_partitions: cute.copy( load_atom, load_gX[(None, (s, 0))], @@ -804,7 +809,7 @@ class BarrierStorage: # pylint: disable=missing-class-docstring if warp_idx == 0: # The tile quantize function has already written the quantized outputs to SMEM, # so we can issue TMA stores from SMEM to GMEM - for store_atom, store_sO, store_gO in TMA_atom_and_partitions: + for store_atom, store_sO, store_gO in TMA_store_atom_and_partitions: cute.copy( store_atom, store_sO[(None, stage_idx)], @@ -819,7 +824,7 @@ class BarrierStorage: # pylint: disable=missing-class-docstring next_tile_idx = tile_idx + NUM_STAGES if next_tile_idx < NUM_TILES: mainloop_pipeline.producer_acquire(prod_state) - for load_atom, load_sX, load_gX in inputs: + for load_atom, load_sX, load_gX in TMA_load_atom_and_partitions: cute.copy( load_atom, load_gX[(None, (next_tile_idx, 0))], From 9c8534a5dfb55ab4e74e0daf3e126b86852ccaf0 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 3 Nov 2026 00:58:17 +0000 Subject: [PATCH 40/43] refactor Signed-off-by: Kaining Zhong --- .../common/CuTeDSL/activations.py | 59 +++++++------ .../cast/mxfp8/quantize_mxfp8_common.py | 6 +- .../cast/mxfp8/quantize_mxfp8_gated_g2s.py | 84 ++++++++++--------- transformer_engine/common/CuTeDSL/utils.py | 11 ++- 4 files changed, 86 insertions(+), 74 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/activations.py b/transformer_engine/common/CuTeDSL/activations.py index 636dafe7f82..d65fbca6f49 100644 --- a/transformer_engine/common/CuTeDSL/activations.py +++ b/transformer_engine/common/CuTeDSL/activations.py @@ -8,6 +8,7 @@ # Each function is a 1:1 port of its math.h counterpart (see module docstring). # pylint: disable=missing-function-docstring import os +from typing import Optional from cutlass import cute from cutlass import Float32 @@ -21,8 +22,9 @@ def tanh(x: Float32) -> Float32: return cute.math.tanh(x, fastmath=USE_FAST_MATH) -def exp(x: Float32) -> Float32: - return cute.math.exp(x, fastmath=USE_FAST_MATH) +def exp(x: Float32, fast_math: bool = False) -> Float32: + """exp(x); fast_math, or NVTE_USE_FAST_MATH, selects the fast approximate exp.""" + return cute.math.exp(x, fastmath=fast_math or USE_FAST_MATH) def act_relu(x: Float32) -> Float32: @@ -35,8 +37,11 @@ def act_gelu(x: Float32) -> Float32: return x * (Float32(0.5) + Float32(0.5) * tanh(x * (A + B * x * x))) -def act_silu(x: Float32) -> Float32: - return x / (Float32(1.0) + exp(-x)) +def act_silu(x: Float32, sig: Optional[Float32] = None, fast_math: bool = False) -> Float32: + # Recompute the sigmoid if not provided one + if sig is None: + sig = sigmoid(x, fast_math) + return x * sig def act_qgelu(x: Float32) -> Float32: @@ -57,17 +62,20 @@ def dact_dsrelu(x: Float32) -> Float32: return cute.arch.fmax(Float32(2.0) * x, Float32(0.0)) -def sigmoid(x: Float32) -> Float32: - return Float32(1.0) / (Float32(1.0) + exp(-x)) +def sigmoid(x: Float32, fast_math: bool = False) -> Float32: + return Float32(1.0) / (Float32(1.0) + exp(-x, fast_math)) -def dsigmoid(x: Float32) -> Float32: - s = sigmoid(x) +def dsigmoid(x: Float32, fast_math: bool = False) -> Float32: + s = sigmoid(x, fast_math) return s * (Float32(1.0) - s) -def dact_dsilu(x: Float32) -> Float32: - return fma_f32(x, dsigmoid(x), sigmoid(x)) +def dact_dsilu(x: Float32, sig: Optional[Float32] = None, fast_math: bool = False) -> Float32: + # Recompute the sigmoid if not provided one + if sig is None: + return fma_f32(x, dsigmoid(x, fast_math), sigmoid(x, fast_math)) + return x * sig * (Float32(1.0) - sig) + sig def dact_dqgelu(x: Float32) -> Float32: @@ -86,8 +94,6 @@ def dact_dgelu(x: Float32) -> Float32: # The parameterized ones take their ClampedSwiGLUParam / SiTUGLUParam fields as # trailing runtime arguments. -LOG2E = 1.4426950408889634 # log2(e), what __expf scales its argument by - def act_sigmoid(x: Float32) -> Float32: return sigmoid(x) @@ -103,34 +109,33 @@ def act_clamped_silu(x: Float32, limit: Float32, alpha: Float32) -> Float32: def act_situ_gate(x: Float32, beta1: Float32) -> Float32: + """SiTU-GLU's activation, applied to the act half.""" return beta1 * tanh(x / beta1) * sigmoid(x) def dact_dsitu_gate(x: Float32, beta1: Float32) -> Float32: + """Derivative of act_situ_gate.""" t = tanh(x / beta1) s = sigmoid(x) return (Float32(1.0) - t * t) * s + beta1 * t * s * (Float32(1.0) - s) +def identity(x: Float32) -> Float32: + """The default gate: gating is ideally the identity.""" + return x + + +def didentity(x: Float32) -> Float32: # pylint: disable=unused-argument + """Derivative of identity.""" + return Float32(1.0) + + def situ_up(x: Float32, beta2: Float32) -> Float32: + """SiTU-GLU's gate transform, applied to the gate half.""" return beta2 * tanh(x / beta2) def dsitu_up(x: Float32, beta2: Float32) -> Float32: + """Derivative of situ_up.""" t = tanh(x / beta2) return Float32(1.0) - t * t - - -@cute.jit -def sigmoidf(x: Float32) -> Float32: - """math.h's sigmoidf: __frcp_rn(1.0f + __expf(-x)).""" - e = cute.arch.inline_ptx( - "ex2.approx.f32 {$w0}, {$r0};", - write_only_types=[Float32], - read_only_args=[-x * Float32(LOG2E)], - ) - return cute.arch.inline_ptx( - "rcp.rn.f32 {$w0}, {$r0};", - write_only_types=[Float32], - read_only_args=[Float32(1.0) + e], - ) 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 8395063ca6a..ff8f4f4dc9e 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -69,17 +69,13 @@ "dsrelu": dact_dsrelu, } - -# Gated activations of gated_mxfp8.cuh, as (ActOP, DActOP) ports of the math.h functions the -# CUDA dispatcher instantiates them with. swiglu and clamped_swiglu backward, and situglu in -# both directions, do not go through this table: like the CUDA kernel they take a special path. GATED_ACTIVATIONS = { "geglu": (act_gelu, dact_dgelu), "qgeglu": (act_qgelu, dact_dqgelu), "reglu": (act_relu, dact_drelu), "sreglu": (act_srelu, dact_dsrelu), "glu": (act_sigmoid, dact_dsigmoid), - "swiglu": (act_silu, None), + "swiglu": (act_silu, dact_dsilu), "clamped_swiglu": (None, None), "situglu": (None, None), } diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py index 70a1b65cf93..1a005585543 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py @@ -23,12 +23,13 @@ import cutlass from cutlass import cute from cutlass import pipeline -from cutlass import Boolean, Float32, Int32 +from cutlass import Float32, Int32 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( fabs_f32, is_packed16, + select_f32, truncate_f32_bf16, truncate_f32_f16, ) @@ -36,9 +37,13 @@ act_clamped_silu, act_situ_gate, dact_dsitu_gate, + identity, + didentity, situ_up, dsitu_up, - sigmoidf, + act_silu, + dact_dsilu, + sigmoid, ) from .quantize_mxfp8_common import ( CUTEDSL_DEBUG_LOGGING, @@ -54,23 +59,15 @@ @cute.jit -def _select_f32(pred: Boolean, if_true: Float32, if_false: Float32) -> Float32: - result = if_false - if pred: - result = if_true - return result - - -@cute.jit -def _gated_element( +def gated_elementwise_op( act_element: Float32, gate_element: Float32, y_grad_element: Optional[Float32], - limit: Optional[Float32], - alpha: Optional[Float32], - glu_linear_offset: Optional[Float32], - beta1: Optional[Float32], - beta2: Optional[Float32], + limit: Optional[Float32], # Clamped SwiGLU only + alpha: Optional[Float32], # Clamped SwiGLU only + glu_linear_offset: Optional[Float32], # Clamped SwiGLU only + beta1: Optional[Float32], # SiTU-GLU only + beta2: Optional[Float32], # SiTU-GLU only ACTIVATION: cutlass.Constexpr[str], IS_BWD: cutlass.Constexpr[bool], ): @@ -79,16 +76,23 @@ def _gated_element( BWD: (act_grad, gate_grad) act_grad = y_grad * gate(g) * dact(a), gate_grad = y_grad * act(a) * dgate(g) """ - # Transform the gate. Gating is ideally the identity, with derivative 1. - gate = gate_element - dgate = None + if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): - # |g| > limit is gated_mxfp8.cuh's (g > limit || g < -limit), NaN included - dgate = _select_f32(fabs_f32(gate_element) > limit, Float32(0.0), Float32(1.0)) + # FWD gate function for clamped_swiglu: + # gate(g) = clamp(g, -limit, limit) + glu_linear_offset gate = cute.arch.fmin(cute.arch.fmax(-limit, gate_element), limit) + glu_linear_offset + # BWD dgate function for clamped_swiglu: + # dgate(g) = 0 if g < -limit or g > limit + # = 1 if -limit <= g <= limit + dgate = select_f32(fabs_f32(gate_element) > limit, Float32(0.0), Float32(1.0)) elif cutlass.const_expr(ACTIVATION == "situglu"): - dgate = dsitu_up(gate_element, beta2) gate = situ_up(gate_element, beta2) + dgate = dsitu_up(gate_element, beta2) + else: + gate = identity(gate_element) + dgate = didentity(gate_element) + + act_fn, dact_fn = GATED_ACTIVATIONS[ACTIVATION] if cutlass.const_expr(not IS_BWD): if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): @@ -96,33 +100,31 @@ def _gated_element( elif cutlass.const_expr(ACTIVATION == "situglu"): act = act_situ_gate(act_element, beta1) else: - act = GATED_ACTIVATIONS[ACTIVATION][0](act_element) + act = act_fn(act_element) y = act * gate outputs = (y, None) else: if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): x = cute.arch.fmin(act_element, limit) - s = sigmoidf(alpha * x) + # Somehow CUDA always uses fast math for this and we should too to match its output + s = sigmoid(alpha * x, fast_math=True) act = x * s - dact = _select_f32( + dact = select_f32( act_element <= limit, s + s * (Float32(1.0) - s) * alpha * x, Float32(0.0) ) elif cutlass.const_expr(ACTIVATION == "situglu"): act = act_situ_gate(act_element, beta1) dact = dact_dsitu_gate(act_element, beta1) elif cutlass.const_expr(ACTIVATION == "swiglu"): - s = sigmoidf(act_element) - act = act_element * s - dact = act_element * s * (Float32(1.0) - s) + s + # Match CUDA implementation's behavior too + s = sigmoid(act_element, fast_math=True) + act = act_silu(act_element, s) + dact = dact_dsilu(act_element, s) else: - act_fn, dact_fn = GATED_ACTIVATIONS[ACTIVATION] act = act_fn(act_element) dact = dact_fn(act_element) act_grad = dact * y_grad_element * gate - if cutlass.const_expr(dgate is None): - gate_grad = act * y_grad_element - else: - gate_grad = dgate * act * y_grad_element + gate_grad = dgate * act * y_grad_element outputs = (act_grad, gate_grad) return outputs @@ -131,12 +133,12 @@ def _gated_element( def gated_elementwise_smem_tile( sAct_tile: cute.Tensor, # also the y tile in FWD / the act_grad tile in BWD sGate_tile: cute.Tensor, # also the gate_grad tile in BWD - sYGrad_tile: Optional[cute.Tensor], - limit: Optional[Float32], - alpha: Optional[Float32], - glu_linear_offset: Optional[Float32], - beta1: Optional[Float32], - beta2: Optional[Float32], + sYGrad_tile: Optional[cute.Tensor], # BWD only + limit: Optional[Float32], # Clamped SwiGLU only + alpha: Optional[Float32], # Clamped SwiGLU only + glu_linear_offset: Optional[Float32], # Clamped SwiGLU only + beta1: Optional[Float32], # SiTU-GLU only + beta2: Optional[Float32], # SiTU-GLU only ACTIVATION: cutlass.Constexpr[str], IS_BWD: cutlass.Constexpr[bool], ROWWISE: cutlass.Constexpr[bool], @@ -166,7 +168,7 @@ def gated_elementwise_smem_tile( # DTYPE below is then exact. if cutlass.const_expr(IS_BWD): y_grad_element = Float32(sYGrad_tile[i]) - act_grad_element, gate_grad_element = _gated_element( + act_grad_element, gate_grad_element = gated_elementwise_op( act_element, gate_element, y_grad_element, @@ -184,7 +186,7 @@ def gated_elementwise_smem_tile( sAct_tile[i] = DTYPE(act_grad_element) sGate_tile[i] = DTYPE(gate_grad_element) else: - y_element, _ = _gated_element( + y_element, _ = gated_elementwise_op( act_element, gate_element, None, diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index 90c61d03052..83ff738c6c7 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -9,7 +9,7 @@ import cutlass from cutlass import cute -from cutlass import Float32, Int64, Int32, Uint16, Uint32 +from cutlass import Boolean, Float32, Int64, Int32, Uint16, Uint32 # This must match `inline std::string to_string(const DType type)`` in transformer_engine/common/common.h exactly # which is how we encode types to strings in C++ @@ -69,6 +69,15 @@ def fabs_f32(val: Float32) -> Float32: return (val.bitcast(Int32) & Int32(0x7FFFFFFF)).bitcast(Float32) +@cute.jit +def select_f32(pred: Boolean, if_true: Float32, if_false: Float32) -> Float32: + """Return if_true if the runtime predicate pred holds, else if_false.""" + result = if_false + if pred: + result = if_true + return result + + @cute.jit def fma_f32(a: Float32, b: Float32, c: Float32) -> Float32: """Compute the fused multiply-add of three float32 values: a * b + c.""" From 24220c06f97ca900c6a66308a8a3c6c6f4bc42a3 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Tue, 3 Nov 2026 01:27:01 +0000 Subject: [PATCH 41/43] nit Signed-off-by: Kaining Zhong --- .../common/cast/mxfp8/gated_mxfp8_cutedsl.cuh | 21 ++++++++----------- 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh index cb6a3a44d07..adb3ced0f78 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh @@ -29,13 +29,6 @@ namespace cutedsl_backend { using namespace tvm_ffi_bridge; -namespace detail { - -// TMA needs 16-byte aligned global addresses. -inline bool is_16B_aligned(const void *ptr) { return reinterpret_cast(ptr) % 16 == 0; } - -} // namespace detail - struct MXFP8GatedQuantConfig { static constexpr const char *kEntrypointName = "get_mxfp8_gated_quantization_function"; @@ -133,11 +126,13 @@ bool mxfp8_quantize_gated_cutedsl(const Tensor &gated_input, const Tensor &grad, "output halves to be divisible by 128."); return false; } - bool aligned = detail::is_16B_aligned(gated_input.data.dptr) && - (!config.rowwise || detail::is_16B_aligned(output->data.dptr)) && - (!config.colwise || detail::is_16B_aligned(output->columnwise_data.dptr)); + // TMA needs aligned global addresses + bool aligned = + is_aligned_ptr(gated_input.data.dptr, TMA_GMEM_ALIGNMENT) && + (!config.rowwise || is_aligned_ptr(output->data.dptr, TMA_GMEM_ALIGNMENT)) && + (!config.colwise || is_aligned_ptr(output->columnwise_data.dptr, TMA_GMEM_ALIGNMENT)); if constexpr (IS_BWD) { - aligned = aligned && detail::is_16B_aligned(grad.data.dptr); + aligned = aligned && is_aligned_ptr(grad.data.dptr, TMA_GMEM_ALIGNMENT); } if (!aligned) { maybe_warn_cutedsl_not_chosen("an input or output is not 16-byte aligned."); @@ -150,13 +145,15 @@ bool mxfp8_quantize_gated_cutedsl(const Tensor &gated_input, const Tensor &grad, } checkCuDriverContext(stream); - // Sanity checks and swizzled-scale padding, mirroring mxfp8::quantize_gated + // Sanity checks, mirroring mxfp8::quantize_gated if (config.rowwise) { NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); } if (config.colwise) { NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); } + + // Zero out swizzled scales if padding is needed. if (config.swizzled && (cols % 128 != 0 || rows % 128 != 0)) { if (config.rowwise) { NVTE_CHECK_CUDA(cudaMemsetAsync(output->scale_inv.dptr, 0, From 29a1dd229fd8adb19031707f2adfaaf8c7f5c6a7 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 5 Nov 2026 23:13:14 +0000 Subject: [PATCH 42/43] refactor Signed-off-by: Kaining Zhong --- .../cast/mxfp8/quantize_mxfp8_gated_g2s.py | 30 ++++++++----------- 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py index 1a005585543..8b3208e8ce6 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py @@ -78,7 +78,7 @@ def gated_elementwise_op( """ if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): - # FWD gate function for clamped_swiglu: + # FWD gate function for clamped_swiglu: # gate(g) = clamp(g, -limit, limit) + glu_linear_offset gate = cute.arch.fmin(cute.arch.fmax(-limit, gate_element), limit) + glu_linear_offset # BWD dgate function for clamped_swiglu: @@ -272,43 +272,38 @@ def _split_gated_scales( COLWISE: cutlass.Constexpr[bool], SWIZZLED: cutlass.Constexpr[bool], ): - """Views of the backward output's scales split into its act_grad and gate_grad halves; returns - (mS_row_act, mS_row_gate, mS_col_act, mS_col_gate). - - The [M, 2N] output is [act_grad | gate_grad] and its scales are laid out for the whole - output. Each view has the shape of an [M, N] tensor's scales and the strides of the whole - output's, and starts where the whole output's layout places that half's first scale. For - swizzled scales this needs the halves to be whole 128x4 tiles (N % 128 == 0, required by the - C++ dispatcher). - """ + """Split the scales of a gated input with logical shape [M, 2N] into the act and gate halves, each with logical shape [M, N].""" mS_row_act = mS_row_gate = mS_col_act = mS_col_gate = None if cutlass.const_expr(SWIZZLED): mS_row, mS_col = derive_swizzled_scale_layout( M, 2 * N, ROWWISE, COLWISE, mS_row, mS_col ) - half_tiles = N // 128 + # To slice the left half (act) and right half (gate) of the swizzled scales, we need to tile using the tiler + # with the same mode pattern as the swizzled layout's shape: + # mS_row: ((32, 4, ceil(M/128)), (4, ceil(N/64))) : ((16, 4, 512·ceil(N/64)), (1, 512)) + # mS_col: ((4, ceil(M/128)), (32, 4, ceil(N/64))) : ((1, 512), (16, 4, 512·ceil(M/128))) if cutlass.const_expr(ROWWISE): row_tiler = ( - cute.make_layout((32, 4, cute.ceil_div(M, 128))), # all rows - cute.make_layout((4, half_tiles)), # one half's scale columns + cute.make_layout((32, 4, cute.ceil_div(M, 128))), # slice all rows + cute.make_layout((4, cute.ceil_div(N, 128))), # slice half columns so cdiv uses 128 instead of 64 ) mS_row_act = cute.local_tile(mS_row, row_tiler, (0, 0)) mS_row_gate = cute.local_tile(mS_row, row_tiler, (0, 1)) if cutlass.const_expr(COLWISE): col_tiler = ( - cute.make_layout((4, cute.ceil_div(cute.ceil_div(M, 32), 4))), # all scale rows - cute.make_layout((32, 4, half_tiles)), # one half's columns + cute.make_layout((4, cute.ceil_div(M, 128))), # slice all rows + cute.make_layout((32, 4, cute.ceil_div(N, 128))), # slice half columns so cdiv uses 128 instead of 64 ) mS_col_act = cute.local_tile(mS_col, col_tiler, (0, 0)) mS_col_gate = cute.local_tile(mS_col, col_tiler, (0, 1)) else: # Compact scales are row-major, so each half is a block of scale columns if cutlass.const_expr(ROWWISE): - row_tiler = (mS_row.shape[0], N // MXFP8_BLOCK_SCALING_SIZE) + row_tiler = (M, cute.ceil_div(N, MXFP8_BLOCK_SCALING_SIZE)) mS_row_act = cute.local_tile(mS_row, row_tiler, (0, 0)) mS_row_gate = cute.local_tile(mS_row, row_tiler, (0, 1)) if cutlass.const_expr(COLWISE): - col_tiler = (mS_col.shape[0], N) + col_tiler = (cute.ceil_div(M, MXFP8_BLOCK_SCALING_SIZE), N) mS_col_act = cute.local_tile(mS_col, col_tiler, (0, 0)) mS_col_gate = cute.local_tile(mS_col, col_tiler, (0, 1)) return mS_row_act, mS_row_gate, mS_col_act, mS_col_gate @@ -375,7 +370,6 @@ def __call__( mS_row_act_grad = mS_row_gate_grad = mS_col_act_grad = mS_col_gate_grad = None # BWD only if cutlass.const_expr(self.cfg.IS_BWD): # The [M, 2N] output is [act_grad | gate_grad]; the kernel takes its two halves - # TODO(kainingz): how to implement this in a more elegant way mS_row_act_grad, mS_row_gate_grad, mS_col_act_grad, mS_col_gate_grad = ( _split_gated_scales( mS_row, From a3474cc6f234b33b06f33085b52bcb1dabddeeb0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 21:13:38 +0000 Subject: [PATCH 43/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/jax/test_mxfp8_cutedsl_backend.py | 4 +- .../mxfp8/test_mxfp8_cutedsl_backend.py | 4 +- .../cast/mxfp8/quantize_mxfp8_gated_g2s.py | 104 ++++++++++-------- 3 files changed, 67 insertions(+), 45 deletions(-) diff --git a/tests/jax/test_mxfp8_cutedsl_backend.py b/tests/jax/test_mxfp8_cutedsl_backend.py index 2529cbfadfa..c44b93b5614 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -346,7 +346,9 @@ def get_gated_cfg_key(act_desc, in_dtype, fp8_dtype, q_layout, is_bwd): major, minor = device_compute_capability() return ( f"cutedsl_mxfp8_gated_sm{major * 10 + minor}_{DTYPE_TO_STR[in_dtype]}_" - f"{FP8_TO_KEY[fp8_dtype]}_" + "_".join("1" if f else "0" for f in flags) + f"_{act_desc}" + f"{FP8_TO_KEY[fp8_dtype]}_" + + "_".join("1" if f else "0" for f in flags) + + f"_{act_desc}" ) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index 97712afed84..fa51d292cb7 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -493,7 +493,9 @@ def get_gated_cfg_key(act, in_dtype, fp8_dtype, rowwise, colwise, swizzled, is_b flags = (rowwise, colwise, swizzled, is_bwd) return ( f"cutedsl_mxfp8_gated_sm{major * 10 + minor}_{DTYPE_TO_STR[in_dtype]}_" - f"{FP8_TO_KEY[fp8_dtype]}_" + "_".join("1" if f else "0" for f in flags) + f"_{act['desc']}" + f"{FP8_TO_KEY[fp8_dtype]}_" + + "_".join("1" if f else "0" for f in flags) + + f"_{act['desc']}" ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py index 8b3208e8ce6..dea96563611 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py @@ -63,11 +63,11 @@ def gated_elementwise_op( act_element: Float32, gate_element: Float32, y_grad_element: Optional[Float32], - limit: Optional[Float32], # Clamped SwiGLU only - alpha: Optional[Float32], # Clamped SwiGLU only - glu_linear_offset: Optional[Float32], # Clamped SwiGLU only - beta1: Optional[Float32], # SiTU-GLU only - beta2: Optional[Float32], # SiTU-GLU only + limit: Optional[Float32], # Clamped SwiGLU only + alpha: Optional[Float32], # Clamped SwiGLU only + glu_linear_offset: Optional[Float32], # Clamped SwiGLU only + beta1: Optional[Float32], # SiTU-GLU only + beta2: Optional[Float32], # SiTU-GLU only ACTIVATION: cutlass.Constexpr[str], IS_BWD: cutlass.Constexpr[bool], ): @@ -131,14 +131,14 @@ def gated_elementwise_op( @cute.jit def gated_elementwise_smem_tile( - sAct_tile: cute.Tensor, # also the y tile in FWD / the act_grad tile in BWD - sGate_tile: cute.Tensor, # also the gate_grad tile in BWD - sYGrad_tile: Optional[cute.Tensor], # BWD only - limit: Optional[Float32], # Clamped SwiGLU only - alpha: Optional[Float32], # Clamped SwiGLU only - glu_linear_offset: Optional[Float32], # Clamped SwiGLU only - beta1: Optional[Float32], # SiTU-GLU only - beta2: Optional[Float32], # SiTU-GLU only + sAct_tile: cute.Tensor, # also the y tile in FWD / the act_grad tile in BWD + sGate_tile: cute.Tensor, # also the gate_grad tile in BWD + sYGrad_tile: Optional[cute.Tensor], # BWD only + limit: Optional[Float32], # Clamped SwiGLU only + alpha: Optional[Float32], # Clamped SwiGLU only + glu_linear_offset: Optional[Float32], # Clamped SwiGLU only + beta1: Optional[Float32], # SiTU-GLU only + beta2: Optional[Float32], # SiTU-GLU only ACTIVATION: cutlass.Constexpr[str], IS_BWD: cutlass.Constexpr[bool], ROWWISE: cutlass.Constexpr[bool], @@ -272,27 +272,29 @@ def _split_gated_scales( COLWISE: cutlass.Constexpr[bool], SWIZZLED: cutlass.Constexpr[bool], ): - """Split the scales of a gated input with logical shape [M, 2N] into the act and gate halves, each with logical shape [M, N].""" + """Split the scales of a gated input with logical shape [M, 2N] into the act and gate halves, each with logical shape [M, N].""" mS_row_act = mS_row_gate = mS_col_act = mS_col_gate = None if cutlass.const_expr(SWIZZLED): - mS_row, mS_col = derive_swizzled_scale_layout( - M, 2 * N, ROWWISE, COLWISE, mS_row, mS_col - ) + mS_row, mS_col = derive_swizzled_scale_layout(M, 2 * N, ROWWISE, COLWISE, mS_row, mS_col) # To slice the left half (act) and right half (gate) of the swizzled scales, we need to tile using the tiler # with the same mode pattern as the swizzled layout's shape: # mS_row: ((32, 4, ceil(M/128)), (4, ceil(N/64))) : ((16, 4, 512·ceil(N/64)), (1, 512)) # mS_col: ((4, ceil(M/128)), (32, 4, ceil(N/64))) : ((1, 512), (16, 4, 512·ceil(M/128))) if cutlass.const_expr(ROWWISE): row_tiler = ( - cute.make_layout((32, 4, cute.ceil_div(M, 128))), # slice all rows - cute.make_layout((4, cute.ceil_div(N, 128))), # slice half columns so cdiv uses 128 instead of 64 + cute.make_layout((32, 4, cute.ceil_div(M, 128))), # slice all rows + cute.make_layout( + (4, cute.ceil_div(N, 128)) + ), # slice half columns so cdiv uses 128 instead of 64 ) mS_row_act = cute.local_tile(mS_row, row_tiler, (0, 0)) mS_row_gate = cute.local_tile(mS_row, row_tiler, (0, 1)) if cutlass.const_expr(COLWISE): col_tiler = ( - cute.make_layout((4, cute.ceil_div(M, 128))), # slice all rows - cute.make_layout((32, 4, cute.ceil_div(N, 128))), # slice half columns so cdiv uses 128 instead of 64 + cute.make_layout((4, cute.ceil_div(M, 128))), # slice all rows + cute.make_layout( + (32, 4, cute.ceil_div(N, 128)) + ), # slice half columns so cdiv uses 128 instead of 64 ) mS_col_act = cute.local_tile(mS_col, col_tiler, (0, 0)) mS_col_gate = cute.local_tile(mS_col, col_tiler, (0, 1)) @@ -338,11 +340,11 @@ def __call__( mO_col: Optional[cute.Tensor], # [M, C], not transposed mS_col: Optional[cute.Tensor], # [roundup(ceil(M / 32), 4), roundup(C, 128)] # Special activation params, None if not used by the activation - limit: Optional[Float32], # Clamped SwiGLU only - alpha: Optional[Float32], # Clamped SwiGLU only - glu_linear_offset: Optional[Float32], # Clamped SwiGLU only - beta1: Optional[Float32], # SiTUGLU only - beta2: Optional[Float32], # SiTUGLU only + limit: Optional[Float32], # Clamped SwiGLU only + alpha: Optional[Float32], # Clamped SwiGLU only + glu_linear_offset: Optional[Float32], # Clamped SwiGLU only + beta1: Optional[Float32], # SiTUGLU only + beta2: Optional[Float32], # SiTUGLU only stream: CUstream, ): """ @@ -354,7 +356,9 @@ def __call__( [act_grad | gate_grad] = quantize([ y_grad * gate(g) * dact(a) | y_grad * act(a) * dgate(g) ]) with shape [M, 2N] """ if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): - cute.printf(f"[CuTeDSL] MXFP8GatedQuantizeG2SKernel.__call__() with config: {self.cfg}\n") + cute.printf( + f"[CuTeDSL] MXFP8GatedQuantizeG2SKernel.__call__() with config: {self.cfg}\n" + ) M = mGatedInput.shape[0] # Pick any mO that is used just to get N @@ -650,8 +654,20 @@ class BarrierStorage: # pylint: disable=missing-class-docstring ) # In BWD mode we output gradient for both the act and gate halves, which are stored in sAct and sGate respectively tile_quantizer_params = [ - (sAct, sO_row_act_grad, mS_row_act_grad, sO_col_act_grad, mS_col_act_grad), # We will write unquantized act_grad to sAct in-place - (sGate, sO_row_gate_grad, mS_row_gate_grad, sO_col_gate_grad, mS_col_gate_grad), # We will write unquantized gate_grad to sGate in-place + ( + sAct, + sO_row_act_grad, + mS_row_act_grad, + sO_col_act_grad, + mS_col_act_grad, + ), # We will write unquantized act_grad to sAct in-place + ( + sGate, + sO_row_gate_grad, + mS_row_gate_grad, + sO_col_gate_grad, + mS_col_gate_grad, + ), # We will write unquantized gate_grad to sGate in-place ] else: if cutlass.const_expr(cfg.ROWWISE): @@ -659,7 +675,9 @@ class BarrierStorage: # pylint: disable=missing-class-docstring if cutlass.const_expr(cfg.COLWISE): mS_col = cute.local_tile(mS_col, cta_block_col_scales, cta_coord) # In FWD mode we only output the result after the gated op, which is stored in sAct - tile_quantizer_params = [(sAct, sO_row, mS_row, sO_col, mS_col)] # We will write unquantized output to sAct in-place + tile_quantizer_params = [ + (sAct, sO_row, mS_row, sO_col, mS_col) + ] # We will write unquantized output to sAct in-place # Prepare input TMA (atom, smem partition, gmem partition) triples: act and gate, and BWD's y_grad TMA_load_atom_and_partitions = [] @@ -681,12 +699,12 @@ class BarrierStorage: # pylint: disable=missing-class-docstring # Prepare output TMA (atom, smem partition, gmem partition) triples of the outputs this config has TMA_store_atom_and_partitions = [] for atom_out, dst_out, sO_out in ( - (tma_atom_row, tma_dst_row, sO_row), # FWD only - (tma_atom_col, tma_dst_col, sO_col), # FWD only - (tma_atom_row_act_grad, tma_dst_row_act_grad, sO_row_act_grad), # BWD only - (tma_atom_row_gate_grad, tma_dst_row_gate_grad, sO_row_gate_grad), # BWD only - (tma_atom_col_act_grad, tma_dst_col_act_grad, sO_col_act_grad), # BWD only - (tma_atom_col_gate_grad, tma_dst_col_gate_grad, sO_col_gate_grad), # BWD only + (tma_atom_row, tma_dst_row, sO_row), # FWD only + (tma_atom_col, tma_dst_col, sO_col), # FWD only + (tma_atom_row_act_grad, tma_dst_row_act_grad, sO_row_act_grad), # BWD only + (tma_atom_row_gate_grad, tma_dst_row_gate_grad, sO_row_gate_grad), # BWD only + (tma_atom_col_act_grad, tma_dst_col_act_grad, sO_col_act_grad), # BWD only + (tma_atom_col_gate_grad, tma_dst_col_gate_grad, sO_col_gate_grad), # BWD only ): if cutlass.const_expr(atom_out is not None): tXsO_out, tXgO_out = cute.nvgpu.cpasync.tma_partition( @@ -743,8 +761,8 @@ class BarrierStorage: # pylint: disable=missing-class-docstring # FWD: write the output into sAct_tile # BWD: write the act half's gradient to sAct_tile, and write the gate half's gradient to sGate_tile gated_elementwise_smem_tile( - sAct_tile, # we also store the FWD output / BWD act half's gradient here - sGate_tile, # we also store the BWD gate half's gradient here + sAct_tile, # we also store the FWD output / BWD act half's gradient here + sGate_tile, # we also store the BWD gate half's gradient here sYGrad_tile, limit, alpha, @@ -777,11 +795,11 @@ class BarrierStorage: # pylint: disable=missing-class-docstring ) # Use the general MXFP8 SMEM tile quantize function to quantize it quantize_mxfp8_smem_tile_general_partial( - sX_out[(None, stage_idx)], # The SMEM tile to quantize - sO_row_tile, # SMEM tile for rowwise quantized output - mS_row_stage, # GMEM tile for rowwise quantized scales - sO_col_tile, # SMEM tile for columnwise quantized output - mS_col_stage, # GMEM tile for columnwise quantized scales + sX_out[(None, stage_idx)], # The SMEM tile to quantize + sO_row_tile, # SMEM tile for rowwise quantized output + mS_row_stage, # GMEM tile for rowwise quantized scales + sO_col_tile, # SMEM tile for columnwise quantized output + mS_col_stage, # GMEM tile for columnwise quantized scales # The tile's first row and column in the whole tensor, for bounds checks (bidy * NUM_TILES + tile_idx) * TILE_ROWS, bidx * TILE_COLS,