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/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 4671ea06869..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) @@ -283,3 +286,33 @@ 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). + "g2r_rowwise2lane_tier0": ((1024, 4096), 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 + "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..5d7feeb6ce6 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -59,7 +59,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 @@ -250,31 +251,39 @@ 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 -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( @@ -370,3 +379,55 @@ 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, 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. + ("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, 0), + # Row tiles past grid.y's limit spill into grid.z. + ("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, 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, 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, 0), + # G2SBidim: fp16, or bf16 with N % 256 != 0. + ("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,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, input_offset): + run_test_case( + "CAST_ONLY", + IDENTITY, + shape, + block_size, + in_dtype, + tex.DType.kFloat8E4M3, + swizzled, + input_offset, + ) 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.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py deleted file mode 100644 index 824ed18aa59..00000000000 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py +++ /dev/null @@ -1,2846 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""MXFP8 quantization kernel implemented in CuTeDSL. - -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). - -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 - MXFP8_BLOCK_SCALING_SIZE = 32 (elements per MXFP8 scaling block) - -Grid: (ceil(N / 64), ceil(M / 64)) -Each block processes a 64x64 chunk in 2 stages of 32x64 tiles loaded into -shared memory. -""" - -# 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 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, - max_x2_bf16, - max_x2_f16, - max_scalar_bf16, - max_scalar_f16, - abs_max_scalar_bf16, - abs_max_scalar_f16, - to_f32_bf16, - to_f32_f16, - x2_lo_to_f32_bf16, - x2_lo_to_f32_f16, - x2_hi_to_f32_bf16, - x2_hi_to_f32_f16, - truncate_f32_bf16, - truncate_f32_f16, - fabs_f32, - exp2f_rcp, - 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 - - -@cute.jit -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) - mS_row_stage: cute.Tensor, # rowwise scale tensor (1D swizzled, or 2D linear) - MAX_NORM_RCP: cutlass.Constexpr[float], - tile_row_start: Int32, # global row index of this stage's row 0 - tile_col_start: Int32, # global col index of this CTA's col 0 - M: Int32, - N: Int32, - ACTIVATION: cutlass.Constexpr[str | None], - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - 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, - WITH_DACT: cutlass.Constexpr[bool] = False, - WITH_DBIAS: cutlass.Constexpr[bool] = False, - dbias_acc: Optional[cute.Tensor] = None, # only needed when WITH_DBIAS is True -): - """Quantize one SMEM tile rowwise to MXFP8 (per-row 32-elt block scales); returns the tile amax.""" - tidx, _, _ = cute.arch.thread_idx() - - 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) - - _, tv_layout = 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, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), - ) - - sX_tv = cute.composition(sX_tile, tv_layout) - sO_tv = cute.composition(sO_row_tile, tv_layout) - - # I/O Elements that belong to this thread - sX_thread = sX_tv[tidx, None] # shape (32,) bf16 - sO_thread = sO_tv[tidx, None] # shape (32,) fp8 - - sO_thread_u32_ptr = cute.recast_ptr(sO_thread.iterator, dtype=Uint32) - sO_thread_u32 = cute.make_tensor( - sO_thread_u32_ptr, - cute.make_layout( - (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) - ), # 1 uint32 is 4 fp8 elements - ) - - # PTX allows to fuse relu activation in `cvt.rn.satfinite` unless we need to reduction for dbias - FUSE_RELU = cutlass.const_expr(ACTIVATION == "relu") and not WITH_DBIAS - # For this fast path we can read in pack of 2 instead of reading individual f16 / bf16 element, - # and keep the elements in half precision instead of upcasting every one of them to f32. - USE_HALF_PRECISION = is_packed16(DTYPE) and (ACTIVATION is None or FUSE_RELU) - - amax_r = Float32(0.0) - - # Each thread start reading from the specfic bank based on its thread ID so they can do their best to access different banks - # to avoid bank conflict. - bank_group = (tidx % THREADS_PER_WARP) // THREADS_PER_BANK - # The offset this thread should start reading from based on what's its first bank to access. - offset = bank_group * PACK_SIZE - if cutlass.const_expr(USE_HALF_PRECISION): - # If no activation, f16 / bf16 and rowwise quantization, we can read 2 f16 / bf16 at once in a pack - # and use max.xorsign.abs.f16x2 / max.xorsign.abs.bf16x2 to compute - max_x2 = max_x2_f16 if DTYPE is cutlass.Float16 else max_x2_bf16 - 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 - sX_thread_rw_i64 = cute.make_tensor( - cute.recast_ptr(sX_thread.iterator, dtype=Int64), - cute.make_layout( - (1, MXFP8_BLOCK_SCALING_SIZE // 4), stride=(0, 1) - ), # 1 int64 is 4 fp16/bf16 elements - ) - # Each wave reads its 4 elements (PACK_SIZE) as one 8-byte vectorized load - in_r = [[None, None] for _ in range(WAVES)] - for w in cutlass.range_constexpr(WAVES): - idx = (w + offset // 4) % (MXFP8_BLOCK_SCALING_SIZE // 4) - in_r[w][0], in_r[w][1] = unpack_i64_to_i32x2(sX_thread_rw_i64[0, idx]) - - if cutlass.const_expr(WITH_DBIAS): - for w in cutlass.range_constexpr(WAVES): - dbias_acc[w * PACK_SIZE + 0] += x2_lo_to_f32(in_r[w][0]) - dbias_acc[w * PACK_SIZE + 1] += x2_hi_to_f32(in_r[w][0]) - dbias_acc[w * PACK_SIZE + 2] += x2_lo_to_f32(in_r[w][1]) - dbias_acc[w * PACK_SIZE + 3] += x2_hi_to_f32(in_r[w][1]) - - amax_2x = in_r[0][0] - # Each wave will use max.xorsign.abs.f16x2 or max.xorsign.abs.bf16x2 to compare 2 packed elements in parallel - for w in cutlass.range_constexpr(WAVES): - if cutlass.const_expr(FUSE_RELU): - # Skip the first iteration because we assigned in_r[0][0] to amax_2x already - if cutlass.const_expr(w > 0): - # If we fuse relu then we don't want to do abs since negative value will be set to 0 and they will lose comparison automatically - amax_2x = max_x2(amax_2x, in_r[w][0]) - amax_2x = max_x2(amax_2x, in_r[w][1]) - else: - # Skip the first iteration because we assigned in_r[0][0] to amax_2x already - if cutlass.const_expr(w > 0): - amax_2x = abs_max_x2(amax_2x, in_r[w][0]) - amax_2x = abs_max_x2(amax_2x, in_r[w][1]) - if cutlass.const_expr(FUSE_RELU): - # Compare the 2 packed max without abs - amax_r = cute.arch.fmax( - x2_lo_to_f32(amax_2x), - x2_hi_to_f32(amax_2x), - ) - # For relu the max is at least 0 - amax_r = cute.arch.fmax(amax_r, Float32(0.0)) - else: - # Compare the 2 packed abs max - amax_r = cute.arch.fmax( - fabs_f32(x2_lo_to_f32(amax_2x)), - fabs_f32(x2_hi_to_f32(amax_2x)), - ) - else: - # Since we need to do computation on individual f16 / bf16 elements, we can't read in pack - sX_thread_rw = cute.make_tensor( - sX_thread.iterator, - cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), - ) - - if cutlass.const_expr(WITH_DACT): - # Backward: out = grad · act'(act_input). sX is grad, sA is act_input. - dop = SUPPORTED_DACTIVATIONS[ACTIVATION] - sA_thread = cute.composition(sActInput_tile, tv_layout)[tidx, None] - sA_thread_rw = cute.make_tensor( - sA_thread.iterator, - cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), - ) - elif cutlass.const_expr(WITH_ACT): - op = SUPPORTED_ACTIVATIONS[ACTIVATION] - - if cutlass.const_expr(is_packed16(DTYPE) and ACTIVATION is not None): - truncate_f32 = truncate_f32_f16 if DTYPE is cutlass.Float16 else truncate_f32_bf16 - - # Each wave we read PACK_SIZE elements, and we have WAVES waves, so we read WAVES * PACK_SIZE (= MXFP8_BLOCK_SCALING_SIZE) elements in total. - in_r = [[None] * PACK_SIZE for _ in range(WAVES)] - # Don't repeat these OOB masking computation in the loop - thread_row_start = tile_row_start + tidx // CTA_THREADS_X - thread_col_start = tile_col_start + (tidx % CTA_THREADS_X) * MXFP8_BLOCK_SCALING_SIZE - thread_row_oob = thread_row_start >= M - # PACK_SIZE must divide every term of `thread_col_start + start` and N itself, or a - # pack could straddle the N boundary and the check below would be wrong. TILE_X is a - # tunable, so it is asserted rather than assumed. - assert MXFP8_BLOCK_SCALING_SIZE % PACK_SIZE == 0 - assert SYM_N_DIVISIBILITY % PACK_SIZE == 0 - assert TILE_X % PACK_SIZE == 0 - for w in cutlass.range_constexpr(WAVES): - start = (w * PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE - # thread_col_start + start >= N means "if the first element in this pack is OOB" - # If the first element is OOB, then (obviously) elements after it within this pack must be OOB - # If the first element is not OOB, then elements after it within this pack must be not OOB either. - # Proof: - # (1) when the first element is not OOB, we have thread_col_start + start < N, - # which implies N - (thread_col_start + start) > 0 - # (2) with the assertions above: - # - N is divisible by SYM_N_DIVISIBILITY, so N is also divisible by PACK_SIZE - # - thread_col_start = tile_col_start + (tidx % CTA_THREADS_X) * MXFP8_BLOCK_SCALING_SIZE is divisible by PACK_SIZE - # because TILE_X and MXFP8_BLOCK_SCALING_SIZE are both divisible by PACK_SIZE - # - start = (w * PACK_SIZE + bank_group * PACK_SIZE) % MXFP8_BLOCK_SCALING_SIZE is also divisible by PACK_SIZE - # because MXFP8_BLOCK_SCALING_SIZE is a multiple of PACK_SIZE - # so N - (thread_col_start + start) must be divisible by PACK_SIZE - # Therefore, N - (thread_col_start + start) >= PACK_SIZE because PACK_SIZE is the minimal number - # for (1) and (2) to hold, which implies thread_col_start + start + PACK_SIZE -1 < N - # Hence, we only need to check if the first element is OOB to determine if the whole pack is OOB - thread_col_oob = thread_col_start + start >= N - for i in cutlass.range_constexpr(PACK_SIZE): - x = Float32(sX_thread_rw[0, start + i]) - if cutlass.const_expr(WITH_DACT): - # out = grad · act'(act_input) - x = x * dop(Float32(sA_thread_rw[0, start + i])) - # If IS_ACT, apply activation function to x in f32 - elif cutlass.const_expr(WITH_ACT): - # If it's relu, we can handle it later - if not cutlass.const_expr(FUSE_RELU): - x = op(x) - if not cutlass.const_expr(SKIP_INPUT_MASKING): - # If the input shape is not divisible by the tile size, - # TMA would zero-fills the input tile outside its logical MxN bounds. - # This is fine for non-activation cases, but for activation cases, - # op(0) might not be 0 which will pollute the amax and dbias. - # So we must manually mask the OOB region here. - if thread_row_oob or thread_col_oob: - x = Float32(0.0) - # Accumulate to the per-thread dbias register buffer for this tile if WITH_DBIAS - if cutlass.const_expr(WITH_DBIAS): - # dbias_acc is register buffer so we can just write without bank conflict - dbias_acc[w * PACK_SIZE + i] += x - # If 16-bit input with activation, truncate to IType - if cutlass.const_expr(is_packed16(DTYPE) and ACTIVATION is not None): - x = truncate_f32(x) - in_r[w][i] = x - if cutlass.const_expr(FUSE_RELU): - amax_r = cute.arch.fmax( - amax_r, x - ) # For relu cases, we don't need abs since negative values will be 0 so they lose comparison automatically - else: - amax_r = cute.arch.fmax(amax_r, fabs_f32(x)) - if cutlass.const_expr(FUSE_RELU): - amax_r = cute.arch.fmax(amax_r, Float32(0.0)) # If relu, the amax is at least 0 - - biased_exp_r = cvt_f32_to_fp8e8m0fnu(amax_r * MAX_NORM_RCP) - - # mS_row_stage has logical shape (32, 2) and we have 64 threads where each is mapped to one scale factor - # The TV layout is equivalent to TV layout with thr_layout=(32, 2):(2, 1), val_layout=(1,) - # but it's too trival so let's just index it directly without using layout - # Note this is the logical layout, which is on top of the swizzled / non-swizzled scale factor layout - # that mappes the logical index to the physical offset - - # 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 cutlass.const_expr(SKIP_SCALE_BOUNDS): - mS_row_stage[(tidx // CTA_THREADS_X, tidx % CTA_THREADS_X)] = biased_exp_r - else: - scale_row = tile_row_start + tidx // CTA_THREADS_X - scale_col_first_elt = tile_col_start + (tidx % CTA_THREADS_X) * MXFP8_BLOCK_SCALING_SIZE - if scale_row < M and scale_col_first_elt < N: - mS_row_stage[(tidx // CTA_THREADS_X, tidx % CTA_THREADS_X)] = biased_exp_r - - inv_scale_r = exp2f_rcp(biased_exp_r) # f32 reciprocal of the scale - scale_2x = pack_f32x2(inv_scale_r, inv_scale_r) - if cutlass.const_expr(USE_HALF_PRECISION): - mul_cvt_x4_func = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, FP8_DTYPE, FUSE_RELU) - else: - mul_cvt_x4_func = mul_f32x2_cvt_f32x4_to_fp8x4(FP8_DTYPE, FUSE_RELU) - - for w in cutlass.range_constexpr(WAVES): - idx = (w * 4 + offset) % MXFP8_BLOCK_SCALING_SIZE - idx = idx // 4 - if cutlass.const_expr(USE_HALF_PRECISION): - # Convert 2 packed f16/bf16 pairs to 4 fp8 in one fused op - sO_thread_u32[idx] = mul_cvt_x4_func(in_r[w][0], in_r[w][1], scale_2x) - else: - # Convert 4 f32 to 4 fp8 in one fused op - sO_thread_u32[idx] = mul_cvt_x4_func( - in_r[w][0], in_r[w][1], in_r[w][2], in_r[w][3], scale_2x - ) - - return amax_r - - -@cute.jit -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) - mS_col_stage: cute.Tensor, # colwise scale tensor (1D swizzled, or 2D linear) - MAX_NORM_RCP: cutlass.Constexpr[float], - tile_row_start: Int32, # global row index of this stage's row 0 - tile_col_start: Int32, # global col index of this CTA's col 0 - M: Int32, - N: Int32, - ACTIVATION: cutlass.Constexpr[str | None], - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - SWIZZLE: cutlass.Constexpr[bool], - TILE_X: cutlass.Constexpr[int], - TILE_Y: cutlass.Constexpr[int], # pylint: disable=unused-argument # kept for API consistency - SKIP_INPUT_MASKING: cutlass.Constexpr[bool] = False, - SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool] = False, - WITH_ACT: cutlass.Constexpr[bool] = False, - WITH_DACT: cutlass.Constexpr[bool] = False, - WITH_DBIAS: cutlass.Constexpr[bool] = False, - CACHE_ACTIVATION: cutlass.Constexpr[bool] = False, # cache post-activation values to sX_tile -): - """Quantize one SMEM tile colwise to MXFP8 (per-column 32-elt block scales); returns (amax, dbias_partial).""" - 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)), - ) - - sX_tv = cute.composition(sX_tile, tv_layout) - sO_tv = cute.composition(sO_col_tile, tv_layout) - - # I/O Elements that belong to this thread - sX_thread = sX_tv[tidx, None] - sO_thread = sO_tv[tidx, None] - - # PTX allows to fuse relu activation in `cvt.rn.satfinite` unless we need to reduction for dbias - FUSE_RELU = cutlass.const_expr(ACTIVATION == "relu") and not WITH_DBIAS - # Keep input in half precision format if possible - USE_HALF_PRECISION = is_packed16(DTYPE) and (ACTIVATION is None or FUSE_RELU) - dbias_partial = Float32(0.0) - - if cutlass.const_expr(USE_HALF_PRECISION): - max_scalar = max_scalar_f16 if DTYPE is cutlass.Float16 else max_scalar_bf16 - abs_max_scalar = abs_max_scalar_f16 if DTYPE is cutlass.Float16 else abs_max_scalar_bf16 - to_f32 = to_f32_f16 if DTYPE is cutlass.Float16 else to_f32_bf16 - # If we can use the half precision format, then use the input tile directly since there is no need to upcast - sX_thread_packed16 = cute.make_tensor( - sX_thread.iterator, - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(TILE_X,)), - ) - # Stash the strided column reads in registers (CUDA's in_colwise_IType): - # the cvt loop below reuses them instead of re-reading smem. - in_c = [None] * MXFP8_BLOCK_SCALING_SIZE - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - in_c[i] = sX_thread_packed16[i] - - amax_value = in_c[0] - # Skip the first iteration because we assigned in_c[0] to amax_value already - for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE): - if cutlass.const_expr(FUSE_RELU): - # If we fuse relu then we don't want to do abs since negative value will be set to 0 - # and they will lose comparison automatically - amax_value = max_scalar(amax_value, in_c[i]) - else: - amax_value = abs_max_scalar(amax_value, in_c[i]) - if cutlass.const_expr(FUSE_RELU): - amax_c = cute.arch.fmax(to_f32(amax_value), Float32(0.0)) - else: - amax_c = fabs_f32(to_f32(amax_value)) - else: - # Otherwise we need to case input values to fp32. Allocate the register tensor and load from SMEM input tiles. - rX_thread_f32 = cute.make_rmem_tensor( - layout_or_shape=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(1,)), - dtype=Float32, - ) - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - rX_thread_f32[i] = Float32(sX_thread[i]) - # Apply activation (fwd) or grad·act'(act_input) (bwd dact) in f32. - if cutlass.const_expr(WITH_DACT): - dop = SUPPORTED_DACTIVATIONS[ACTIVATION] - sA_thread = cute.composition(sActInput_tile, tv_layout)[tidx, None] - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - rX_thread_f32[i] = rX_thread_f32[i] * dop(Float32(sA_thread[i])) - elif cutlass.const_expr(WITH_ACT): - op = SUPPORTED_ACTIVATIONS[ACTIVATION] - # Don't repeat these OOB masking computation in the loop - thread_col_start = tile_col_start + tidx - thread_col_oob = thread_col_start >= N - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - # If it's relu, we can handle it later in the cvt - if not cutlass.const_expr(FUSE_RELU): - rX_thread_f32[i] = op(rX_thread_f32[i]) - # If the input shape is not divisible by the tile size, - # TMA would zero-fills the input tile outside its logical MxN bounds. - # This is fine for non-activation cases, but for activation cases, - # op(0) might not be 0 which will pollute the amax and dbias. - # So we must manually mask the OOB region here. - if not cutlass.const_expr(SKIP_INPUT_MASKING): - if tile_row_start + i >= M or thread_col_oob: - rX_thread_f32[i] = Float32(0.0) - # Accumulate fp32 activations to DBIAS before we truncate to half precision when the input is half precision - if cutlass.const_expr(WITH_DBIAS): - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - dbias_partial += rX_thread_f32[i] - # Truncate the activation (after we apply op) back to the half precision type if input is also half precision. - if cutlass.const_expr(is_packed16(DTYPE) and ACTIVATION is not None): - truncate_f32 = truncate_f32_f16 if DTYPE is cutlass.Float16 else truncate_f32_bf16 - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - rX_thread_f32[i] = truncate_f32(rX_thread_f32[i]) - # Columnwise is the preferred direction so it runs first. If it needs to cache the activation in the input tile - # to let the rowwise pass read it, we need to cast and overwrite the input data in-place here - if cutlass.const_expr(CACHE_ACTIVATION): - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - sX_thread[i] = DTYPE(rX_thread_f32[i]) - amax_c = Float32(0.0) - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - if cutlass.const_expr(FUSE_RELU): - amax_c = cute.arch.fmax(amax_c, rX_thread_f32[i]) - else: - amax_c = cute.arch.fmax(amax_c, fabs_f32(rX_thread_f32[i])) - - # Irregular shapes: skip when this stage's row range or this thread's - # 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 cutlass.const_expr(SKIP_SCALE_BOUNDS): - if cutlass.const_expr(SWIZZLE): - mS_col_stage[(0, tidx % 32, tidx // 32)] = biased_exp_c - else: - mS_col_stage[(0, tidx)] = biased_exp_c - else: - scale_col = tile_col_start + tidx - if tile_row_start < M and scale_col < N: - if cutlass.const_expr(SWIZZLE): - mS_col_stage[(0, tidx % 32, tidx // 32)] = biased_exp_c - else: - mS_col_stage[(0, tidx)] = biased_exp_c - - inv_scale_c = exp2f_rcp(biased_exp_c) - # cvt.rn.satfinite can be vectorized to convert 2 f32 to 2 fp8 in one instruction - cvt_x2_func = get_cvt_f32x2_to_fp8x2_func(FP8_DTYPE, FUSE_RELU) - sO_thread_fp8 = cute.make_tensor( - cute.recast_ptr(sO_thread.iterator, dtype=FP8_DTYPE), sO_thread.layout - ) - if cutlass.const_expr(USE_HALF_PRECISION): - to_f32 = to_f32_f16 if DTYPE is cutlass.Float16 else to_f32_bf16 - for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): - lo, hi = 2 * j, 2 * j + 1 - v_lo = to_f32(in_c[lo]) - v_hi = to_f32(in_c[hi]) - # Accumulate the per-thread column partial for dbias if WITH_DBIAS. - # Kept as two adds in element order: f32 addition is not associative - if cutlass.const_expr(WITH_DBIAS): - dbias_partial += v_lo - dbias_partial += v_hi - pair = cvt_x2_func(v_hi * inv_scale_c, v_lo * inv_scale_c) - sO_thread_fp8[lo] = Uint8(pair & Uint16(0xFF)).bitcast(FP8_DTYPE) - sO_thread_fp8[hi] = Uint8(pair >> Uint16(8)).bitcast(FP8_DTYPE) - else: - for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): - lo, hi = 2 * j, 2 * j + 1 - pair = cvt_x2_func(rX_thread_f32[hi] * inv_scale_c, rX_thread_f32[lo] * inv_scale_c) - sO_thread_fp8[lo] = Uint8(pair & Uint16(0xFF)).bitcast(FP8_DTYPE) - sO_thread_fp8[hi] = Uint8(pair >> Uint16(8)).bitcast(FP8_DTYPE) - - # Return this stage's per-column partial alongside amax; the caller accumulates - # it across stages (a scalar can't be updated in-place through the arg). - return amax_c, dbias_partial - - -@cute.jit -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], - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - TILE_X: cutlass.Constexpr[int], - WITH_DACT: cutlass.Constexpr[bool] = False, - CACHE_ACTIVATION: cutlass.Constexpr[bool] = False, # cache post-activation values to sX_tile -): - """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. - _, 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)), - ) - sX_thread = cute.composition(sX_tile, tv_layout)[tidx, None] - - dbias_partial = Float32(0.0) - if cutlass.const_expr(WITH_DACT): - dop = SUPPORTED_DACTIVATIONS[ACTIVATION] - sA_thread = cute.composition(sA_tile, tv_layout)[tidx, None] - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - x = Float32(sX_thread[i]) * dop(Float32(sA_thread[i])) - dbias_partial += x - if cutlass.const_expr(CACHE_ACTIVATION): - sX_thread[i] = DTYPE(x) - else: - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - dbias_partial += Float32(sX_thread[i]) - 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.). - `__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__( - self, - cfg: MXFP8QuantizeConfig, - SKIP_INPUT_MASKING: bool = False, - SKIP_SCALE_BOUNDS: bool = False, - ): - self.cfg = cfg - # If the input shape is divisible by the tile size or f(0)=0 holds for activaions, - # we can skip masking inputs with zero in the kernel and save some instructions. - self.SKIP_INPUT_MASKING = SKIP_INPUT_MASKING - # If the input shape is divisible by the tile size, we can skip bounds check for scale writes and save some instructions. - self.SKIP_SCALE_BOUNDS = SKIP_SCALE_BOUNDS - # Only honor the noop flag when no activation or dbias is fused to match CUDA C++'s implementation - self.CHECK_NOOP_FLAG: cutlass.const_expr = ( - not self.cfg.WITH_ACT and not self.cfg.WITH_DACT and not self.cfg.WITH_DBIAS - ) - cast_dbias_only = cfg.WITH_DBIAS and not cfg.WITH_DACT and not cfg.WITH_ACT - # Use a different tile size for dbias only config - # No matter what tile size we use, each thread always handles a (1, MXFP8_BLOCK_SCALING_SIZE) chunk - if cast_dbias_only: - self._NUM_TILES = 4 # Each CTA handles 4 tiles stacked vertically - self._THREADS_PER_CTA = 128 - else: - self._NUM_TILES = 2 # Each CTA handles 2 tiles stacked vertically - self._THREADS_PER_CTA = 64 - # Each thread handles a (1, MXFP8_BLOCK_SCALING_SIZE) chunk - self._TILE_COLS = self._THREADS_PER_CTA - self._TILE_ROWS = MXFP8_BLOCK_SCALING_SIZE - self._NUM_WARPS = self._THREADS_PER_CTA // 32 - # We prefer to do dbias reduction in colwise which is easier (no cross-thread reduction needed). - # Only do rowwise reduction when we don't quantize columnwisely when WITH_DBIAS is True. - - # If columnwise is not quantized, and dbias is fused, we will do a columnwise reduction only pass - # (may also cache dact values) without quantization because doing dbias in rowwise requires a SMEM shuffle - # which is even slower - self.COLWISE_DBIAS_REDUCTION_ONLY = cfg.WITH_DBIAS and not cfg.COLWISE - # If columnwise is quantized and dbias is fused, do dbias reduction in colwise which is easier - # (no cross-thread reduction needed). - self.DBIAS_REDUCTION_IN_COLWISE = cfg.WITH_DBIAS and cfg.COLWISE - # Never do dbias reduction in rowwise because it is slow. Leave this so in case it's enabled in the future - self.DBIAS_REDUCTION_IN_ROWWISE = False - - # Cache activation in-place in the SMEM input tile when we process both rowwise and colwise passes - # so the activation is only computed once in the direction we favor (columnwise) and the other direction (rowwise) - # reads the cached value instead of recomputing it. - # Note: if activation is relu, there is no standalong relu applied because it's already fused into `cvt.rn.satfinite` - # so it should be treated as "no activation" - self.CACHE_ACTIVATION = ( - (cfg.WITH_ACT or cfg.WITH_DACT) - and cfg.ROWWISE - # Always cache activation in the colwise quantization or dbias reduction only pass - 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__( - self, - mX: cute.Tensor, # Input tensor to quantize - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], # Rowwise output and scale tensors - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], # Colwise output and scale tensors - mAmax: Optional[cute.Tensor], # Global amax accumulator, only used when WITH_AMAX is True - mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device - mDActInput: Optional[ - cute.Tensor - ], # Activation input for activation derivative fusion, only used when WITH_DACT is True - mWorkspace: Optional[ - cute.Tensor - ], # Workspace for the dbias reduction, only used when WITH_DBIAS is True - stream: CUstream, - ): - if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): - cute.printf(f"[CuTeDSL] MXFP8QuantizeKernel.__call__() with config: {self.cfg}\n") - - M = mX.shape[0] - N = mX.shape[1] - cfg = self.cfg - - # If WITH_GEMM_SWIZZLED_SCALES is enabled, the output must satisfy cublas's swizzled layout - # This is expressed as a CuTe layout applied to the output tensor so it can be transparent throughout the kernel implementation. - # See https://docs.nvidia.com/cuda/cublas/#d-block-scaling-factors-layout for more details. - if cutlass.const_expr(cfg.WITH_GEMM_SWIZZLED_SCALES): - mS_row, mS_col = derive_swizzled_scale_layout( - M, N, cfg.ROWWISE, cfg.COLWISE, mS_row, mS_col - ) - - # We have 2 stages in our pipeline where each stage loads / computes a (TILE_Y, TILE_X) tile - smem_tile_layout = cute.make_ordered_layout( - (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) - ) - cta_tiler = (self._TILE_ROWS, self._TILE_COLS) - - # Input TMA atoms - op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() - tma_atom, tma_src = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_load, - mX, - smem_tile_layout, - cta_tiler, - num_multicast=1, - ) - - # Activation input TMA atoms for activation derivative fusion - tma_atom_act = None - tma_src_act = None - if cutlass.const_expr(cfg.WITH_DACT): - tma_atom_act, tma_src_act = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_load, - mDActInput, - smem_tile_layout, - cta_tiler, - num_multicast=1, - ) - - # Output TMA atoms - op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() - out_smem_layout = cute.make_ordered_layout((self._TILE_ROWS, self._TILE_COLS), order=(1, 0)) - tma_atom_out_row = None - tma_dst_out_row = None - tma_atom_out_col = None - tma_dst_out_col = None - if cutlass.const_expr(cfg.ROWWISE): - 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, - ) - if cutlass.const_expr(cfg.COLWISE): - 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 = [ - cute.ceil_div(Int32(N), self._TILE_COLS), - cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES), - ] - block = [ - self._THREADS_PER_CTA, - ] - - self.kernel( - mX, - mS_row, - mS_col, - mAmax, - mNoop, - mWorkspace, - mX.element_type, - tma_atom, - tma_src, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - tma_atom_act, - tma_src_act, - ).launch( - grid=grid, - block=block, - stream=stream, - ) - - @cute.kernel - def kernel( - self, - mX: cute.Tensor, - mS_row: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, - mWorkspace: Optional[cute.Tensor], - dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - tma_atom: cute.CopyAtom, - tma_src: cute.Tensor, # Input TMA atoms - tma_atom_out_row: Optional[cute.CopyAtom], - tma_dst_out_row: Optional[cute.Tensor], # Rowwise output TMA atoms - tma_atom_out_col: Optional[cute.CopyAtom], - tma_dst_out_col: Optional[cute.Tensor], # Colwise output TMA atoms - tma_atom_act: Optional[cute.CopyAtom], - tma_src_act: Optional[ - cute.Tensor - ], # Activation derivative TMA atoms, None unless WITH_DACT - ): - """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" - - skip_execution = False - if cutlass.const_expr(self.CHECK_NOOP_FLAG): - skip_execution = noop_flag_is_set(mNoop) - if not skip_execution: - self._kernel_main( - mX, - mS_row, - mS_col, - mAmax, - mWorkspace, - dtype, - tma_atom, - tma_src, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - tma_atom_act, - tma_src_act, - ) - - @cute.jit - def _kernel_main( - self, - mX: cute.Tensor, - mS_row: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - mAmax: Optional[cute.Tensor], - mWorkspace: Optional[cute.Tensor], - dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - tma_atom: cute.CopyAtom, - tma_src: cute.Tensor, # Input TMA atoms - tma_atom_out_row: Optional[cute.CopyAtom], - tma_dst_out_row: Optional[cute.Tensor], # Rowwise output TMA atoms - tma_atom_out_col: Optional[cute.CopyAtom], - tma_dst_out_col: Optional[cute.Tensor], # Colwise output TMA atoms - tma_atom_act: Optional[cute.CopyAtom], - tma_src_act: Optional[ - cute.Tensor - ], # Activation derivative TMA atoms, None unless WITH_DACT - ): - cfg = self.cfg - - if cutlass.const_expr(cfg.ROWWISE): - mS_row = cute.zipped_divide( - mS_row, (self._TILE_ROWS, self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) - ) - if cutlass.const_expr(cfg.COLWISE): - mS_col = cute.zipped_divide( - mS_col, (self._TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, self._TILE_COLS) - ) - - # Allocate shared memory for the input and rowwise / columnwise outputs - FP8_DTYPE = cfg.FP8_DTYPE - - if cutlass.const_expr(cfg.ROWWISE and cfg.COLWISE): - - @cute.struct - class SharedStorage: - mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] - 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, - ] - sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] - - elif cutlass.const_expr(cfg.ROWWISE and not cfg.COLWISE): - - @cute.struct - class SharedStorage: - mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] - 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, - ] - sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] - - else: - - @cute.struct - class SharedStorage: - mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] - sX: cute.struct.Align[ - cute.struct.MemRange[ - 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, - ] - sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] - - smem = cutlass.utils.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 - sX = storage.sX.get_tensor( - cute.make_layout( - ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), - stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), - ) - ) - if cutlass.const_expr(cfg.ROWWISE): - sO_row = storage.sO_row.get_tensor( - cute.make_layout( - ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), - stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), - ) - ) - if cutlass.const_expr(cfg.COLWISE): - sO_col = storage.sO_col.get_tensor( - cute.make_layout( - ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), - stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), - ) - ) - - # Allocate shared memory for the activation input used for the activation derivative fusion. - if cutlass.const_expr(cfg.WITH_DACT): - - @cute.struct - class DactStorage: - sActInput: cute.struct.Align[ - cute.struct.MemRange[ - dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - - dact_storage = smem.allocate(DactStorage) - # Apply the same layout as the input - sActInput = dact_storage.sActInput.get_tensor( - cute.make_layout( - ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), - stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), - ) - ) - - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - - # Prefetch TMA descriptors - if warp_idx == 0: - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom) - if cutlass.const_expr(cfg.WITH_DACT): - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_act) - - tidx, _, _ = cute.arch.thread_idx() - bidx, bidy, _ = cute.arch.block_idx() - - # Only warp 0 is the producer (issues TMA) - producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) - # Every warp is the consumer (reads the data loaded by TMA) - consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._NUM_WARPS) - - # Bytes transferred per TMA copy: one (TILE_Y, TILE_X) tile of dtype. - tx_count = self._TILE_ROWS * self._TILE_COLS * dtype.width // 8 - # dact loads two tiles (grad + act_input) under the same per-stage barrier, - # so the barrier must expect both copies' bytes. - if cutlass.const_expr(cfg.WITH_DACT): - tx_count *= 2 - - 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, # single-CTA, no cluster/multicast - ) - - prod_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self._NUM_STAGES - ) - cons_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self._NUM_STAGES - ) - - 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), - ) - - # 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)) - - # Partition sX/gX for the TMA atom (single-CTA, no cluster/multicast). - tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( - tma_atom, - 0, # Use the only CTA to do the TMA copy - cute.make_layout(1), # This cluster only has 1 CTAs - sX, - gX_tiled, - ) - - # If WITH_DACT, partition the activation input for TMA as well in the same way - if cutlass.const_expr(cfg.WITH_DACT): - gA_tiled = cute.zipped_divide(tma_src_act, (self._TILE_ROWS, self._TILE_COLS)) - tXsA, tXgA = cute.nvgpu.cpasync.tma_partition( - tma_atom_act, - 0, - cute.make_layout(1), - sActInput, - gA_tiled, - ) - - # Partitioning for rowwise / columnwise outputs - if cutlass.const_expr(cfg.ROWWISE): - 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, - ) - if cutlass.const_expr(cfg.COLWISE): - 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, - ) - - # Ensure barrier init is visible to all threads before the pipeline is used. - 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: - mainloop_pipeline.producer_acquire(prod_state) - tile_y = bidy * self._NUM_TILES + s - cute.copy( - tma_atom, - tXgX[(None, (tile_y, bidx))], - tXsX[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - if cutlass.const_expr(cfg.WITH_DACT): - cute.copy( - tma_atom_act, - tXgA[(None, (tile_y, bidx))], - tXsA[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - mainloop_pipeline.producer_commit(prod_state) - prod_state.advance() - - # Per-thread amax accumulator - if cutlass.const_expr(cfg.WITH_AMAX): - per_thread_amax = Float32(0.0) - - # Prepare thread-level register accumulators for rowwise dbias reduction. - # Each thread will process two (1, MXFP8_BLOCK_SCALING_SIZE) rows in two stages, and in each stage the thread will add the - # (after dact applied) value to this register array with the same shape so it carries the the two stages' partial sum. - # Then it will be written to a SMEM buffer to let the whole CTA do the reduction separately to yield - # the final (1, TILE_X) dbias workspace output. - rowwise_dbias_acc = None - if cutlass.const_expr(self.DBIAS_REDUCTION_IN_ROWWISE): - rowwise_dbias_acc = cute.make_rmem_tensor( - layout_or_shape=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(1,)), - dtype=Float32, - ) - # Zero the accumulator registers. - for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - rowwise_dbias_acc[c] = Float32(0.0) - block_dbias = Float32(0.0) - # Prepare thread-level register for columnwise dbias reduction. - # Each thread will process two (MXFP8_BLOCK_SCALING_SIZE, 1) columns in two stages, and in each stage the thread will reduce the - # (after dact applied) column to (1,) and add to this register. - # Then this partial sum scalar will be written to the GMEM workspace buffer directly. - if cutlass.const_expr(self.DBIAS_REDUCTION_IN_COLWISE or self.COLWISE_DBIAS_REDUCTION_ONLY): - block_dbias = Float32(0.0) - - # 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 - sX_tile = sX[(None, stage_idx)] - # Also fetch the activation input if WITH_DACT - sActInput_tile = None - if cutlass.const_expr(cfg.WITH_DACT): - sActInput_tile = sActInput[(None, stage_idx)] - # Each CTA handles `NUM_TILES` tiles stacked vertically, so tile_idx_x is just the block index along X dimension - # 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() - 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). - # Each CTA owns NUM_TILES consecutive row-tiles of GRID. cute - # auto-decomposes the flat row coord `bidy * NUM_TILES + tile_idx` - # onto GRID's hierarchical row modes — which is the - # (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, - ) - - if cutlass.const_expr(self.AMAX_FROM_ROWWISE): - per_thread_amax = cute.arch.fmax(per_thread_amax, amax_r) - - # Make the shared-memory writes visible to the TMA's async proxy before the TMA reads them. - 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) - - # Warp 0 issues TMA copy to write the quantized output tile from shared memory to global memory and then commits - if warp_idx == 0: - tile_y = bidy * self._NUM_TILES + tile_idx - if cutlass.const_expr(cfg.ROWWISE): - cute.copy( - tma_atom_out_row, - tXsO_row[(None, stage_idx)], - tXgO_row[(None, (tile_y, bidx))], - ) - if cutlass.const_expr(cfg.COLWISE): - cute.copy( - tma_atom_out_col, - tXsO_col[(None, stage_idx)], - tXgO_col[(None, (tile_y, bidx))], - ) - cute.arch.cp_async_bulk_commit_group() - - cons_state.advance() - - # The pipeline is no longer fully filled after we consume this tile, so we fetch a new tile to fill the pipeline. - # The next _NUM_STAGES-1 tiles are already in-flight, so the next tile to fetch is after _NUM_STAGES tiles. - if warp_idx == 0: - next_tile_idx = tile_idx + self._NUM_STAGES - if next_tile_idx < num_tiles: - mainloop_pipeline.producer_acquire(prod_state) - tile_y = bidy * self._NUM_TILES + next_tile_idx - cute.copy( - tma_atom, - tXgX[(None, (tile_y, bidx))], - tXsX[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - if cutlass.const_expr(cfg.WITH_DACT): - cute.copy( - tma_atom_act, - tXgA[(None, (tile_y, bidx))], - tXsA[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - mainloop_pipeline.producer_commit(prod_state) - prod_state.advance() - # End of the main pipeline loop - - # Complete the cross-thread dbias reduction after each thread has its own per-thread partial sum after the rowwise quantization. - if cutlass.const_expr(self.DBIAS_REDUCTION_IN_ROWWISE): - # If we do the dbias reduction in the rowwise pass, each thread will have a (1, MXFP8_BLOCK_SCALING_SIZE) partial sum - # and we need to write these to a SMEM buffer and let each thread reduce it in the columnwise direction - block_dbias = self._dbias_reduction_rowwise_epilouge(smem, tidx, rowwise_dbias_acc) - - # Write the per-tile reduced dbias to the global workspace. - if cutlass.const_expr(cfg.WITH_DBIAS): - dbias_col = bidx * self._TILE_COLS + tidx - if dbias_col < N: - mWorkspace[(bidy, dbias_col)] = block_dbias - - if cutlass.const_expr(cfg.WITH_AMAX): - 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) - - @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.utils.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 = ( - self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE * (MXFP8_BLOCK_SCALING_SIZE + 1) - ) - - # Allocate the SMEM buffer that all threads use to reduce the two-stage partial sum (per thread) to the - # partial sum (per block). - @cute.struct - class DbiasStorage: - sDbias: cute.struct.MemRange[Float32, self._TILE_ROWS * DBIAS_BUFF_WIDTH] - - dbias_storage = smem.allocate(DbiasStorage) - sDbias = dbias_storage.sDbias.get_tensor( - cute.make_layout((self._TILE_ROWS, self._TILE_COLS), stride=(DBIAS_BUFF_WIDTH, 1)), - ) - _, tv_layout_dbias_write = cute.make_layout_tv( - thr_layout=cute.make_layout( - (self._TILE_ROWS, self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE), - stride=(self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE, 1), - ), - val_layout=cute.make_layout( - (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) - ), - ) - sDbias_write = cute.composition(sDbias, tv_layout_dbias_write) - # 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 - # The offset this thread should start reading from based on what's its first bank to access. - offset = bank_group * self._PACK_SIZE - for w in cutlass.range_constexpr( - self._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): - # 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] - 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( - thr_layout=cute.make_layout((1, self._TILE_COLS), stride=(self._TILE_COLS, 1)), - val_layout=cute.make_layout((self._TILE_ROWS, 1), stride=(1, 1)), - ) - sDbias_reduce = cute.composition(sDbias, tv_layout_dbias_reduce) - # make_layout_tv yields a (thread, value) layout: thread=tidx -> column tidx, - # value=i -> row i. So index [tidx, i] (thread first), summing the column's rows. - block_dbias = Float32(0.0) - for i in cutlass.range_constexpr(self._TILE_ROWS): - block_dbias += sDbias_reduce[tidx, i] - return block_dbias - - @cute.jit - def _amax_epilogue( - self, - sAmax: cute.Tensor, - mAmax: cute.Tensor, - tidx: Int32, - warp_idx: Int32, - per_thread_amax: Float32, - ): - # Reduce and get the per-warp amax. - warp_amax = cute.arch.warp_redux_sync(per_thread_amax, kind="fmax") - # Write the per-warp amax to shared memory - lane_idx = tidx % 32 - if lane_idx == 0: - sAmax[warp_idx] = warp_amax - cute.arch.sync_threads() - if tidx == 0: - cta_amax = Float32(0.0) - # The first thread reduces all the per-warp amax to the per-CTA amax - for w in cutlass.range_constexpr(self._NUM_WARPS): - cta_amax = cute.arch.fmax(cta_amax, sAmax[w]) - amax_i32 = cute.make_tensor( - cute.recast_ptr(mAmax.iterator, dtype=Int32), - cute.make_layout(1), - ) - # The first thread updates the global amax with an atomic max on the bitcasted float value - cute.arch.atomic_max( - 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, - ) - - -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/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py new file mode 100644 index 00000000000..0df60e0714e --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -0,0 +1,600 @@ +# 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 +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, + device_compute_capability, +) +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 + +# 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__ + + 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, + mDActInput: Optional[cute.Tensor], + ) -> 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, + 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 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_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 + + self.cfg = cfg + # These activation functions satisfy f(0) = 0 so we don't need to mask with OOB regions + # (zeros filled by TMA are still zeros without applying activation to them) + # Note that as the time of writing, all activations supported by MXFP8 satisfy this property, so !ACT_NEED_MASKING is never reachable. + # This is only a precaution for future use cases if someone wants to add a new activation function. + self.ACT_NEED_MASKING = cfg.WITH_ACT and cfg.ACTIVATION not in ( + "relu", + "gelu", + "silu", + "qgelu", + "srelu", + ) + # Instantiate all possible kernels at compile time, + # and we will pick the right one at runtime based on the input shape and config. + self.general_skip_input_masking_kernel = MXFP8QuantizeGeneralG2SKernel( + cfg, SKIP_INPUT_MASKING=True + ) + self.general_skip_input_masking_and_scale_bounds_kernel = MXFP8QuantizeGeneralG2SKernel( + cfg, SKIP_INPUT_MASKING=True, SKIP_SCALE_BOUNDS=True + ) + # We only need to mask when WITH_ACT is enabled because with activations applied zeros filled by TMA affect block statistics + self.general_kernel = MXFP8QuantizeGeneralG2SKernel(cfg) if self.ACT_NEED_MASKING else None + + # Specialized 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 ( + MXFP8QuantizeRowwise2LaneG2RKernel, + MXFP8QuantizeBidimensionalG2RKernel, + MXFP8QuantizeRowwise1LaneG2RKernel, + MXFP8QuantizeBidimensionalG2SKernel, + ) + if kernel_cls.supports(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, + ): + 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, mDActInput + ): + dispatched_to_specialized = True + kernel( + 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 (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 + 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_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py new file mode 100644 index 00000000000..61598a7d37b --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -0,0 +1,919 @@ +# 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 Callable, Optional, Type + +import cutlass +from cutlass import cute +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 ( + abs_max_x2_bf16, + bf16_pair_magnitude, + bf16_mx_scale_reciprocal, + bf16_mx_scale_reciprocal_pair, + 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, +) + + +@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: 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[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 + + @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_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_pair, + rX_row_i32[2 * pack_idx + 1], + 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_pair)).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[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 + + 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 + + # 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( + 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( + 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( + sAmaxs_col_scratch: cute.Tensor, + # Constexprs + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], +) -> 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 + + 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), + which is the return value + """ + tidx, _, _ = cute.arch.thread_idx() + + TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD + + # 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_PAIRS = TILE_COLS // REDUCTION_THREADS_NUM // 2 + + # 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( + 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) + sS_col_i32 = cute.recast_tensor(sS_col, Int32) + + # 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_PAIRS, CTA_THREADS_Y), stride=(1, REDUCTION_i32_PAIRS) + ), + ) + # 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_PAIRS, 1), stride=(1, REDUCTION_i32_PAIRS)), + ) + + sS_col_i32 = cute.composition(sS_col_i32, tv_col_rcp_layout) + + # 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 + if is_reduction_thread: + 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_PAIRS), stride=(REDUCTION_i32_PAIRS, 1)), + ) + + # 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 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 pair + sS_col_i32[tidx, pair_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) + cute.arch.sync_threads() + + return sS_col + + +@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_X: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + 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_u32 = cute.recast_tensor(rS_col_rcp, Uint32) + + # 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) + # 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 + 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 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 _apply_colwise_scales( + 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: cutlass.Constexpr[Callable[..., Uint32]], + # Constexprs + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + 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) + + 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 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) + + +@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_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], + 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 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 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 + + # 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 + + SCALE_DTYPE = cutlass.Float8E8M0FNU + + tidx, _, _ = cute.arch.thread_idx() + + 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( + (ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD), stride=(ELEMENTS_X_PER_THREAD, 1) + ), + ) + + gX_thread = cute.composition(gX_tile, tv_layout)[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), DTYPE) + + 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(), + ), + # 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(), + FP8_DTYPE, + num_bits_per_copy=ELEMENTS_X_PER_THREAD * FP8_DTYPE.width, + ) + + 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): + cute.copy( + 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] + + # 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] + + row_scale_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8 + ) + + _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, + ) + + 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, + 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 MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): + """Load one CTA tile into registers and quantize it along both axes. + + 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 and N divisible by 256; writes packed or GEMM-swizzled + scales. + """ + + _CTA_THREADS_Y = 8 + _CTA_THREADS_X = 32 + + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + # BF16 packed math. + return ( + not cfg.WITH_AMAX + and not cfg.WITH_DBIAS + and cfg.DTYPE is cutlass.BFloat16 + and 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, + 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 ( + M % 32 == 0 + 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 + ) + + 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 + 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, # Unused: the dispatcher routes noop calls to the general kernel + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + M, N = mX.shape + 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( + mX, + mDActInput, + mO_row, + mS_row, + mO_col, + mS_col, + stream, + ELEMENTS_X_PER_THREAD=16, + CLUSTER_WIDTH=4, + ) + else: + self._launch( + mX, + mDActInput, + mO_row, + mS_row, + mO_col, + mS_col, + stream, + ELEMENTS_X_PER_THREAD=16, + CLUSTER_WIDTH=1, + ) + else: + self._launch( + mX, + mDActInput, + mO_row, + mS_row, + mO_col, + mS_col, + stream, + ELEMENTS_X_PER_THREAD=8, + CLUSTER_WIDTH=8, + ) + + @cute.jit + def _launch( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: cute.Tensor, + mS_col: cute.Tensor, + stream: CUstream, + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + CLUSTER_WIDTH: cutlass.Constexpr[int], + ): + 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) + 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 + 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] + + # 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) + 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) + 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, 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)), + ) + if cutlass.const_expr(self.cfg.ROWWISE) + else None + ) + gS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout((cute.size(mS_col, mode=[0]), 2048), stride=(2048, 1)), + ) + self.kernel( + gX, + gA, + gO_row, + gS_row, + gO_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) + elif N == 4096: + payload_layout = cute.make_layout((M, 4096), stride=(4096, 1)) + gX = cute.make_tensor(mX.iterator, payload_layout) + 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) + 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, 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)), + ) + if cutlass.const_expr(self.cfg.ROWWISE) + else None + ) + gS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout((cute.size(mS_col, mode=[0]), 4096), stride=(4096, 1)), + ) + self.kernel( + gX, + gA, + gO_row, + gS_row, + gO_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) + else: + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + 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( + mX, + mA, + mO_row, + gS_row, + mO_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) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mA: Optional[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], + ): + """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, + mA, + 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, + mA: Optional[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) + 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): + 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_col_tile = cute.local_tile(mS_col, COL_SCALE_TILER, cta_coord) + + cache_policy = create_l2_policy(True) + + 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, + 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 new file mode 100644 index 00000000000..682ba02c611 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py @@ -0,0 +1,334 @@ +# 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 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, + 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, +) + + +@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_mxfp8_gmem_tile_rowwise_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 MXFP8QuantizeRowwise1LaneG2RKernel(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 + + @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, + mDActInput: Optional[cute.Tensor], + ) -> 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, MXFP8QuantizeRowwise1LaneG2RKernel._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, + # 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, # 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, + ): + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + f"[CuTeDSL] MXFP8QuantizeRowwise1LaneG2RKernel.__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, + ).launch(grid=grid, block=block, stream=stream) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + ): + """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) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + ): + """Device entry for the specialized rowwise-only cast kernel (vectorized global loads/stores, no TMA).""" + bidx, bidy, _ = cute.arch.block_idx() + M = mX.shape[0] + N = mX.shape[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)) + + quantize_mxfp8_gmem_tile_rowwise_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. + 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, + ) 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..03d5d510ab8 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -0,0 +1,689 @@ +# 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, Type + +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 + +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, + 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, + mul_bf16x2x2_cvt_bf16x4_to_fp8x4, +) +from .quantize_mxfp8_common import ( + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + SUPPORTED_ACTIVATIONS, + SUPPORTED_DACTIVATIONS, + derive_swizzled_scale_layout, +) + + +@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, + 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], + 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, + cta_row: Int32, + cta_col: Int32, + num_elems: 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), 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 + # 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 * DTYPE.width, + max_num_bits_per_copy_ld_global(), + ) + load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + 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 + 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, + num_bits_per_copy=ELEMENTS_PER_THREAD * FP8_DTYPE.width, + ) + + 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): + # 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): + 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 + # 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 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: + """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_mx_block_amax(rX_tile_i32) + 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 + # 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) + 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 MXFP8QuantizeRowwise2LaneG2RKernel(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 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. + """ + + _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 + + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + # BF16 packed math. + return ( + not cfg.WITH_AMAX + and not cfg.WITH_DBIAS + and cfg.DTYPE is cutlass.BFloat16 + 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, + 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. + 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 act_input_aligned + and scale_stride_matches + and mNoop.toint() == 0 + ) + + 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, # Unused: the dispatcher routes noop calls to the general kernel + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + 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)) + 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, mDActInput, mO_row, mS_row, stream, num_elems) + else: + 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, + 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, + mA, + mO_row, + mS_row, + stream, + num_elems, + 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( + mX, + mA, + mO_row, + mS_row, + stream, + num_elems, + 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( + mX, + mA, + mO_row, + mS_row, + stream, + num_elems, + 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( + mX, + mA, + mO_row, + mS_row, + stream, + num_elems, + 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 + def _launch( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], + mO: cute.Tensor, + mS: cute.Tensor, + stream: CUstream, + num_elems: Int64, + 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 <= 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 or num_elems != Int64(M) * N + grid = Int64(grid_rows) * grid_cols + 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] + 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: + 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: + 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( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], + 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], + 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.""" + self._kernel_main( + 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=CHECK_BOUNDS, + STASH_SCALE_TO_SMEM=STASH_SCALE_TO_SMEM, + ) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], + 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], + 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. + 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 + cta_coord = (cta_row, bidx) + 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 + # 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): + gS, _ = derive_swizzled_scale_layout(M, N, True, False, mS, None) + + 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) + 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(L2_CACHED_CTA_PERCENT > 0): + fraction = Float32(0.0) + # 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) + 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) + # 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)) + + 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: + 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, + 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], + 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, + ) 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 new file mode 100644 index 00000000000..1192d98cb3a --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py @@ -0,0 +1,737 @@ +# 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 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 ( + 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, +) + + +@cute.jit +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 + 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 MXFP8QuantizeBidimensionalG2SKernel(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 + + @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, + mDActInput: Optional[cute.Tensor], + ) -> Boolean: + 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 + + 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, # Unused: the dispatcher routes noop calls to the general kernel + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + stream: CUstream, + ): + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + "[CuTeDSL] MXFP8QuantizeBidimensionalG2SKernel.__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, + 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, + 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. 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( + 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.memory.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 thiHs stage + quantize_mxfp8_smem_tile_bidimensional( + 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_g2s_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py new file mode 100644 index 00000000000..9801d82a259 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py @@ -0,0 +1,1400 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MXFP8 quantization kernel implemented in CuTeDSL. + +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). + +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 + BUFF_DIM_Y = 32, BUFF_DIM_X = 64, STAGES = 2 + MXFP8_BLOCK_SCALING_SIZE = 32 (elements per MXFP8 scaling block) + +Grid: (ceil(N / 64), ceil(M / 64)) +Each block processes a 64x64 chunk in 2 stages of 32x64 tiles loaded into +shared memory. +""" + +# Local @cute.struct classes are SMEM-layout descriptors that need no docstrings. +# pylint: disable=missing-class-docstring + +from typing import Optional, Type + +import cutlass +from cutlass import cute +from cutlass import pipeline +from cutlass import Float32, Int32, Int64, Uint16, Uint32, Uint8 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + is_packed16, + abs_max_x2_bf16, + abs_max_x2_f16, + max_x2_bf16, + max_x2_f16, + max_scalar_bf16, + max_scalar_f16, + abs_max_scalar_bf16, + abs_max_scalar_f16, + to_f32_bf16, + to_f32_f16, + x2_lo_to_f32_bf16, + x2_lo_to_f32_f16, + x2_hi_to_f32_bf16, + x2_hi_to_f32_f16, + truncate_f32_bf16, + truncate_f32_f16, + fabs_f32, + exp2f_rcp, + pack_f32x2, + unpack_i64_to_i32x2, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + get_cvt_f32x2_to_fp8x2_func, + cvt_f32_to_fp8e8m0fnu, + mul_f32x2_cvt_f32x4_to_fp8x4, + mul_f32x2_cvt_packed16x4_to_fp8x4, +) +from .quantize_mxfp8_common import ( + 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 +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) + mS_row_stage: cute.Tensor, # rowwise scale tensor (1D swizzled, or 2D linear) + MAX_NORM_RCP: cutlass.Constexpr[float], + tile_row_start: Int32, # global row index of this stage's row 0 + tile_col_start: Int32, # global col index of this CTA's col 0 + M: Int32, + N: Int32, + ACTIVATION: cutlass.Constexpr[str | None], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + TILE_X: cutlass.Constexpr[int], + TILE_Y: cutlass.Constexpr[int], + SKIP_INPUT_MASKING: cutlass.Constexpr[bool] = False, + SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool] = False, + WITH_ACT: cutlass.Constexpr[bool] = False, + WITH_DACT: cutlass.Constexpr[bool] = False, + WITH_DBIAS: cutlass.Constexpr[bool] = False, + dbias_acc: Optional[cute.Tensor] = None, # only needed when WITH_DBIAS is True +): + """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) + + _, tv_layout = 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, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), + ) + + sX_tv = cute.composition(sX_tile, tv_layout) + sO_tv = cute.composition(sO_row_tile, tv_layout) + + # I/O Elements that belong to this thread + sX_thread = sX_tv[tidx, None] # shape (32,) bf16 + sO_thread = sO_tv[tidx, None] # shape (32,) fp8 + + sO_thread_u32_ptr = cute.recast_ptr(sO_thread.iterator, dtype=Uint32) + sO_thread_u32 = cute.make_tensor( + sO_thread_u32_ptr, + cute.make_layout( + (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) + ), # 1 uint32 is 4 fp8 elements + ) + + # PTX allows to fuse relu activation in `cvt.rn.satfinite` unless we need to reduction for dbias + FUSE_RELU = cutlass.const_expr(ACTIVATION == "relu") and not WITH_DBIAS + # For this fast path we can read in pack of 2 instead of reading individual f16 / bf16 element, + # and keep the elements in half precision instead of upcasting every one of them to f32. + USE_HALF_PRECISION = is_packed16(DTYPE) and (ACTIVATION is None or FUSE_RELU) + + amax_r = Float32(0.0) + + # Each thread start reading from the specfic bank based on its thread ID so they can do their best to access different banks + # to avoid bank conflict. + bank_group = (tidx % THREADS_PER_WARP) // THREADS_PER_BANK + # The offset this thread should start reading from based on what's its first bank to access. + offset = bank_group * PACK_SIZE + if cutlass.const_expr(USE_HALF_PRECISION): + # If no activation, f16 / bf16 and rowwise quantization, we can read 2 f16 / bf16 at once in a pack + # and use max.xorsign.abs.f16x2 / max.xorsign.abs.bf16x2 to compute + max_x2 = max_x2_f16 if DTYPE is cutlass.Float16 else max_x2_bf16 + 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 + sX_thread_rw_i64 = cute.make_tensor( + cute.recast_ptr(sX_thread.iterator, dtype=Int64), + cute.make_layout( + (1, MXFP8_BLOCK_SCALING_SIZE // 4), stride=(0, 1) + ), # 1 int64 is 4 fp16/bf16 elements + ) + # Each wave reads its 4 elements (PACK_SIZE) as one 8-byte vectorized load + in_r = [[None, None] for _ in range(WAVES)] + for w in cutlass.range_constexpr(WAVES): + idx = (w + offset // 4) % (MXFP8_BLOCK_SCALING_SIZE // 4) + in_r[w][0], in_r[w][1] = unpack_i64_to_i32x2(sX_thread_rw_i64[0, idx]) + + if cutlass.const_expr(WITH_DBIAS): + for w in cutlass.range_constexpr(WAVES): + dbias_acc[w * PACK_SIZE + 0] += x2_lo_to_f32(in_r[w][0]) + dbias_acc[w * PACK_SIZE + 1] += x2_hi_to_f32(in_r[w][0]) + dbias_acc[w * PACK_SIZE + 2] += x2_lo_to_f32(in_r[w][1]) + dbias_acc[w * PACK_SIZE + 3] += x2_hi_to_f32(in_r[w][1]) + + amax_2x = in_r[0][0] + # Each wave will use max.xorsign.abs.f16x2 or max.xorsign.abs.bf16x2 to compare 2 packed elements in parallel + for w in cutlass.range_constexpr(WAVES): + if cutlass.const_expr(FUSE_RELU): + # Skip the first iteration because we assigned in_r[0][0] to amax_2x already + if cutlass.const_expr(w > 0): + # If we fuse relu then we don't want to do abs since negative value will be set to 0 and they will lose comparison automatically + amax_2x = max_x2(amax_2x, in_r[w][0]) + amax_2x = max_x2(amax_2x, in_r[w][1]) + else: + # Skip the first iteration because we assigned in_r[0][0] to amax_2x already + if cutlass.const_expr(w > 0): + amax_2x = abs_max_x2(amax_2x, in_r[w][0]) + amax_2x = abs_max_x2(amax_2x, in_r[w][1]) + if cutlass.const_expr(FUSE_RELU): + # Compare the 2 packed max without abs + amax_r = cute.arch.fmax( + x2_lo_to_f32(amax_2x), + x2_hi_to_f32(amax_2x), + ) + # For relu the max is at least 0 + amax_r = cute.arch.fmax(amax_r, Float32(0.0)) + else: + # Compare the 2 packed abs max + amax_r = cute.arch.fmax( + fabs_f32(x2_lo_to_f32(amax_2x)), + fabs_f32(x2_hi_to_f32(amax_2x)), + ) + else: + # Since we need to do computation on individual f16 / bf16 elements, we can't read in pack + sX_thread_rw = cute.make_tensor( + sX_thread.iterator, + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), + ) + + if cutlass.const_expr(WITH_DACT): + # Backward: out = grad · act'(act_input). sX is grad, sA is act_input. + dop = SUPPORTED_DACTIVATIONS[ACTIVATION] + sA_thread = cute.composition(sActInput_tile, tv_layout)[tidx, None] + sA_thread_rw = cute.make_tensor( + sA_thread.iterator, + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), + ) + elif cutlass.const_expr(WITH_ACT): + op = SUPPORTED_ACTIVATIONS[ACTIVATION] + + if cutlass.const_expr(is_packed16(DTYPE) and ACTIVATION is not None): + truncate_f32 = truncate_f32_f16 if DTYPE is cutlass.Float16 else truncate_f32_bf16 + + # Each wave we read PACK_SIZE elements, and we have WAVES waves, so we read WAVES * PACK_SIZE (= MXFP8_BLOCK_SCALING_SIZE) elements in total. + in_r = [[None] * PACK_SIZE for _ in range(WAVES)] + # Don't repeat these OOB masking computation in the loop + thread_row_start = tile_row_start + tidx // CTA_THREADS_X + thread_col_start = tile_col_start + (tidx % CTA_THREADS_X) * MXFP8_BLOCK_SCALING_SIZE + thread_row_oob = thread_row_start >= M + # PACK_SIZE must divide every term of `thread_col_start + start` and N itself, or a + # pack could straddle the N boundary and the check below would be wrong. TILE_X is a + # tunable, so it is asserted rather than assumed. + assert MXFP8_BLOCK_SCALING_SIZE % PACK_SIZE == 0 + assert SYM_N_DIVISIBILITY % PACK_SIZE == 0 + assert TILE_X % PACK_SIZE == 0 + for w in cutlass.range_constexpr(WAVES): + start = (w * PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE + # thread_col_start + start >= N means "if the first element in this pack is OOB" + # If the first element is OOB, then (obviously) elements after it within this pack must be OOB + # If the first element is not OOB, then elements after it within this pack must be not OOB either. + # Proof: + # (1) when the first element is not OOB, we have thread_col_start + start < N, + # which implies N - (thread_col_start + start) > 0 + # (2) with the assertions above: + # - N is divisible by SYM_N_DIVISIBILITY, so N is also divisible by PACK_SIZE + # - thread_col_start = tile_col_start + (tidx % CTA_THREADS_X) * MXFP8_BLOCK_SCALING_SIZE is divisible by PACK_SIZE + # because TILE_X and MXFP8_BLOCK_SCALING_SIZE are both divisible by PACK_SIZE + # - start = (w * PACK_SIZE + bank_group * PACK_SIZE) % MXFP8_BLOCK_SCALING_SIZE is also divisible by PACK_SIZE + # because MXFP8_BLOCK_SCALING_SIZE is a multiple of PACK_SIZE + # so N - (thread_col_start + start) must be divisible by PACK_SIZE + # Therefore, N - (thread_col_start + start) >= PACK_SIZE because PACK_SIZE is the minimal number + # for (1) and (2) to hold, which implies thread_col_start + start + PACK_SIZE -1 < N + # Hence, we only need to check if the first element is OOB to determine if the whole pack is OOB + thread_col_oob = thread_col_start + start >= N + for i in cutlass.range_constexpr(PACK_SIZE): + x = Float32(sX_thread_rw[0, start + i]) + if cutlass.const_expr(WITH_DACT): + # out = grad · act'(act_input) + x = x * dop(Float32(sA_thread_rw[0, start + i])) + # If IS_ACT, apply activation function to x in f32 + elif cutlass.const_expr(WITH_ACT): + # If it's relu, we can handle it later + if not cutlass.const_expr(FUSE_RELU): + x = op(x) + if not cutlass.const_expr(SKIP_INPUT_MASKING): + # If the input shape is not divisible by the tile size, + # TMA would zero-fills the input tile outside its logical MxN bounds. + # This is fine for non-activation cases, but for activation cases, + # op(0) might not be 0 which will pollute the amax and dbias. + # So we must manually mask the OOB region here. + if thread_row_oob or thread_col_oob: + x = Float32(0.0) + # Accumulate to the per-thread dbias register buffer for this tile if WITH_DBIAS + if cutlass.const_expr(WITH_DBIAS): + # dbias_acc is register buffer so we can just write without bank conflict + dbias_acc[w * PACK_SIZE + i] += x + # If 16-bit input with activation, truncate to IType + if cutlass.const_expr(is_packed16(DTYPE) and ACTIVATION is not None): + x = truncate_f32(x) + in_r[w][i] = x + if cutlass.const_expr(FUSE_RELU): + amax_r = cute.arch.fmax( + amax_r, x + ) # For relu cases, we don't need abs since negative values will be 0 so they lose comparison automatically + else: + amax_r = cute.arch.fmax(amax_r, fabs_f32(x)) + if cutlass.const_expr(FUSE_RELU): + amax_r = cute.arch.fmax(amax_r, Float32(0.0)) # If relu, the amax is at least 0 + + biased_exp_r = cvt_f32_to_fp8e8m0fnu(amax_r * MAX_NORM_RCP) + + # mS_row_stage has logical shape (32, 2) and we have 64 threads where each is mapped to one scale factor + # The TV layout is equivalent to TV layout with thr_layout=(32, 2):(2, 1), val_layout=(1,) + # but it's too trival so let's just index it directly without using layout + # Note this is the logical layout, which is on top of the swizzled / non-swizzled scale factor layout + # that mappes the logical index to the physical offset + + # 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 cutlass.const_expr(SKIP_SCALE_BOUNDS): + mS_row_stage[(tidx // CTA_THREADS_X, tidx % CTA_THREADS_X)] = biased_exp_r + else: + scale_row = tile_row_start + tidx // CTA_THREADS_X + scale_col_first_elt = tile_col_start + (tidx % CTA_THREADS_X) * MXFP8_BLOCK_SCALING_SIZE + if scale_row < M and scale_col_first_elt < N: + mS_row_stage[(tidx // CTA_THREADS_X, tidx % CTA_THREADS_X)] = biased_exp_r + + inv_scale_r = exp2f_rcp(biased_exp_r) # f32 reciprocal of the scale + scale_2x = pack_f32x2(inv_scale_r, inv_scale_r) + if cutlass.const_expr(USE_HALF_PRECISION): + mul_cvt_x4_func = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, FP8_DTYPE, FUSE_RELU) + else: + mul_cvt_x4_func = mul_f32x2_cvt_f32x4_to_fp8x4(FP8_DTYPE, FUSE_RELU) + + for w in cutlass.range_constexpr(WAVES): + idx = (w * 4 + offset) % MXFP8_BLOCK_SCALING_SIZE + idx = idx // 4 + if cutlass.const_expr(USE_HALF_PRECISION): + # Convert 2 packed f16/bf16 pairs to 4 fp8 in one fused op + sO_thread_u32[idx] = mul_cvt_x4_func(in_r[w][0], in_r[w][1], scale_2x) + else: + # Convert 4 f32 to 4 fp8 in one fused op + sO_thread_u32[idx] = mul_cvt_x4_func( + in_r[w][0], in_r[w][1], in_r[w][2], in_r[w][3], scale_2x + ) + + return amax_r + + +@cute.jit +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) + mS_col_stage: cute.Tensor, # colwise scale tensor (1D swizzled, or 2D linear) + MAX_NORM_RCP: cutlass.Constexpr[float], + tile_row_start: Int32, # global row index of this stage's row 0 + tile_col_start: Int32, # global col index of this CTA's col 0 + M: Int32, + N: Int32, + ACTIVATION: cutlass.Constexpr[str | None], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + SWIZZLE: cutlass.Constexpr[bool], + TILE_X: cutlass.Constexpr[int], + TILE_Y: cutlass.Constexpr[int], # pylint: disable=unused-argument # kept for API consistency + SKIP_INPUT_MASKING: cutlass.Constexpr[bool] = False, + SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool] = False, + WITH_ACT: cutlass.Constexpr[bool] = False, + WITH_DACT: cutlass.Constexpr[bool] = False, + WITH_DBIAS: cutlass.Constexpr[bool] = False, + CACHE_ACTIVATION: cutlass.Constexpr[bool] = False, # cache post-activation values to sX_tile +): + """Quantize one SMEM tile colwise to MXFP8 (per-column 32-elt block scales); returns (amax, dbias_partial).""" + 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)), + ) + + sX_tv = cute.composition(sX_tile, tv_layout) + sO_tv = cute.composition(sO_col_tile, tv_layout) + + # I/O Elements that belong to this thread + sX_thread = sX_tv[tidx, None] + sO_thread = sO_tv[tidx, None] + + # PTX allows to fuse relu activation in `cvt.rn.satfinite` unless we need to reduction for dbias + FUSE_RELU = cutlass.const_expr(ACTIVATION == "relu") and not WITH_DBIAS + # Keep input in half precision format if possible + USE_HALF_PRECISION = is_packed16(DTYPE) and (ACTIVATION is None or FUSE_RELU) + dbias_partial = Float32(0.0) + + if cutlass.const_expr(USE_HALF_PRECISION): + max_scalar = max_scalar_f16 if DTYPE is cutlass.Float16 else max_scalar_bf16 + abs_max_scalar = abs_max_scalar_f16 if DTYPE is cutlass.Float16 else abs_max_scalar_bf16 + to_f32 = to_f32_f16 if DTYPE is cutlass.Float16 else to_f32_bf16 + # If we can use the half precision format, then use the input tile directly since there is no need to upcast + sX_thread_packed16 = cute.make_tensor( + sX_thread.iterator, + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(TILE_X,)), + ) + # Stash the strided column reads in registers (CUDA's in_colwise_IType): + # the cvt loop below reuses them instead of re-reading smem. + in_c = [None] * MXFP8_BLOCK_SCALING_SIZE + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + in_c[i] = sX_thread_packed16[i] + + amax_value = in_c[0] + # Skip the first iteration because we assigned in_c[0] to amax_value already + for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE): + if cutlass.const_expr(FUSE_RELU): + # If we fuse relu then we don't want to do abs since negative value will be set to 0 + # and they will lose comparison automatically + amax_value = max_scalar(amax_value, in_c[i]) + else: + amax_value = abs_max_scalar(amax_value, in_c[i]) + if cutlass.const_expr(FUSE_RELU): + amax_c = cute.arch.fmax(to_f32(amax_value), Float32(0.0)) + else: + amax_c = fabs_f32(to_f32(amax_value)) + else: + # Otherwise we need to case input values to fp32. Allocate the register tensor and load from SMEM input tiles. + rX_thread_f32 = cute.make_rmem_tensor( + layout_or_shape=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(1,)), + dtype=Float32, + ) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + rX_thread_f32[i] = Float32(sX_thread[i]) + # Apply activation (fwd) or grad·act'(act_input) (bwd dact) in f32. + if cutlass.const_expr(WITH_DACT): + dop = SUPPORTED_DACTIVATIONS[ACTIVATION] + sA_thread = cute.composition(sActInput_tile, tv_layout)[tidx, None] + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + rX_thread_f32[i] = rX_thread_f32[i] * dop(Float32(sA_thread[i])) + elif cutlass.const_expr(WITH_ACT): + op = SUPPORTED_ACTIVATIONS[ACTIVATION] + # Don't repeat these OOB masking computation in the loop + thread_col_start = tile_col_start + tidx + thread_col_oob = thread_col_start >= N + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + # If it's relu, we can handle it later in the cvt + if not cutlass.const_expr(FUSE_RELU): + rX_thread_f32[i] = op(rX_thread_f32[i]) + # If the input shape is not divisible by the tile size, + # TMA would zero-fills the input tile outside its logical MxN bounds. + # This is fine for non-activation cases, but for activation cases, + # op(0) might not be 0 which will pollute the amax and dbias. + # So we must manually mask the OOB region here. + if not cutlass.const_expr(SKIP_INPUT_MASKING): + if tile_row_start + i >= M or thread_col_oob: + rX_thread_f32[i] = Float32(0.0) + # Accumulate fp32 activations to DBIAS before we truncate to half precision when the input is half precision + if cutlass.const_expr(WITH_DBIAS): + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + dbias_partial += rX_thread_f32[i] + # Truncate the activation (after we apply op) back to the half precision type if input is also half precision. + if cutlass.const_expr(is_packed16(DTYPE) and ACTIVATION is not None): + truncate_f32 = truncate_f32_f16 if DTYPE is cutlass.Float16 else truncate_f32_bf16 + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + rX_thread_f32[i] = truncate_f32(rX_thread_f32[i]) + # Columnwise is the preferred direction so it runs first. If it needs to cache the activation in the input tile + # to let the rowwise pass read it, we need to cast and overwrite the input data in-place here + if cutlass.const_expr(CACHE_ACTIVATION): + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + sX_thread[i] = DTYPE(rX_thread_f32[i]) + amax_c = Float32(0.0) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + if cutlass.const_expr(FUSE_RELU): + amax_c = cute.arch.fmax(amax_c, rX_thread_f32[i]) + else: + amax_c = cute.arch.fmax(amax_c, fabs_f32(rX_thread_f32[i])) + + # Irregular shapes: skip when this stage's row range or this thread's + # 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 cutlass.const_expr(SKIP_SCALE_BOUNDS): + if cutlass.const_expr(SWIZZLE): + mS_col_stage[(0, tidx % 32, tidx // 32)] = biased_exp_c + else: + mS_col_stage[(0, tidx)] = biased_exp_c + else: + scale_col = tile_col_start + tidx + if tile_row_start < M and scale_col < N: + if cutlass.const_expr(SWIZZLE): + mS_col_stage[(0, tidx % 32, tidx // 32)] = biased_exp_c + else: + mS_col_stage[(0, tidx)] = biased_exp_c + + inv_scale_c = exp2f_rcp(biased_exp_c) + # cvt.rn.satfinite can be vectorized to convert 2 f32 to 2 fp8 in one instruction + cvt_x2_func = get_cvt_f32x2_to_fp8x2_func(FP8_DTYPE, FUSE_RELU) + sO_thread_fp8 = cute.make_tensor( + cute.recast_ptr(sO_thread.iterator, dtype=FP8_DTYPE), sO_thread.layout + ) + if cutlass.const_expr(USE_HALF_PRECISION): + to_f32 = to_f32_f16 if DTYPE is cutlass.Float16 else to_f32_bf16 + for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): + lo, hi = 2 * j, 2 * j + 1 + v_lo = to_f32(in_c[lo]) + v_hi = to_f32(in_c[hi]) + # Accumulate the per-thread column partial for dbias if WITH_DBIAS. + # Kept as two adds in element order: f32 addition is not associative + if cutlass.const_expr(WITH_DBIAS): + dbias_partial += v_lo + dbias_partial += v_hi + pair = cvt_x2_func(v_hi * inv_scale_c, v_lo * inv_scale_c) + sO_thread_fp8[lo] = Uint8(pair & Uint16(0xFF)).bitcast(FP8_DTYPE) + sO_thread_fp8[hi] = Uint8(pair >> Uint16(8)).bitcast(FP8_DTYPE) + else: + for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): + lo, hi = 2 * j, 2 * j + 1 + pair = cvt_x2_func(rX_thread_f32[hi] * inv_scale_c, rX_thread_f32[lo] * inv_scale_c) + sO_thread_fp8[lo] = Uint8(pair & Uint16(0xFF)).bitcast(FP8_DTYPE) + sO_thread_fp8[hi] = Uint8(pair >> Uint16(8)).bitcast(FP8_DTYPE) + + # Return this stage's per-column partial alongside amax; the caller accumulates + # it across stages (a scalar can't be updated in-place through the arg). + return amax_c, dbias_partial + + +@cute.jit +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], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + TILE_X: cutlass.Constexpr[int], + WITH_DACT: cutlass.Constexpr[bool] = False, + CACHE_ACTIVATION: cutlass.Constexpr[bool] = False, # cache post-activation values to sX_tile +): + """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. + _, 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)), + ) + sX_thread = cute.composition(sX_tile, tv_layout)[tidx, None] + + dbias_partial = Float32(0.0) + if cutlass.const_expr(WITH_DACT): + dop = SUPPORTED_DACTIVATIONS[ACTIVATION] + sA_thread = cute.composition(sA_tile, tv_layout)[tidx, None] + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + x = Float32(sX_thread[i]) * dop(Float32(sA_thread[i])) + dbias_partial += x + if cutlass.const_expr(CACHE_ACTIVATION): + sX_thread[i] = DTYPE(x) + else: + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + dbias_partial += Float32(sX_thread[i]) + return dbias_partial + + +@cute.jit +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) + 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 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 + """ + + _NUM_STAGES = 2 # The pipeline depth is always 2 + + def __init__( + self, + cfg: MXFP8QuantizeConfig, + SKIP_INPUT_MASKING: bool = False, + SKIP_SCALE_BOUNDS: bool = False, + ): + self.cfg = cfg + # If the input shape is divisible by the tile size or f(0)=0 holds for activaions, + # we can skip masking inputs with zero in the kernel and save some instructions. + self.SKIP_INPUT_MASKING = SKIP_INPUT_MASKING + # If the input shape is divisible by the tile size, we can skip bounds check for scale writes and save some instructions. + self.SKIP_SCALE_BOUNDS = SKIP_SCALE_BOUNDS + # Only honor the noop flag when no activation or dbias is fused to match CUDA C++'s implementation + self.CHECK_NOOP_FLAG: cutlass.const_expr = ( + not self.cfg.WITH_ACT and not self.cfg.WITH_DACT and not self.cfg.WITH_DBIAS + ) + cast_dbias_only = cfg.WITH_DBIAS and not cfg.WITH_DACT and not cfg.WITH_ACT + # Use a different tile size for dbias only config + # No matter what tile size we use, each thread always handles a (1, MXFP8_BLOCK_SCALING_SIZE) chunk + if cast_dbias_only: + self._NUM_TILES = 4 # Each CTA handles 4 tiles stacked vertically + self._THREADS_PER_CTA = 128 + else: + self._NUM_TILES = 2 # Each CTA handles 2 tiles stacked vertically + self._THREADS_PER_CTA = 64 + # Each thread handles a (1, MXFP8_BLOCK_SCALING_SIZE) chunk + self._TILE_COLS = self._THREADS_PER_CTA + self._TILE_ROWS = MXFP8_BLOCK_SCALING_SIZE + self._NUM_WARPS = self._THREADS_PER_CTA // 32 + # We prefer to do dbias reduction in colwise which is easier (no cross-thread reduction needed). + # Only do rowwise reduction when we don't quantize columnwisely when WITH_DBIAS is True. + + # If columnwise is not quantized, and dbias is fused, we will do a columnwise reduction only pass + # (may also cache dact values) without quantization because doing dbias in rowwise requires a SMEM shuffle + # which is even slower + self.COLWISE_DBIAS_REDUCTION_ONLY = cfg.WITH_DBIAS and not cfg.COLWISE + # If columnwise is quantized and dbias is fused, do dbias reduction in colwise which is easier + # (no cross-thread reduction needed). + self.DBIAS_REDUCTION_IN_COLWISE = cfg.WITH_DBIAS and cfg.COLWISE + # Never do dbias reduction in rowwise because it is slow. Leave this so in case it's enabled in the future + self.DBIAS_REDUCTION_IN_ROWWISE = False + + # Cache activation in-place in the SMEM input tile when we process both rowwise and colwise passes + # so the activation is only computed once in the direction we favor (columnwise) and the other direction (rowwise) + # reads the cached value instead of recomputing it. + # Note: if activation is relu, there is no standalong relu applied because it's already fused into `cvt.rn.satfinite` + # so it should be treated as "no activation" + self.CACHE_ACTIVATION = ( + (cfg.WITH_ACT or cfg.WITH_DACT) + and cfg.ROWWISE + # Always cache activation in the colwise quantization or dbias reduction only pass + and (cfg.COLWISE or self.COLWISE_DBIAS_REDUCTION_ONLY) + and cfg.ACTIVATION != "relu" + ) + + @cute.jit + def __call__( + self, + mX: cute.Tensor, # Input tensor to quantize + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], # Rowwise output and scale tensors + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], # Colwise output and scale tensors + mAmax: Optional[cute.Tensor], # Global amax accumulator, only used when WITH_AMAX is True + mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device + mDActInput: Optional[ + cute.Tensor + ], # Activation input for activation derivative fusion, only used when WITH_DACT is True + mWorkspace: Optional[ + cute.Tensor + ], # Workspace for the dbias reduction, only used when WITH_DBIAS is True + stream: CUstream, + ): + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + f"[CuTeDSL] MXFP8QuantizeGeneralG2SKernel.__call__() with config: {self.cfg}\n" + ) + + M = mX.shape[0] + N = mX.shape[1] + cfg = self.cfg + + # If WITH_GEMM_SWIZZLED_SCALES is enabled, the output must satisfy cublas's swizzled layout + # This is expressed as a CuTe layout applied to the output tensor so it can be transparent throughout the kernel implementation. + # See https://docs.nvidia.com/cuda/cublas/#d-block-scaling-factors-layout for more details. + if cutlass.const_expr(cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row, mS_col = derive_swizzled_scale_layout( + M, N, cfg.ROWWISE, cfg.COLWISE, mS_row, mS_col + ) + + # We have 2 stages in our pipeline where each stage loads / computes a (TILE_Y, TILE_X) tile + smem_tile_layout = cute.make_ordered_layout( + (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) + ) + cta_tiler = (self._TILE_ROWS, self._TILE_COLS) + + # Input TMA atoms + op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + tma_atom, tma_src = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_load, + mX, + smem_tile_layout, + cta_tiler, + num_multicast=1, + ) + + # Activation input TMA atoms for activation derivative fusion + tma_atom_act = None + tma_src_act = None + if cutlass.const_expr(cfg.WITH_DACT): + tma_atom_act, tma_src_act = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_load, + mDActInput, + smem_tile_layout, + cta_tiler, + num_multicast=1, + ) + + # Output TMA atoms + op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + out_smem_layout = cute.make_ordered_layout((self._TILE_ROWS, self._TILE_COLS), order=(1, 0)) + tma_atom_out_row = None + tma_dst_out_row = None + tma_atom_out_col = None + tma_dst_out_col = None + if cutlass.const_expr(cfg.ROWWISE): + 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, + ) + if cutlass.const_expr(cfg.COLWISE): + 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 = [ + cute.ceil_div(Int32(N), self._TILE_COLS), + cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES), + ] + block = [ + self._THREADS_PER_CTA, + ] + + self.kernel( + mX, + mS_row, + mS_col, + mAmax, + mNoop, + mWorkspace, + mX.element_type, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + tma_atom_act, + tma_src_act, + ).launch( + grid=grid, + block=block, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mS_row: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mNoop: cute.Pointer, + mWorkspace: Optional[cute.Tensor], + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom: cute.CopyAtom, + tma_src: cute.Tensor, # Input TMA atoms + tma_atom_out_row: Optional[cute.CopyAtom], + tma_dst_out_row: Optional[cute.Tensor], # Rowwise output TMA atoms + tma_atom_out_col: Optional[cute.CopyAtom], + tma_dst_out_col: Optional[cute.Tensor], # Colwise output TMA atoms + tma_atom_act: Optional[cute.CopyAtom], + tma_src_act: Optional[ + cute.Tensor + ], # Activation derivative TMA atoms, None unless WITH_DACT + ): + """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" + + skip_execution = False + if cutlass.const_expr(self.CHECK_NOOP_FLAG): + skip_execution = noop_flag_is_set(mNoop) + if not skip_execution: + self._kernel_main( + mX, + mS_row, + mS_col, + mAmax, + mWorkspace, + dtype, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + tma_atom_act, + tma_src_act, + ) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mS_row: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom: cute.CopyAtom, + tma_src: cute.Tensor, # Input TMA atoms + tma_atom_out_row: Optional[cute.CopyAtom], + tma_dst_out_row: Optional[cute.Tensor], # Rowwise output TMA atoms + tma_atom_out_col: Optional[cute.CopyAtom], + tma_dst_out_col: Optional[cute.Tensor], # Colwise output TMA atoms + tma_atom_act: Optional[cute.CopyAtom], + tma_src_act: Optional[ + cute.Tensor + ], # Activation derivative TMA atoms, None unless WITH_DACT + ): + cfg = self.cfg + + if cutlass.const_expr(cfg.ROWWISE): + mS_row = cute.zipped_divide( + mS_row, (self._TILE_ROWS, self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + ) + if cutlass.const_expr(cfg.COLWISE): + mS_col = cute.zipped_divide( + mS_col, (self._TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, self._TILE_COLS) + ) + + # Allocate shared memory for the input and rowwise / columnwise outputs + FP8_DTYPE = cfg.FP8_DTYPE + + if cutlass.const_expr(cfg.ROWWISE and cfg.COLWISE): + + @cute.struct + class SharedStorage: + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] + 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, + ] + sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] + + elif cutlass.const_expr(cfg.ROWWISE and not cfg.COLWISE): + + @cute.struct + class SharedStorage: + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] + 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, + ] + sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] + + else: + + @cute.struct + class SharedStorage: + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] + sX: cute.struct.Align[ + cute.struct.MemRange[ + 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, + ] + sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] + + 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 + sX = storage.sX.get_tensor( + cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + ) + if cutlass.const_expr(cfg.ROWWISE): + sO_row = storage.sO_row.get_tensor( + cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + ) + if cutlass.const_expr(cfg.COLWISE): + sO_col = storage.sO_col.get_tensor( + cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + ) + + # Allocate shared memory for the activation input used for the activation derivative fusion. + if cutlass.const_expr(cfg.WITH_DACT): + + @cute.struct + class DactStorage: + sActInput: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + + dact_storage = smem.allocate(DactStorage) + # Apply the same layout as the input + sActInput = dact_storage.sActInput.get_tensor( + cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + ) + + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # Prefetch TMA descriptors + if warp_idx == 0: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom) + if cutlass.const_expr(cfg.WITH_DACT): + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_act) + + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + + # Only warp 0 is the producer (issues TMA) + producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + # Every warp is the consumer (reads the data loaded by TMA) + consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._NUM_WARPS) + + # Bytes transferred per TMA copy: one (TILE_Y, TILE_X) tile of dtype. + tx_count = self._TILE_ROWS * self._TILE_COLS * dtype.width // 8 + # dact loads two tiles (grad + act_input) under the same per-stage barrier, + # so the barrier must expect both copies' bytes. + if cutlass.const_expr(cfg.WITH_DACT): + tx_count *= 2 + + 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, # single-CTA, no cluster/multicast + ) + + prod_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._NUM_STAGES + ) + cons_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._NUM_STAGES + ) + + 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), + ) + + # 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)) + + # Partition sX/gX for the TMA atom (single-CTA, no cluster/multicast). + tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( + tma_atom, + 0, # Use the only CTA to do the TMA copy + cute.make_layout(1), # This cluster only has 1 CTAs + sX, + gX_tiled, + ) + + # If WITH_DACT, partition the activation input for TMA as well in the same way + if cutlass.const_expr(cfg.WITH_DACT): + gA_tiled = cute.zipped_divide(tma_src_act, (self._TILE_ROWS, self._TILE_COLS)) + tXsA, tXgA = cute.nvgpu.cpasync.tma_partition( + tma_atom_act, + 0, + cute.make_layout(1), + sActInput, + gA_tiled, + ) + + # Partitioning for rowwise / columnwise outputs + if cutlass.const_expr(cfg.ROWWISE): + 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, + ) + if cutlass.const_expr(cfg.COLWISE): + 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, + ) + + # Ensure barrier init is visible to all threads before the pipeline is used. + 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: + mainloop_pipeline.producer_acquire(prod_state) + tile_y = bidy * self._NUM_TILES + s + cute.copy( + tma_atom, + tXgX[(None, (tile_y, bidx))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + if cutlass.const_expr(cfg.WITH_DACT): + cute.copy( + tma_atom_act, + tXgA[(None, (tile_y, bidx))], + tXsA[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + # Per-thread amax accumulator + if cutlass.const_expr(cfg.WITH_AMAX): + per_thread_amax = Float32(0.0) + + # Prepare thread-level register accumulators for rowwise dbias reduction. + # Each thread will process two (1, MXFP8_BLOCK_SCALING_SIZE) rows in two stages, and in each stage the thread will add the + # (after dact applied) value to this register array with the same shape so it carries the the two stages' partial sum. + # Then it will be written to a SMEM buffer to let the whole CTA do the reduction separately to yield + # the final (1, TILE_X) dbias workspace output. + rowwise_dbias_acc = None + if cutlass.const_expr(self.DBIAS_REDUCTION_IN_ROWWISE): + rowwise_dbias_acc = cute.make_rmem_tensor( + layout_or_shape=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(1,)), + dtype=Float32, + ) + # Zero the accumulator registers. + for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + rowwise_dbias_acc[c] = Float32(0.0) + block_dbias = Float32(0.0) + # Prepare thread-level register for columnwise dbias reduction. + # Each thread will process two (MXFP8_BLOCK_SCALING_SIZE, 1) columns in two stages, and in each stage the thread will reduce the + # (after dact applied) column to (1,) and add to this register. + # Then this partial sum scalar will be written to the GMEM workspace buffer directly. + if cutlass.const_expr(self.DBIAS_REDUCTION_IN_COLWISE or self.COLWISE_DBIAS_REDUCTION_ONLY): + block_dbias = Float32(0.0) + + # 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 + sX_tile = sX[(None, stage_idx)] + # Also fetch the activation input if WITH_DACT + sActInput_tile = None + if cutlass.const_expr(cfg.WITH_DACT): + sActInput_tile = sActInput[(None, stage_idx)] + # Each CTA handles `NUM_TILES` tiles stacked vertically, so tile_idx_x is just the block index along X dimension + # 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 + 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). + # Each CTA owns NUM_TILES consecutive row-tiles of GRID. cute + # auto-decomposes the flat row coord `bidy * NUM_TILES + tile_idx` + # onto GRID's hierarchical row modes — which is the + # (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))]) + 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))]) + + amax, dbias = quantize_mxfp8_smem_tile_general( + 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( + "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) + + # Warp 0 issues TMA copy to write the quantized output tile from shared memory to global memory and then commits + if warp_idx == 0: + tile_y = bidy * self._NUM_TILES + tile_idx + if cutlass.const_expr(cfg.ROWWISE): + cute.copy( + tma_atom_out_row, + tXsO_row[(None, stage_idx)], + tXgO_row[(None, (tile_y, bidx))], + ) + if cutlass.const_expr(cfg.COLWISE): + cute.copy( + tma_atom_out_col, + tXsO_col[(None, stage_idx)], + tXgO_col[(None, (tile_y, bidx))], + ) + cute.arch.cp_async_bulk_commit_group() + + cons_state.advance() + + # The pipeline is no longer fully filled after we consume this tile, so we fetch a new tile to fill the pipeline. + # The next _NUM_STAGES-1 tiles are already in-flight, so the next tile to fetch is after _NUM_STAGES tiles. + if warp_idx == 0: + next_tile_idx = tile_idx + self._NUM_STAGES + if next_tile_idx < num_tiles: + mainloop_pipeline.producer_acquire(prod_state) + tile_y = bidy * self._NUM_TILES + next_tile_idx + cute.copy( + tma_atom, + tXgX[(None, (tile_y, bidx))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + if cutlass.const_expr(cfg.WITH_DACT): + cute.copy( + tma_atom_act, + tXgA[(None, (tile_y, bidx))], + tXsA[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + # End of the main pipeline loop + + # Complete the cross-thread dbias reduction after each thread has its own per-thread partial sum after the rowwise quantization. + if cutlass.const_expr(self.DBIAS_REDUCTION_IN_ROWWISE): + # If we do the dbias reduction in the rowwise pass, each thread will have a (1, MXFP8_BLOCK_SCALING_SIZE) partial sum + # and we need to write these to a SMEM buffer and let each thread reduce it in the columnwise direction + block_dbias = self._dbias_reduction_rowwise_epilouge(smem, tidx, rowwise_dbias_acc) + + # Write the per-tile reduced dbias to the global workspace. + if cutlass.const_expr(cfg.WITH_DBIAS): + dbias_col = bidx * self._TILE_COLS + tidx + if dbias_col < N: + mWorkspace[(bidy, dbias_col)] = block_dbias + + if cutlass.const_expr(cfg.WITH_AMAX): + 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) + + @cute.jit + def _dbias_reduction_rowwise_epilouge( + 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 = ( + self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE * (MXFP8_BLOCK_SCALING_SIZE + 1) + ) + + # Allocate the SMEM buffer that all threads use to reduce the two-stage partial sum (per thread) to the + # partial sum (per block). + @cute.struct + class DbiasStorage: + sDbias: cute.struct.MemRange[Float32, self._TILE_ROWS * DBIAS_BUFF_WIDTH] + + dbias_storage = smem.allocate(DbiasStorage) + sDbias = dbias_storage.sDbias.get_tensor( + cute.make_layout((self._TILE_ROWS, self._TILE_COLS), stride=(DBIAS_BUFF_WIDTH, 1)), + ) + _, tv_layout_dbias_write = cute.make_layout_tv( + thr_layout=cute.make_layout( + (self._TILE_ROWS, self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE), + stride=(self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE, 1), + ), + val_layout=cute.make_layout( + (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) + ), + ) + 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) // THREADS_PER_BANK + # The offset this thread should start reading from based on what's its first bank to access. + offset = bank_group * PACK_SIZE + for w in cutlass.range_constexpr( + WAVES + ): # Each thread starts from this offset when writing into SMEM to avoid bank conflict + 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 * 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( + thr_layout=cute.make_layout((1, self._TILE_COLS), stride=(self._TILE_COLS, 1)), + val_layout=cute.make_layout((self._TILE_ROWS, 1), stride=(1, 1)), + ) + sDbias_reduce = cute.composition(sDbias, tv_layout_dbias_reduce) + # make_layout_tv yields a (thread, value) layout: thread=tidx -> column tidx, + # value=i -> row i. So index [tidx, i] (thread first), summing the column's rows. + block_dbias = Float32(0.0) + for i in cutlass.range_constexpr(self._TILE_ROWS): + block_dbias += sDbias_reduce[tidx, i] + return block_dbias + + @cute.jit + def _amax_epilogue( + self, + sAmax: cute.Tensor, + mAmax: cute.Tensor, + tidx: Int32, + warp_idx: Int32, + per_thread_amax: Float32, + ): + # Reduce and get the per-warp amax. + warp_amax = cute.arch.warp_redux_sync(per_thread_amax, kind="fmax") + # Write the per-warp amax to shared memory + lane_idx = tidx % 32 + if lane_idx == 0: + sAmax[warp_idx] = warp_amax + cute.arch.sync_threads() + if tidx == 0: + cta_amax = Float32(0.0) + # The first thread reduces all the per-warp amax to the per-CTA amax + for w in cutlass.range_constexpr(self._NUM_WARPS): + cta_amax = cute.arch.fmax(cta_amax, sAmax[w]) + amax_i32 = cute.make_tensor( + cute.recast_ptr(mAmax.iterator, dtype=Int32), + cute.make_layout(1), + ) + # The first thread updates the global amax with an atomic max on the bitcasted float value + cute.arch.atomic_max( + amax_i32.iterator, + cta_amax.bitcast(Int32), + ) diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index ed05aa1c3c6..90c61d03052 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 @@ -83,6 +93,72 @@ 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 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.""" + 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). @@ -210,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.""" @@ -258,3 +344,15 @@ 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 + + +# 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.""" + 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]: